Add embedding ops executorch (#1137)

Summary:

Adds ExecuTorch ops for embedding

Reviewed By: digantdesai, manuelcandales

Differential Revision: D64477043

Author: metascroy

torchao/experimental/CMakeLists.txt

@@ -48,9 +48,16 @@ if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     )
     if(TORCHAO_BUILD_EXECUTORCH_OPS)
         add_library(torchao_ops_executorch STATIC)
-        target_link_libraries(torchao_ops_executorch PRIVATE torchao_ops_linear_8bit_act_xbit_weight_executorch)
+        target_link_libraries(torchao_ops_executorch PRIVATE
+            torchao_ops_linear_8bit_act_xbit_weight_executorch
+            torchao_ops_embedding_xbit_executorch
+        )
         install(
-            TARGETS torchao_ops_executorch torchao_ops_linear_8bit_act_xbit_weight_executorch torchao_kernels_aarch64
+            TARGETS
+                torchao_ops_executorch
+                torchao_kernels_aarch64
+                torchao_ops_linear_8bit_act_xbit_weight_executorch
+                torchao_ops_embedding_xbit_executorch
             EXPORT _targets
             DESTINATION lib
         )

torchao/experimental/ops/embedding_xbit/CMakeLists.txt

@@ -17,3 +17,20 @@ target_link_libraries(torchao_ops_embedding_xbit_aten PRIVATE torchao_kernels_aa
 target_include_directories(torchao_ops_embedding_xbit_aten PRIVATE "${TORCH_INCLUDE_DIRS}")
 target_link_libraries(torchao_ops_embedding_xbit_aten PRIVATE "${TORCH_LIBRARIES}")
 target_compile_definitions(torchao_ops_embedding_xbit_aten PRIVATE USE_ATEN=1)
+
+if(TORCHAO_BUILD_EXECUTORCH_OPS)
+    # ExecuTorch package is not required, but EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES must
+    # be defined and EXECUTORCH_LIBRARIES must include the following libraries installed by ExecuTorch:
+    # libexecutorch.a
+    # libextension_threadpool.a
+    # libcpuinfo.a
+    # libpthreadpool.a
+    add_library(torchao_ops_embedding_xbit_executorch OBJECT
+        op_embedding_xbit_executorch.cpp
+    )
+    target_link_torchao_parallel_backend(torchao_ops_embedding_xbit_executorch executorch)
+    target_include_directories(torchao_ops_embedding_xbit_executorch PRIVATE "${EXECUTORCH_INCLUDE_DIRS}")
+    target_compile_definitions(torchao_ops_embedding_xbit_executorch PRIVATE USE_EXECUTORCH=1)
+    target_link_libraries(torchao_ops_embedding_xbit_executorch PRIVATE "${EXECUTORCH_LIBRARIES}")
+    target_link_libraries(torchao_ops_embedding_xbit_executorch PRIVATE torchao_kernels_aarch64)
+endif()

torchao/experimental/ops/embedding_xbit/op_embedding_xbit_executorch.cpp

@@ -0,0 +1,38 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <torchao/experimental/ops/embedding_xbit/op_embedding_xbit-impl.h>
+
+#define DEFINE_OP(weight_nbit)             \
+  Tensor _op_out_##weight_nbit(            \
+      RuntimeContext& ctx,                 \
+      const Tensor& packed_weight_qvals,   \
+      const Tensor& num_embeddings_tensor, \
+      const Tensor& embedding_dim_tensor,  \
+      const Tensor& weight_scales,         \
+      const Tensor& weight_zeros,          \
+      const Tensor& indices,               \
+      Tensor& out) {                       \
+    (void)ctx;                             \
+    embedding_out_cpu<weight_nbit>(        \
+        packed_weight_qvals,               \
+        num_embeddings_tensor,             \
+        embedding_dim_tensor,              \
+        weight_scales,                     \
+        weight_zeros,                      \
+        indices,                           \
+        out);                              \
+    return out;                            \
+  }                                        \
+  EXECUTORCH_LIBRARY(                      \
+      torchao, "_embedding_" #weight_nbit "bit.out", _op_out_##weight_nbit)
+
+DEFINE_OP(1);
+DEFINE_OP(2);
+DEFINE_OP(3);
+DEFINE_OP(4);
+DEFINE_OP(5);
+DEFINE_OP(6);