Ship the TorchAO kernels on the Linux aarch64 wheel too

.ci/scripts/wheel/test_shared_libraries.py

@@ -162,8 +162,8 @@
     "executorch::backends::coreml::CoreMLBackendDelegate::get_registered_delegate()",
 )
 
-# A representative symbol from the TorchAO kernels. These are Apple Silicon only, so
-# most wheels ship no such library and the row below is not required.
+# A representative symbol from the TorchAO kernels. These need aarch64, so the x86 and Windows
+# wheels ship no such library and the row below is not required.
 _TORCHAO_KERNEL_SYMBOLS = ("torchao::quantization::get_qvals_range",)
 
 # A representative symbol from the profiler. A second definer means two event

tools/cmake/executorch-wheel-config.cmake

@@ -75,7 +75,7 @@
 # executorch::backend_mlx        The MLX delegate. macOS on Apple Silicon only.
 #                                Its Metal kernel archive is published as
 #                                MLX_METALLIB_PATH, see below.
-# executorch::kernels_torchao    The TorchAO kernels. macOS on Apple Silicon only.
+# executorch::kernels_torchao    The TorchAO kernels. aarch64 rows only.
 # executorch::backend_cuda       The CUDA delegate. Linux only.
 # executorch::extension_cuda     The CUDA stream extension. Linux only.
 # executorch::backend_openvino   The OpenVINO delegate. Linux only. Opens the
@@ -594,8 +594,8 @@ _executorch_define_component(threadpool executorch_threadpool)
 # checks, so it has to be defined here or a consumer following the documentation
 # gets a bare name that CMake hands to the linker as a literal flag.
 _executorch_define_component(kernels_optimized executorch_kernels_optimized)
-# The TorchAO kernels, present only in a wheel built for Apple Silicon, which is
-# the only architecture they build for.
+# The TorchAO kernels, present only in a wheel built for aarch64, which is the
+# only architecture they build for.
 _executorch_define_component(kernels_torchao executorch_kernels_torchao)
 # The quantized kernels, optional in the same way: a wheel built without them
 # simply has no such library and the component is not defined.

tools/cmake/preset/pybind.cmake

@@ -59,11 +59,10 @@ if(CMAKE_SYSTEM_NAME STREQUAL "Darwin")
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_TRAINING ON)
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER ON)
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM ON)
-  # Both of these are Apple Silicon only. The TorchAO kernels build only for
-  # aarch64, which is what TORCHAO_BUILD_CPU_AARCH64 selects; the Apple
-  # framework build already ships them and this brings the wheel in line. MLX
-  # additionally needs the Metal compiler (xcrun -sdk macosx metal), which comes
-  # with Xcode and not with the Command Line Tools.
+  # MLX needs the Metal compiler (xcrun -sdk macosx metal), which comes with
+  # Xcode and not with the Command Line Tools, so it is probed rather than
+  # assumed. The TorchAO kernels are enabled on every aarch64 row, here and
+  # under Linux below, since aarch64 is what they require.
   if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
     set_overridable_option(EXECUTORCH_BUILD_KERNELS_TORCHAO ON)
     execute_process(
@@ -89,6 +88,13 @@ elseif(CMAKE_SYSTEM_NAME STREQUAL "Linux")
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_TRAINING ON)
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER ON)
   set_overridable_option(EXECUTORCH_BUILD_EXTENSION_LLM ON)
+  # The same aarch64 kernels the macOS arm64 wheel gets above. The hardware is
+  # what these need, not Apple: they are selected by TORCHAO_BUILD_CPU_AARCH64
+  # and reach their NEON paths through TORCHAO_ENABLE_ARM_NEON_DOT, both of
+  # which hold here too.
+  if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$")
+    set_overridable_option(EXECUTORCH_BUILD_KERNELS_TORCHAO ON)
+  endif()
   if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|i.86)$")
     # Auto-enable QNN on Linux x86 when the SDK is available. - QNN_SDK_ROOT set
     # explicitly → always enable - GitHub Actions CI → skip (avoids flaky 1.3GB