Ship the TorchAO kernels in the wheel

.ci/scripts/wheel/test_shared_libraries.py

@@ -156,6 +156,10 @@
 _MLX_SYMBOLS = ("executorch::backends::mlx::mutable_state_note_handle",)
 _BUNDLED_MLX_SYMBOLS = ("mlx::core::allocator::free",)
 
+# 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.
+_TORCHAO_KERNEL_SYMBOLS = ("torchao::quantization::get_qvals_range",)
+
 # A representative symbol from the profiler. A second definer means two event
 # tracers, so a trace records only part of what ran.
 _ETDUMP_SYMBOLS = ("executorch::etdump::ETDumpGen::ETDumpGen",)
@@ -963,6 +967,12 @@ def _resolve_required(required):
         _library_file_name("libexecutorch_kernels_quantized"),
         True,
     ),
+    (
+        "set of TorchAO kernels",
+        _TORCHAO_KERNEL_SYMBOLS,
+        _library_file_name("libexecutorch_kernels_torchao"),
+        False,
+    ),
     # The CUDA components. Required exactly when the wheel says it is a CUDA wheel,
     # which is decided at check time rather than here: a fixed False meant a wheel
     # tagged +cu126 carrying no CUDA library at all passed every check in this file,
@@ -2195,7 +2205,9 @@ def test_extension_contains_no_component() -> None:
     # Not every shipped library serves Python. The quantized kernels exist for a C++
     # application, since Python registers those operators through the torch-linked
     # ahead-of-time library at export time, and requiring a dependency would demand the
-    # extension link code it has no use for.
+    # extension link code it has no use for. The TorchAO kernels are in that same
+    # category: torchao registers its operators itself at export time, so the extension
+    # has no reason to link them either.
     #
     # The CUDA delegate is NOT in that category. The build deliberately links it into the
     # extension with a retention option, so it does carry a dependency, and excluding it
