Add buck2 example for custom ops

larryliu0820 · facebook-github-bot · commit 95810861b346 · 2023-08-16T13:21:02.000-07:00
Summary:
Similar to D48103652 and D48184410 but focusing on buck2 support.

This example demonstrates the ability to build a `exir_custom_ops_aot_lib` by using only C++ implementations of both functional and out variants of the custom op.

Specifically, let's assume a user has an op `my_ops::mul4` and already has a C++ implementation of it. The user registers this op through PyTorch C++ op registration API.

Now they want to run the same op on Executorch, naturally they need `my_ops::mul4.out` and a C++ implementation of it so that Executorch runtime can consume. `exir_custom_ops_aot_lib` is able to reuse this C++ impl for Executorch runtime, link it with `libtorch` and generate code to register it into PyTorch JIT runtime. This way the user doesn't have to write an extra Python impl.

This is the exact behavior demonstrated in D48184410, but instead of relying on CMake this diff enables BUCK2 targets.

Reviewed By: kimishpatel

Differential Revision: D48296936

fbshipit-source-id: 58ab7127f853549d29e7b25648dd7098219fe5c3
diff --git a/examples/custom_ops/custom_ops_2.py b/examples/custom_ops/custom_ops_2.py
@@ -8,11 +8,12 @@
 library that calls PyTorch C++ op registration API.
 """
 
+import argparse
+
 import torch
 
 from examples.export.export_example import export_to_pte
 
-
 # example model
 class Model(torch.nn.Module):
     def forward(self, a):
@@ -22,21 +23,21 @@ def forward(self, a):
 def main():
     m = Model()
     input = torch.randn(2, 3)
-    # load shared library
-    from sys import platform
-
-    if platform == "linux" or platform == "linux2":
-        extension = ".so"
-    elif platform == "darwin":
-        extension = ".dylib"
-    else:
-        raise RuntimeError(f"Unsupported platform {platform}")
-    torch.ops.load_library(
-        f"cmake-out/examples/custom_ops/libcustom_ops_aot_lib{extension}"
-    )
+
     # capture and lower
     export_to_pte("custom_ops_2", m, (input,))
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-s",
+        "--so_library",
+        required=True,
+        help="Provide path to so library. E.g., cmake-out/examples/custom_ops/libcustom_ops_aot_lib.so",
+    )
+    args = parser.parse_args()
+    torch.ops.load_library(args.so_library)
+    print(args.so_library)
+
     main()
diff --git a/examples/custom_ops/targets.bzl b/examples/custom_ops/targets.bzl
@@ -1,5 +1,5 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
-load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib")
+load("@fbsource//xplat/executorch/codegen:codegen.bzl", "et_operator_library", "executorch_generated_lib", "exir_custom_ops_aot_lib")
 
 def define_common_targets():
     """Defines targets that should be shared between fbcode and xplat.
@@ -54,3 +54,75 @@ def define_common_targets():
     )
 
     # ~~~ END of custom ops 1 `my_ops::mul3` library definitions ~~~
