up

Author: metascroy

setup.py

@@ -76,19 +76,53 @@ def use_debug_mode():
     IS_WINDOWS,
 )
 
-build_torchao_experimental_mps = (
-    os.getenv("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1"
-    and build_torchao_experimental
-    and torch.mps.is_available()
-)
 
-if os.getenv("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1":
-    if use_cpp != "1":
-        print("Building experimental MPS ops requires USE_CPP=1")
-    if not platform.machine().startswith("arm64") or platform.system() != "Darwin":
-        print("Experimental MPS ops require Apple Silicon.")
-    if not torch.mps.is_available():
-        print("MPS not available. Skipping compilation of experimental MPS ops.")
+class BuildOptions:
+    def __init__(self):
+        # TORCHAO_BUILD_CPU_AARCH64 is enabled by default on Arm-based Apple machines
+        # The kernels require sdot/udot, which are not required on Arm until Armv8.4 or later,
+        # but are available on Arm-based Apple machines.  On non-Apple machines, the kernels
+        # can be built by explicitly setting TORCHAO_BUILD_CPU_AARCH64=1
+        self.build_cpu_aarch64 = self._os_bool_var(
+            "TORCHAO_BUILD_CPU_AARCH64",
+            default=(self._is_arm64() and self._is_macos()),
+        )
+        if self.build_cpu_aarch64:
+            assert (
+                self._is_arm64()
+            ), "TORCHAO_BUILD_CPU_AARCH64 requires an arm64 machine"
+
+        # TORCHAO_BUILD_KLEIDIAI is disabled by default for now because
+        # 1) It increases the build time
+        # 2) It has some accuracy issues in CI tests due to BF16
+        self.build_kleidi_ai = self._os_bool_var(
+            "TORCHAO_BUILD_KLEIDIAI", default=False
+        )
+        if self.build_kleidi_ai:
+            assert (
+                self.build_cpu_aarch64
+            ), "TORCHAO_BUILD_KLEIDIAI requires TORCHAO_BUILD_CPU_AARCH64 be set"
+
+        # TORCHAO_BUILD_EXPERIMENTAL_MPS is disabled by default.
+        self.build_experimental_mps = self._os_bool_var(
+            "TORCHAO_BUILD_EXPERIMENTAL_MPS", default=False
+        )
+        if self.build_experimental_mps:
+            assert self._is_macos(), "TORCHAO_BUILD_EXPERIMENTAL_MPS requires MacOS"
+            assert (
+                torch.mps.is_available()
+            ), "TORCHAO_BUILD_EXPERIMENTAL_MPS requires MPS be available"
+
+    def _is_arm64(self) -> bool:
+        return platform.machine().startswith("arm64")
+
+    def _is_macos(self) -> bool:
+        return platform.system() == "Darwin"
+
+    def _os_bool_var(self, var, default) -> bool:
+        default_val = "1" if default else "0"
+        return os.getenv(var, default_val) == "1"
+
 
 # Constant known variables used throughout this file
 cwd = os.path.abspath(os.path.curdir)
@@ -303,6 +337,11 @@ def get_extensions():
         )
 
     if build_torchao_experimental:
+        build_options = BuildOptions()
+
+        def bool_to_on_off(value):
+            return "ON" if value else "OFF"
+
         from distutils.sysconfig import get_python_lib
 
         torch_dir = get_python_lib() + "/torch/share/cmake/Torch"
@@ -313,23 +352,17 @@ def get_extensions():
                 cmake_lists_dir="torchao/experimental",
                 cmake_args=(
                     [
-                        (
-                            "-DCMAKE_BUILD_TYPE=" + "Debug"
-                            if use_debug_mode()
-                            else "Release"
-                        ),
-                        # Disable now because 1) KleidiAI increases build time, and 2) KleidiAI has accuracy issues due to BF16
-                        "-DTORCHAO_BUILD_KLEIDIAI=ON",
-                        (
-                            "-DTORCHAO_BUILD_MPS_OPS=" + "ON"
-                            if build_torchao_experimental_mps
-                            else "OFF"
-                        ),
+                        f"-DCMAKE_BUILD_TYPE={'Debug' if use_debug_mode() else 'Release'}",
+                        f"-DTORCHAO_BUILD_CPU_AARCH64={bool_to_on_off(build_options.build_cpu_aarch64)}",
+                        f"-DTORCHAO_BUILD_KLEIDIAI={bool_to_on_off(build_options.build_kleidi_ai)}",
+                        f"-DTORCHAO_BUILD_MPS_OPS={bool_to_on_off(build_options.build_experimental_mps)}",
                         "-DTorch_DIR=" + torch_dir,
                     ]
-                    + ["-DCMAKE_INSTALL_PREFIX=cmake-out"]
-                    if build_torchao_experimental_mps
-                    else []
+                    + (
+                        ["-DCMAKE_INSTALL_PREFIX=cmake-out"]
+                        if build_options.build_experimental_mps
+                        else []
+                    )
                 ),
             )
         )

