1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <cstring>
10
11 #include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
12 #include <executorch/runtime/kernel/kernel_includes.h>
13
14 // A simple lookup table that looks up embeddings in a fixed dictionary and
15 // size.
16
17 // This module is often used to retrieve word embeddings using indices. The
18 // input to the module is a list of indices, and the embedding matrix, and the
19 // output is the corresponding word embeddings.
20 namespace torch {
21 namespace executor {
22 namespace native {
23
24 using Tensor = exec_aten::Tensor;
25 using ScalarType = exec_aten::ScalarType;
26
27 namespace {
28
29 template <typename CTYPE>
embedding_kernel(KernelRuntimeContext & ctx,const Tensor & weight,const Tensor & indices,Tensor & out)30 void embedding_kernel(
31 KernelRuntimeContext& ctx,
32 const Tensor& weight,
33 const Tensor& indices,
34 Tensor& out) {
35 int64_t nbytes_per_entry = weight.size(1) * weight.element_size();
36 const char* w_data = weight.const_data_ptr<char>();
37 char* out_data = out.mutable_data_ptr<char>();
38 const CTYPE* indices_ptr = indices.const_data_ptr<CTYPE>();
39 ssize_t weight_height = weight.size(0);
40 const auto indices_numel = indices.numel();
41 for (int i = 0; i < indices_numel; i++) {
42 // Ensure index is larger than 0 and smaller than weight.size(0)
43 ET_KERNEL_CHECK_MSG(
44 ctx,
45 indices_ptr[i] < weight_height,
46 InvalidArgument,
47 ,
48 "indices_ptr[%d] %ld >= weight.size(0) %zd",
49 i,
50 static_cast<long>(indices_ptr[i]),
51 weight_height);
52 ET_KERNEL_CHECK_MSG(
53 ctx,
54 indices_ptr[i] >= 0,
55 InvalidArgument,
56 ,
57 "indices_ptr[%d] %ld < 0",
58 i,
59 static_cast<long>(indices_ptr[i]));
60 if (w_data != nullptr) {
61 memcpy(
62 out_data,
63 w_data + nbytes_per_entry * indices_ptr[i],
64 nbytes_per_entry);
65 }
66 out_data += nbytes_per_entry;
67 }
68 }
69 } // namespace
70
71 // embedding.out(Tensor weight, Tensor indices, int padding_idx=-1, bool
72 // scale_grad_by_freq=False, bool sparse=False, *, Tensor(a!) out) -> Tensor(a!)
embedding_out(KernelRuntimeContext & ctx,const Tensor & weight,const Tensor & indices,int64_t padding_idx,bool scale_grad_by_freq,bool sparse,Tensor & out)73 Tensor& embedding_out(
74 KernelRuntimeContext& ctx,
75 const Tensor& weight,
76 const Tensor& indices,
77 int64_t padding_idx,
78 bool scale_grad_by_freq,
79 bool sparse,
80 Tensor& out) {
81 (void)ctx;
82 (void)padding_idx;
83 (void)scale_grad_by_freq;
84 (void)sparse;
85
86 ET_KERNEL_CHECK(
87 ctx, check_embedding_args(weight, indices, out), InvalidArgument, out);
88
89 ET_KERNEL_CHECK(
90 ctx,
91 resize_embedding_output(weight, indices, out) == Error::Ok,
92 InvalidArgument,
93 out);
94
95 ET_KERNEL_CHECK_MSG(
96 ctx,
97 out.size(out.dim() - 1) == weight.size(1),
98 InvalidArgument,
99 out,
100 "out.size(%zd) %zd != weight.size(1) %zd",
101 out.dim() - 1,
102 out.size(1),
103 weight.size(1));
104
105 ET_KERNEL_CHECK(
106 ctx,
107 tensors_have_same_dim_order(weight, indices, out),
108 InvalidArgument,
109 out);
110
111 ET_KERNEL_CHECK(
112 ctx, tensor_is_default_dim_order(weight), InvalidArgument, out);
113
114 ScalarType ix_type = indices.scalar_type();
115 ET_CHECK_MSG(
116 ix_type == ScalarType::Long || ix_type == ScalarType::Int,
117 "Expected indices tensor to have Long or Int scalar types");
118
119 ET_SWITCH_TWO_TYPES(
120 Long, Int, ix_type, ctx, "op_embedding.out", CTYPE, [&]() {
121 embedding_kernel<CTYPE>(ctx, weight, indices, out);
122 });
123
124 return out;
125 }
126
127 } // namespace native
128 } // namespace executor
129 } // namespace torch
130