+    # ~~~ START of custom ops 2 `my_ops::mul4` library definitions ~~~
+
+    et_operator_library(
+        name = "select_custom_ops_2",
+        ops = [
+            "my_ops::mul4.out",
+        ],
+        define_static_targets = True,
+        visibility = [
+            "//executorch/codegen/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+    )
+
+    runtime.cxx_library(
+        name = "custom_ops_2",
+        srcs = ["custom_ops_2_out.cpp"],
+        deps = [
+            "//executorch/runtime/kernel:kernel_includes",
+        ],
+        visibility = [
+            "//executorch/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+    )
+
+    runtime.cxx_library(
+        name = "custom_ops_2_aten",
+        srcs = [
+            "custom_ops_2.cpp",
+            "custom_ops_2_out.cpp",
+        ],
+        deps = [
+            "//executorch/runtime/kernel:kernel_includes_aten",
+        ],
+        visibility = [
+            "//executorch/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+        external_deps = ["libtorch"],
+        # @lint-ignore BUCKLINT link_whole
+        link_whole = True,
+        # WARNING: using a deprecated API to avoid being built into a shared
+        # library. In the case of dynamically loading so library we don't want
+        # it to depend on other so libraries because that way we have to
+        # specify library directory path.
+        force_static = True,
+    )
+
+    exir_custom_ops_aot_lib(
+        name = "custom_ops_aot_lib_2",
+        yaml_target = ":custom_ops.yaml",
+        visibility = ["//executorch/..."],
+        kernels = [":custom_ops_2_aten"],
+        deps = [
+            ":select_custom_ops_2",
+        ],
+    )
+
+    executorch_generated_lib(
+        name = "lib_2",
+        deps = [
+            ":select_custom_ops_2",
+            ":custom_ops_2",
+        ],
+        custom_ops_yaml_target = ":custom_ops.yaml",
+        visibility = [
+            "//executorch/...",
+            "@EXECUTORCH_CLIENTS",
+        ],
+    )
+    # ~~~ END of custom ops 2 `my_ops::mul4` library definitions ~~~
diff --git a/examples/custom_ops/test_custom_ops.sh b/examples/custom_ops/test_custom_ops.sh
@@ -19,7 +19,7 @@ test_buck2_custom_op_1() {
 
   echo 'Running executor_runner'
   buck2 run //examples/executor_runner:executor_runner \
-      --config=executorch.register_custom_op_1=1 -- --model_path="./${model_name}.pte"
+      --config=executorch.register_custom_op=1 -- --model_path="./${model_name}.pte"
   # should give correct result
 
   echo "Removing ${model_name}.pte"
@@ -43,6 +43,38 @@ test_cmake_custom_op_1() {
   cmake-out/executor_runner --model_path="./${model_name}.pte"
 }
 
+test_buck2_custom_op_2() {
+  local model_name='custom_ops_2'
+
+  echo 'Building custom ops shared library'
+  SO_LIB=$(buck2 build //examples/custom_ops:custom_ops_aot_lib_2 --show-output | grep "buck-out" | cut -d" " -f2)
+
+  echo "Exporting ${model_name}.pte"
+  python3 -m "examples.custom_ops.${model_name}" --so_library="$SO_LIB"
+  # should save file custom_ops_2.pte
+
+  buck2 run //examples/executor_runner:executor_runner \
+      --config=executorch.register_custom_op=2 -- --model_path="./${model_name}.pte"
+  # should give correct result
+  echo "Removing ${model_name}.pte"
+  rm "./${model_name}.pte"
+}
+
+get_shared_lib_ext() {
+  UNAME=$(uname)
+  if [[ $UNAME == "Darwin" ]];
+  then
+    EXT=".dylib"
+  elif [[ $UNAME == "Linux" ]];
+  then
+    EXT=".so"
+  else
+    echo "Unsupported platform $UNAME"
+    exit 1
+  fi
+  echo $EXT
+}
+
 test_cmake_custom_op_2() {
   local model_name='custom_ops_2'
   SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
@@ -56,10 +88,11 @@ test_cmake_custom_op_2() {
       -DCMAKE_PREFIX_PATH="$CMAKE_PREFIX_PATH" ..)
 
   echo 'Building executor_runner'
-  cmake --build cmake-out -j9
+  cmake --build cmake-out -j4
 
+  EXT=$(get_shared_lib_ext)
   echo "Exporting ${model_name}.pte"
-  python3 -m "examples.custom_ops.${model_name}"
+  python3 -m "examples.custom_ops.${model_name}" --so_library="cmake-out/examples/custom_ops/libcustom_ops_aot_lib$EXT"
   # should save file custom_ops_2.pte
 
   echo 'Running executor_runner'
@@ -68,4 +101,5 @@ test_cmake_custom_op_2() {
 
 test_buck2_custom_op_1
 test_cmake_custom_op_1
+test_buck2_custom_op_2
 test_cmake_custom_op_2
diff --git a/examples/executor_runner/targets.bzl b/examples/executor_runner/targets.bzl
@@ -7,9 +7,14 @@ def define_common_targets():
     TARGETS and BUCK files that call this function.
     """
 
-    register_custom_op_1 = native.read_config("executorch", "register_custom_op_1", "0") == "1"
+    register_custom_op = native.read_config("executorch", "register_custom_op", "0")
 
-    custom_ops_lib = ["//executorch/examples/custom_ops:lib_1"] if register_custom_op_1 else []
+    if register_custom_op == "1":
+        custom_ops_lib = ["//executorch/examples/custom_ops:lib_1"]
+    elif register_custom_op == "2":
+        custom_ops_lib = ["//executorch/examples/custom_ops:lib_2"]
+    else:
+        custom_ops_lib = []
 
     # Test driver for models, uses all portable kernels.
     runtime.cxx_binary(
diff --git a/runtime/core/exec_aten/util/targets.bzl b/runtime/core/exec_aten/util/targets.bzl
@@ -61,4 +61,9 @@ def define_common_targets():
                 ":scalar_type_util" + aten_suffix,
                 ":dim_order_util" + aten_suffix,
             ],
+            # WARNING: using a deprecated API to avoid being built into a shared
+            # library. In the case of dynamically loading so library we don't want
+            # it to depend on other so libraries because that way we have to
+            # specify library directory path.
+            force_static = True,
         )
diff --git a/runtime/platform/targets.bzl b/runtime/platform/targets.bzl
@@ -50,6 +50,11 @@ def define_common_targets():
         visibility = [
             "//executorch/core/...",
         ],
+        # WARNING: using a deprecated API to avoid being built into a shared
+        # library. In the case of dynamically loading so library we don't want
+        # it to depend on other so libraries because that way we have to
+        # specify library directory path.
+        force_static = True,
     )
 
     # Interfaces for executorch users
@@ -78,6 +83,11 @@ def define_common_targets():
             "//executorch/...",
             "@EXECUTORCH_CLIENTS",
         ],
+        # WARNING: using a deprecated API to avoid being built into a shared
+        # library. In the case of dynamically loading so library we don't want
+        # it to depend on other so libraries because that way we have to
+        # specify library directory path.
+        force_static = True,
     )
 
     # Library for backend implementers to define implementations against.
diff --git a/shim/xplat/executorch/build/env_interface.bzl b/shim/xplat/executorch/build/env_interface.bzl
@@ -107,6 +107,11 @@ def _remove_platform_specific_args(kwargs):
         kwargs.pop(key)
     return kwargs
 
+def _remove_unsupported_kwargs(kwargs):
+    """Removes environment unsupported kwargs
+    """
+    return kwargs
+
 def _patch_headers(kwargs):
     """Patch (add or modify or remove) headers related attributes for this build environment.
     """
@@ -202,6 +207,7 @@ env = struct(
     # @lint-ignore BUCKLINT: native and fb_native are explicitly forbidden in fbcode.
     python_test = native.python_test,
     remove_platform_specific_args = _remove_platform_specific_args,
+    remove_unsupported_kwargs = _remove_unsupported_kwargs,
     resolve_external_dep = _resolve_external_dep,
     target_needs_patch = _target_needs_patch,
     EXTERNAL_DEP_FALLTHROUGH = _EXTERNAL_DEP_FALLTHROUGH,
diff --git a/shim/xplat/executorch/build/runtime_wrapper.bzl b/shim/xplat/executorch/build/runtime_wrapper.bzl
@@ -208,6 +208,7 @@ def _patch_kwargs_common(kwargs):
     return kwargs
 
 def _patch_kwargs_cxx(kwargs):
+    env.remove_unsupported_kwargs(kwargs)
     env.patch_platforms(kwargs)
     env.remove_platform_specific_args(kwargs)
     return _patch_kwargs_common(kwargs)
diff --git a/third-party/TARGETS b/third-party/TARGETS
@@ -218,35 +218,54 @@ runtime.cxx_binary(
 
 runtime.genrule(
     name = "libtorch_gen",
-    out = "libtorch.so",
+    outs = select({
+        "ovr_config//os:macos": {
+            "libtorch": ["libtorch.dylib"],
+            "libc10": ["libc10.dylib"],
+            "libtorch_cpu": ["libtorch_cpu.dylib"],
+            "include": ["include"],
+        },
+        "DEFAULT": {
+            "libtorch": ["libtorch.so"],
+            "libc10": ["libc10.so"],
+            "libtorch_cpu": ["libtorch_cpu.so"],
+            "include": ["include"],
+        },
+    }),
+    default_outs = ["."],
     srcs = ["link_torch.sh"],
-    bash = 'bash $SRCS torch/lib/libtorch.so ${OUT}',
+    bash = select({
+        "ovr_config//os:macos": "bash $SRCS -f torch/lib/libtorch.dylib,torch/lib/libtorch_cpu.dylib,torch/lib/libc10.dylib,torch/include -o ${OUT}",
+        "DEFAULT": "bash $SRCS -f torch/lib/libtorch.so,torch/lib/libtorch_cpu.so,torch/lib/libc10.so,torch/include -o ${OUT}",
+    }),
 )
 
-runtime.genrule(
-    name = "torch_headers_gen",
-    out = "include",
-    srcs = ["link_torch.sh"],
-    bash = 'bash $SRCS torch/include ${OUT}',
+prebuilt_cxx_library(
+    name = "libc10",
+    shared_lib = ":libtorch_gen[libc10]",
 )
 
-runtime.genrule(
-    name = "libtorch_dir",
-    out = "lib",
-    srcs = ["link_torch.sh"],
-    bash = 'bash $SRCS torch/lib ${OUT}',
+prebuilt_cxx_library(
+    name = "libtorch_cpu",
+    shared_lib = ":libtorch_gen[libtorch_cpu]",
 )
 
 prebuilt_cxx_library(
     name = "libtorch",
-    shared_lib = ":libtorch_gen",
-    soname = "libtorch.so",
+    shared_lib = ":libtorch_gen[libtorch]",
     exported_preprocessor_flags = [
-        "-D_GLIBCXX_USE_CXX11_ABI=0", # `libtorch` is built without CXX11_ABI so any target depends on it need to use the same build config.
-        "-I$(location :torch_headers_gen)", # include header directories
-        "-I$(location :torch_headers_gen)/torch/csrc/api/include", # include header directories
+        "-D_GLIBCXX_USE_CXX11_ABI=0",  # `libtorch` is built without CXX11_ABI so any target depends on it need to use the same build config.
+        "-I$(location :libtorch_gen[include])",  # include header directories
+        "-I$(location :libtorch_gen[include])/torch/csrc/api/include",  # include header directories
+    ],
+    exported_linker_flags = select({
+        "ovr_config//os:macos": ["-Xlinker", "-rpath", "$(location :libtorch_gen)", "-Xlinker"],
+        "DEFAULT": ["-Wl,-rpath,$(location :libtorch_gen)"],  # define rpath to locate shared library
+    }),
+    exported_headers = [":libtorch_gen[include]"],
+    exported_deps = [
+        ":libc10",
+        ":libtorch_cpu",
     ],
-    exported_linker_flags = ["-Wl,-rpath,$(location :libtorch_dir)"], # define rpath to locate shared library
-    exported_headers = [":torch_headers_gen"],
     visibility = ["PUBLIC"],
 )