torchao/experimental/CMakeLists.txt

@@ -17,17 +17,13 @@ endif()
 
 option(TORCHAO_BUILD_EXECUTORCH_OPS "Building torchao ops for ExecuTorch." OFF)
 option(TORCHAO_BUILD_MPS_OPS "Building torchao MPS ops" OFF)
-
+option(TORCHAO_BUILD_CPU_AARCH64 "Build torchao's CPU aarch64 kernels" OFF)
+option(TORCHAO_BUILD_KLEIDIAI "Download, build, and link against Arm KleidiAI library (arm64 only)" OFF)
 
 if(NOT TORCHAO_INCLUDE_DIRS)
   set(TORCHAO_INCLUDE_DIRS ${CMAKE_CURRENT_SOURCE_DIR}/../..)
 endif()
 
-option(TORCHAO_BUILD_KLEIDIAI "Download, build, and link against Arm KleidiAI library (arm64 only)" OFF)
-if(TORCHAO_BUILD_KLEIDIAI)
-    message(STATUS "Building with Arm KleidiAI library")
-    add_compile_definitions(TORCHAO_ENABLE_KLEIDI=1)
-endif()
 include(CMakePrintHelpers)
 
 add_compile_options("-Wall" "-Werror" "-Wno-deprecated")
@@ -36,49 +32,56 @@ include(CMakePrintHelpers)
 message("TORCHAO_INCLUDE_DIRS: ${TORCHAO_INCLUDE_DIRS}")
 include_directories(${TORCHAO_INCLUDE_DIRS})
 
-if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+
+if(TORCHAO_BUILD_CPU_AARCH64)
+    message(STATUS "Building with cpu/aarch64")
+    add_compile_definitions(TORCHAO_BUILD_CPU_AARCH64)
+
     if(TORCHAO_BUILD_KLEIDIAI)
         message(STATUS "Building with Arm KleidiAI library")
         add_compile_definitions(TORCHAO_ENABLE_KLEIDI=1)
     endif()
-    # Defines target torchao_kernels_aarch64
+endif()
+
+
+if (TORCHAO_BUILD_CPU_AARCH64)
+    # Defines torchao_kernels_aarch64
     add_subdirectory(kernels/cpu/aarch64)
-    add_subdirectory(ops/linear_8bit_act_xbit_weight)
-    add_subdirectory(ops/embedding_xbit)
-
-    add_library(torchao_ops_aten SHARED)
-    target_link_libraries(
-        torchao_ops_aten PRIVATE
-        torchao_ops_linear_8bit_act_xbit_weight_aten
-        torchao_ops_embedding_xbit_aten
-    )
-    if (TORCHAO_BUILD_MPS_OPS)
-        message(STATUS "Building with MPS support")
-        add_subdirectory(ops/mps)
-        target_link_libraries(torchao_ops_aten PRIVATE torchao_ops_mps_aten)
-    endif()
+endif()
+
+add_subdirectory(ops/linear_8bit_act_xbit_weight)
+add_subdirectory(ops/embedding_xbit)
 
