up

Author: metascroy

.github/workflows/torchao_experimental_test.yml

@@ -34,6 +34,7 @@ jobs:
         run: |
           conda activate venv
           pip install torch --index-url "https://download.pytorch.org/whl/nightly/cpu"
+          pip install executorch
           pip install numpy
           pip install pytest
           pip install parameterized
@@ -57,6 +58,12 @@ jobs:
           sh build_and_run_tests.sh
           rm -rf /tmp/cmake-out
           popd
+      - name: ET ops build
+        run: |
+          conda activate venv
+          pushd torchao/experimental
+          sh build_torchao_ops.sh executorch
+          popd
 
   test-mps-ops:
     strategy:

torchao/experimental/build_torchao_ops.sh

@@ -21,6 +21,7 @@ export CMAKE_OUT=cmake-out
 cmake -DCMAKE_PREFIX_PATH=${CMAKE_PREFIX_PATH} \
     -DCMAKE_INSTALL_PREFIX=${CMAKE_OUT} \
     -DTORCHAO_BUILD_EXECUTORCH_OPS="${TORCHAO_BUILD_EXECUTORCH_OPS}" \
+    -DTORCHAO_BUILD_CPU_AARCH64=ON \
     -S . \
     -B ${CMAKE_OUT}
 cmake --build  ${CMAKE_OUT} -j 16 --target install --config Release

torchao/experimental/ops/embedding_xbit/CMakeLists.txt

@@ -27,6 +27,10 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
     # libextension_threadpool.a
     # libcpuinfo.a
     # libpthreadpool.a
+    if(NOT DEFINED EXECUTORCH_INCLUDE_DIRS AND NOT DEFINED EXECUTORCH_LIBRARIES)
+        message(WARNING "EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES are not defined. Looking for ExecuTorch.")
+        find_package(ExecuTorch HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
+    endif()
     add_library(torchao_ops_embedding_xbit_executorch OBJECT
         op_embedding_xbit_executorch.cpp
     )

torchao/experimental/ops/linear_8bit_act_xbit_weight/CMakeLists.txt

@@ -41,11 +41,15 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
     # libextension_threadpool.a
     # libcpuinfo.a
     # libpthreadpool.a
+    if(NOT DEFINED EXECUTORCH_INCLUDE_DIRS AND NOT DEFINED EXECUTORCH_LIBRARIES)
+        message(WARNING "EXECUTORCH_INCLUDE_DIRS and EXECUTORCH_LIBRARIES are not defined. Looking for ExecuTorch.")
+        find_package(ExecuTorch HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
+    endif()
     # find_package(ExecuTorch HINTS ${CMAKE_PREFIX_PATH}/executorch/share/cmake)
-    file(GLOB _SRCS "${CMAKE_CURRENT_SOURCE_DIR}/op_linear_8bit_act_xbit_weight_executorch/*.cpp")
+    # file(GLOB _SRCS "${CMAKE_CURRENT_SOURCE_DIR}/op_linear_8bit_act_xbit_weight_executorch/*.cpp")
     add_library(torchao_ops_linear_8bit_act_xbit_weight_executorch OBJECT
         linear_8bit_act_xbit_weight.cpp
-        ${_SRCS}
+        op_linear_8bit_act_xbit_weight_executorch.cpp
     )
     target_link_torchao_parallel_backend(torchao_ops_linear_8bit_act_xbit_weight_executorch executorch)
     target_include_directories(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE "${EXECUTORCH_INCLUDE_DIRS}")

torchao/experimental/ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight_executorch.cpp

@@ -0,0 +1,65 @@
+#include <torchao/experimental/ops/linear_8bit_act_xbit_weight/op_linear_8bit_act_xbit_weight-impl.h>
+
+#define DEFINE_OP(weight_nbit)                                                 \
+  Tensor _op_out_0zp_##weight_nbit(                                            \
+      RuntimeContext &ctx, const Tensor &activations,                          \
+      const Tensor &packed_weights, const int64_t &group_size,                 \
+      const int64_t &n, const int64_t &k, Tensor &out) {                       \
+    (void)ctx;                                                                 \
+    linear_out_cpu<weight_nbit, false>(activations, packed_weights,            \
+                                       group_size, n, k, out);                 \
+    return out;                                                                \
+  }                                                                            \
+  Tensor _op_out_zp_##weight_nbit(                                             \
+      RuntimeContext &ctx, const Tensor &activations,                          \
+      const Tensor &packed_weights, const int64_t &group_size,                 \
+      const int64_t &n, const int64_t &k, Tensor &out) {                       \
+    (void)ctx;                                                                 \
+    linear_out_cpu<weight_nbit, true>(activations, packed_weights, group_size, \
+                                      n, k, out);                              \
+    return out;                                                                \
+  }
+
+#define REGISTER_0ZP(weight_nbit)                                              \
+  EXECUTORCH_LIBRARY(torchao,                                                  \
+                     "_linear_8bit_act_" #weight_nbit "bit0zp_weight.out",     \
+                     _op_out_0zp_##weight_nbit)
+
+#define REGISTER_ZP(weight_nbit)                                               \
+  EXECUTORCH_LIBRARY(torchao,                                                  \
+                     "_linear_8bit_act_" #weight_nbit "bit_weight.out",        \
+                     _op_out_zp_##weight_nbit)
+
+// This looks a bit ridiculous, but I could not get it to compile with two
+// EXECUTORCH_LIBRARY nested inside DEFINE_OP
+DEFINE_OP(1)
+REGISTER_0ZP(1);
+REGISTER_ZP(1);
+
+DEFINE_OP(2)
+REGISTER_0ZP(2);
+REGISTER_ZP(2);
+
+DEFINE_OP(3)
+REGISTER_0ZP(3);
+REGISTER_ZP(3);
+
+DEFINE_OP(4)
+REGISTER_0ZP(4);
+REGISTER_ZP(4);
+
+DEFINE_OP(5)
+REGISTER_0ZP(5);
+REGISTER_ZP(5);
+
+DEFINE_OP(6)
+REGISTER_0ZP(6);
+REGISTER_ZP(6);
+
+DEFINE_OP(7)
+REGISTER_0ZP(7);
+REGISTER_ZP(7);
+
+DEFINE_OP(8)
+REGISTER_0ZP(8);
+REGISTER_ZP(8);