Build custom portable ops

.lintrunner.toml

@@ -221,6 +221,8 @@ exclude_patterns = [
     '**/*.gif',
     'extension/llm/tokenizers',
     'extension/llm/tokenizers/**',
+    'kernels/portable',
+    'kernels/quantized',
     'examples/cuda',
     # File contains @generated
     'extension/llm/custom_ops/spinquant/fast_hadamard_transform_special.h',

CMakeLists.txt

@@ -1068,6 +1068,10 @@ if(EXECUTORCH_BUILD_KERNELS_QUANTIZED)
   list(APPEND _executorch_kernels quantized_ops_lib)
 endif()
 
+if(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/kernels/portable/custom_ops)
+endif()
+
 if(EXECUTORCH_BUILD_VULKAN)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/vulkan)
   list(APPEND _executorch_backends vulkan_backend vulkan_schema)

exir/tests/test_memory_planning.py

@@ -11,6 +11,7 @@
 from typing import Any, Callable, List, Optional, Tuple, Type
 
 import executorch.exir as exir
+import executorch.kernels.portable  # noqa: F401
 import torch
 from executorch.exir import ExecutorchBackendConfig, to_edge
 from executorch.exir.capture._capture import patch_forward
@@ -59,22 +60,6 @@
 from torch.nn import functional as F
 from torch.utils import _pytree as pytree
 
-try:
-    torch.ops.load_library("//executorch/kernels/portable:custom_ops_generated_lib")
-except (OSError, RuntimeError):
-    # When running outside of Buck (e.g., CMake/pip), find the shared library
-    # by globbing relative to the kernels/portable directory.
-    from pathlib import Path
-
-    _libs = list(
-        Path(__file__)
-        .parent.parent.parent.resolve()
-        .glob("**/kernels/portable/**/*custom_ops_generated_lib.*")
-    )
-    if _libs:
-        torch.ops.load_library(str(_libs[0]))
-    del Path
-
 
 def swap_modules(
     module: torch.nn.Module,

kernels/portable/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+try:
+    from pathlib import Path
+
+    libs = list(
+        Path(__file__).parent.resolve().glob("**/*portable_custom_ops_aot_lib.*")
+    )
+    del Path
+    assert len(libs) == 1, f"Expected 1 library but got {len(libs)}"
+    import torch as _torch
+
+    _torch.ops.load_library(libs[0])
+    del _torch
+except:  # noqa: E722
+    import logging
+
+    logging.info("portable_custom_ops_aot_lib is not loaded")
+    del logging

kernels/portable/custom_ops/CMakeLists.txt

@@ -0,0 +1,54 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Build the portable custom ops AOT library for registering custom ops into
+# PyTorch. This is separated from kernels/portable/CMakeLists.txt because it
+# requires find_package(Torch), which is not available when the portable kernels
+# are first processed.
+
+cmake_minimum_required(VERSION 3.19)
+
+if(NOT EXECUTORCH_ROOT)
+  set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)
+endif()
+
+include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake)
+include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake)
+
+if(NOT CMAKE_GENERATOR STREQUAL "Xcode")
+  if(NOT EXECUTORCH_BUILD_ARM_BAREMETAL)
+    set(_custom_ops_yaml "${CMAKE_CURRENT_SOURCE_DIR}/../custom_ops.yaml")
+    set(_portable_custom_ops "aten::allclose.out" "aten::allclose.Tensor")
+    gen_selected_ops(
+      LIB_NAME "portable_custom_ops_aot_lib" ROOT_OPS ${_portable_custom_ops}
+    )
+    generate_bindings_for_kernels(
+      LIB_NAME "portable_custom_ops_aot_lib" CUSTOM_OPS_YAML
+      "${_custom_ops_yaml}"
+    )
+    set(_portable_custom_ops_sources
+        "${CMAKE_CURRENT_SOURCE_DIR}/../cpu/op_allclose.cpp"
+        "${EXECUTORCH_ROOT}/runtime/core/exec_aten/util/tensor_util_aten.cpp"
+    )
+    gen_custom_ops_aot_lib(
+      LIB_NAME "portable_custom_ops_aot_lib" KERNEL_SOURCES
+      "${_portable_custom_ops_sources}"
+    )
+    target_include_directories(
+      portable_custom_ops_aot_lib PRIVATE ${_common_include_directories}
+    )
+
+    if(APPLE)
+      set(RPATH "@loader_path/../../../extensions/pybindings")
+    else()
+      set(RPATH "$ORIGIN/../../../extensions/pybindings")
+    endif()
+    set_target_properties(
+      portable_custom_ops_aot_lib PROPERTIES BUILD_RPATH ${RPATH} INSTALL_RPATH
+                                                                  ${RPATH}
+    )
+  endif()
+endif()

src/executorch/kernels/portable

@@ -0,0 +1 @@
+../../../kernels/portable

tools/cmake/preset/default.cmake

@@ -78,6 +78,10 @@ define_overridable_option(
   EXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT
   "Build the optimized ops library for AOT export usage" BOOL OFF
 )
+define_overridable_option(
+  EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT
+  "Build the portable custom ops library for AOT export usage" BOOL OFF
+)
 define_overridable_option(
   EXECUTORCH_BUILD_EXTENSION_ASR_RUNNER "Build the ASR runner extension" BOOL
   OFF

tools/cmake/preset/pybind.cmake

@@ -7,6 +7,7 @@
 set_overridable_option(EXECUTORCH_BUILD_PYBIND ON)
 set_overridable_option(EXECUTORCH_BUILD_KERNELS_QUANTIZED ON)
 set_overridable_option(EXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT ON)
+set_overridable_option(EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT ON)
 # Enable logging even when in release mode. We are building for desktop, where
 # saving a few kB is less important than showing useful error information to
 # users.