+add_library(torchao_ops_aten SHARED)
+target_link_libraries(
+    torchao_ops_aten PRIVATE
+    torchao_ops_linear_8bit_act_xbit_weight_aten
+    torchao_ops_embedding_xbit_aten
+)
+if (TORCHAO_BUILD_MPS_OPS)
+    message(STATUS "Building with MPS support")
+    add_subdirectory(ops/mps)
+    target_link_libraries(torchao_ops_aten PRIVATE torchao_ops_mps_aten)
+endif()
+
+install(
+    TARGETS torchao_ops_aten
+    EXPORT _targets
+    DESTINATION lib
+)
+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
+        torchao_ops_embedding_xbit_executorch
+    )
     install(
-        TARGETS torchao_ops_aten
+        TARGETS
+            torchao_ops_executorch
+            # torchao_kernels_aarch64
+            torchao_ops_linear_8bit_act_xbit_weight_executorch
+            torchao_ops_embedding_xbit_executorch
         EXPORT _targets
         DESTINATION lib
     )
-    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
-            torchao_ops_embedding_xbit_executorch
-        )
-        install(
-            TARGETS
-                torchao_ops_executorch
-                torchao_kernels_aarch64
-                torchao_ops_linear_8bit_act_xbit_weight_executorch
-                torchao_ops_embedding_xbit_executorch
-            EXPORT _targets
-            DESTINATION lib
-        )
-    endif()
-else()
-    message(FATAL_ERROR "Torchao experimental ops can only be built on arm64 CPUs.")
 endif()

torchao/experimental/kernels/cpu/aarch64/CMakeLists.txt

@@ -4,7 +4,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 
-if ((CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") OR (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64"))
+if (TORCHAO_BUILD_CPU_AARCH64)
   add_library(
     torchao_kernels_aarch64
     ${CMAKE_CURRENT_SOURCE_DIR}/reduction/find_min_and_max.cpp
@@ -27,9 +27,9 @@ if ((CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") OR (CMAKE_SYSTEM_PROCESSOR STREQUA
     set(TORCHAO_BUILD_KLEIDI ON PARENT_SCOPE)
     target_link_libraries(torchao_kernels_aarch64 PUBLIC kleidiai)
   endif()
-endif()
 
 install(
   TARGETS torchao_kernels_aarch64
   DESTINATION lib
 )
+endif()

torchao/experimental/kernels/cpu/aarch64/tests/build_and_run_tests.sh

@@ -40,6 +40,7 @@ cmake \
     ${EXTRA_ARGS} \
     -DCMAKE_BUILD_TYPE=Debug \
     -DTORCHAO_LIBRARIES=${TORCHAO_LIBRARIES} \
+    -DTORCHAO_BUILD_CPU_AARCH64=ON \
     -S ${TORCHAO_LIBRARIES}/torchao/experimental/kernels/cpu/aarch64/tests \
     -B ${CMAKE_OUT}
 

torchao/experimental/ops/embedding_xbit/CMakeLists.txt

@@ -13,7 +13,9 @@ add_library(torchao_ops_embedding_xbit_aten OBJECT
     op_embedding_xbit_aten.cpp
 )
 target_link_torchao_parallel_backend(torchao_ops_embedding_xbit_aten "aten_openmp")
-target_link_libraries(torchao_ops_embedding_xbit_aten PRIVATE torchao_kernels_aarch64)
+if (TORCHAO_BUILD_CPU_AARCH64)
+    target_link_libraries(torchao_ops_embedding_xbit_aten PRIVATE torchao_kernels_aarch64)
+endif()
 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)
@@ -32,5 +34,7 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
     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)
+    if (TORCHAO_BUILD_CPU_AARCH64)
+        target_link_libraries(torchao_ops_embedding_xbit_executorch PRIVATE torchao_kernels_aarch64)
+    endif()
 endif()

torchao/experimental/ops/embedding_xbit/op_embedding_xbit-impl.h

@@ -6,9 +6,9 @@
 
 #pragma once
 
-#if defined(__aarch64__) || defined(__ARM_NEON)
+#if defined(TORCHAO_BUILD_CPU_AARCH64)
 #include <torchao/experimental/kernels/cpu/aarch64/embedding/embedding.h>
-#endif // defined(__aarch64__) || defined(__ARM_NEON)
+#endif // TORCHAO_BUILD_CPU_AARCH64
 
 #include <torchao/experimental/ops/embedding_xbit/packed_weights_header.h>
 #include <torchao/experimental/ops/library.h>
@@ -145,7 +145,7 @@ Tensor embedding_out_cpu(
       index = index64_ptr[idx];
     }
     TORCHAO_CHECK(index >= 0 && index < num_embeddings, "index out of bounds");
-#if defined(__aarch64__) || defined(__ARM_NEON)
+#if defined(TORCHAO_BUILD_CPU_AARCH64)
     torchao::kernels::cpu::aarch64::embedding::embedding<weight_nbit>(
         out.mutable_data_ptr<float>() + idx * embedding_dim,
         embedding_dim,
@@ -157,7 +157,7 @@ Tensor embedding_out_cpu(
         index);
 #else
     TORCHAO_CHECK(false, "Unsupported platform");
-#endif // defined(__aarch64__) || defined(__ARM_NEON)
+#endif // TORCHAO_BUILD_CPU_AARCH64
   });
 
   return out;