@@ -2205,7 +2217,10 @@ def test_extension_contains_no_component() -> None:
     expected = {
         name
         for name in shipped
-        if not any(marker in name for marker in ("kernels_quantized", "extension_cuda"))
+        if not any(
+            marker in name
+            for marker in ("kernels_quantized", "kernels_torchao", "extension_cuda")
+        )
     }
     unused = sorted(expected - needed)
     assert not unused, (
@@ -2313,6 +2328,7 @@ def test_shipped_library_names_are_expected() -> None:
         "libexecutorch",
         "libexecutorch_kernels_optimized",
         "libexecutorch_kernels_quantized",
+        "libexecutorch_kernels_torchao",
         "libexecutorch_backend_cuda",
         "libexecutorch_extension_cuda",
         # The same library under the name a non-shared build gives it. The shared

CMakeLists.txt

@@ -1130,6 +1130,25 @@ if(EXECUTORCH_BUILD_KERNELS_TORCHAO)
   executorch_target_link_options_shared_lib(torchao_ops_executorch)
   list(APPEND _executorch_kernels torchao_ops_executorch)
 
+  # The wheel ships one shared library per component so a process holds a single
+  # copy of each. These kernels are static archives declared in the TorchAO
+  # submodule, so wrap them the way the quantized kernels are wrapped rather
+  # than changing a third-party target. extension_threadpool is named explicitly
+  # because the kernels call into it and the wheel ships it as its own library,
+  # so this resolves against that one copy instead of bundling a second.
+  if(EXECUTORCH_BUILD_SHARED)
+    executorch_add_shared_library(
+      executorch_kernels_torchao torchao_ops_executorch extension_threadpool
+      executorch_shared
+    )
+    # Holds registration constructors and no symbol a consumer names, so a link
+    # with --as-needed would drop it and register nothing.
+    executorch_target_link_options_shared_lib(executorch_kernels_torchao)
+    # Ships beside the runtime in the wheel's lib/ directory, and links it, so
+    # it has to be able to find it from wherever the package is installed.
+    executorch_target_shipped_runtime_path(executorch_kernels_torchao)
+  endif()
+
   install(
     TARGETS torchao_ops_executorch torchao_kernels_aarch64
     EXPORT ExecuTorchTargets

docs/source/using-executorch-cpp.md

@@ -178,6 +178,7 @@ These are the components the package provides:
 | `threadpool` | Multi-threaded execution. | Linux, macOS |
 | `etdump` | Profiling, to record what ran and how long it took. | Linux, macOS |
 | `kernels_quantized` | The quantized operator kernels | Linux, macOS |
+| `kernels_torchao` | The TorchAO low-bit quantized kernels | macOS, Apple Silicon |
 | `backend_cuda` | The CUDA delegate | Linux |
 | `extension_cuda` | The CUDA stream extension | Linux |
 | `backend_openvino` | The OpenVINO delegate | Linux |
@@ -188,8 +189,9 @@ To see what your own install offers, ask CMake:
 ```cmake
 find_package(executorch REQUIRED)
 foreach(_component
-        runtime kernels_optimized kernels_quantized backend_xnnpack
-        backend_mlx backend_cuda extension_cuda backend_openvino threadpool etdump)
+        runtime kernels_optimized kernels_quantized kernels_torchao
+        backend_xnnpack backend_mlx backend_cuda extension_cuda
+        backend_openvino threadpool etdump)
   if(TARGET executorch::${_component})
     message(STATUS "have ${_component}")
   endif()

setup.py

@@ -2178,6 +2178,8 @@ def run(self):  # noqa C901
                     cmake_build_args += ["--target", "optimized_native_cpu_ops_lib"]
                 if cmake_cache.is_enabled("EXECUTORCH_BUILD_KERNELS_QUANTIZED"):
                     cmake_build_args += ["--target", "executorch_quantized_ops"]
+                if cmake_cache.is_enabled("EXECUTORCH_BUILD_KERNELS_TORCHAO"):
+                    cmake_build_args += ["--target", "executorch_kernels_torchao"]
                 if cmake_cache.is_enabled("EXECUTORCH_BUILD_XNNPACK"):
                     cmake_build_args += ["--target", "xnnpack_backend"]
 
@@ -2350,6 +2352,20 @@ def run(self):  # noqa C901
                         "EXECUTORCH_BUILD_KERNELS_QUANTIZED",
                     ],
                 ),
+                # The TorchAO kernels, so a C++ application running a model that
+                # uses them can link them from the wheel. The Apple framework build
+                # already ships these; this is the wheel's equivalent. Written to
+                # the cache root because the wrapper target is declared there.
+                BuiltFile(
+                    src_dir="%CMAKE_CACHE_DIR%/",
+                    src_name=get_dynamic_lib_name("executorch_kernels_torchao"),
+                    dst="executorch/lib/"
+                    + get_dynamic_lib_name("executorch_kernels_torchao"),
+                    dependent_cmake_flags=[
+                        "EXECUTORCH_BUILD_SHARED",
+                        "EXECUTORCH_BUILD_KERNELS_TORCHAO",
+                    ],
+                ),
                 # The OpenVINO delegate, so a C++ application can link it from the
                 # wheel. Only the adapter ships here: the OpenVINO runtime itself is
                 # loaded at run time and comes from the openvino extra.

tools/cmake/executorch-wheel-config.cmake

@@ -74,6 +74,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::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
@@ -333,6 +334,7 @@ if(_executorch_runtime_library AND NOT _executorch_targets_supported)
   foreach(
     _executorch_component IN
     ITEMS libexecutorch_kernels_optimized
+          libexecutorch_kernels_torchao
           libexecutorch_backend_xnnpack
           libexecutorch_backend_mlx
           libexecutorch_backend_cuda
@@ -590,6 +592,9 @@ _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.
+_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,9 +59,13 @@ 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)
-  # MLX requires Apple Silicon (ARM64) and the Metal compiler (xcrun -sdk macosx
-  # metal) which is only available with Xcode, not Command Line Tools
+  # 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.
   if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+    set_overridable_option(EXECUTORCH_BUILD_KERNELS_TORCHAO ON)
     execute_process(
       COMMAND xcrun -sdk macosx --find metal
       RESULT_VARIABLE _metal_compiler_result