@@ -234,7 +234,7 @@ Tensor pack_embedding_cpu(const Tensor& weight_qvals) {
   header.write(out.mutable_data_ptr());
 
   torchao::parallel_1d(0, num_embeddings, [&](int64_t idx) {
-#if defined(__aarch64__) || defined(__ARM_NEON)
+#if defined(TORCHAO_BUILD_CPU_AARCH64)
     torchao::kernels::cpu::aarch64::embedding::pack_embedding_weight_qvals<
         weight_nbit>(
         out.mutable_data_ptr<int8_t>() +
@@ -244,7 +244,7 @@ Tensor pack_embedding_cpu(const Tensor& weight_qvals) {
         idx);
 #else
     TORCHAO_CHECK(false, "Unsupported platform");
-#endif // defined(__aarch64__) || defined(__ARM_NEON)
+#endif // defined(TORCHAO_BUILD_CPU_AARCH64)
   });
 
   return out;

torchao/experimental/ops/linear_8bit_act_xbit_weight/CMakeLists.txt

@@ -18,22 +18,18 @@ FetchContent_Declare(cpuinfo
 FetchContent_MakeAvailable(
     cpuinfo)
 
-FetchContent_Declare(
-  glog
-  GIT_REPOSITORY https://github.com/google/glog.git
-  GIT_TAG v0.7.1
-)
-FetchContent_MakeAvailable(glog)
 
 find_package(Torch REQUIRED)
 add_library(torchao_ops_linear_8bit_act_xbit_weight_aten OBJECT
     linear_8bit_act_xbit_weight.cpp
     op_linear_8bit_act_xbit_weight_aten.cpp
 )
 target_link_torchao_parallel_backend(torchao_ops_linear_8bit_act_xbit_weight_aten aten_openmp)
-target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE torchao_kernels_aarch64)
+
+if(TORCHAO_BUILD_CPU_AARCH64)
+    target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE torchao_kernels_aarch64)
+endif()
 target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE cpuinfo)
-# target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE glog)
 target_include_directories(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE "${TORCH_INCLUDE_DIRS}")
 target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE "${TORCH_LIBRARIES}")
 target_compile_definitions(torchao_ops_linear_8bit_act_xbit_weight_aten PRIVATE USE_ATEN=1)
@@ -55,7 +51,8 @@ if(TORCHAO_BUILD_EXECUTORCH_OPS)
     target_include_directories(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE "${EXECUTORCH_INCLUDE_DIRS}")
     target_compile_definitions(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE USE_EXECUTORCH=1)
     target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE "${EXECUTORCH_LIBRARIES}")
-    target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE torchao_kernels_aarch64)
+    if(TORCHAO_BUILD_CPU_AARCH64)
+        target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE torchao_kernels_aarch64)
+    endif()
     target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE cpuinfo)
-    # target_link_libraries(torchao_ops_linear_8bit_act_xbit_weight_executorch PRIVATE glog)
 endif()