Introduce pybind11 and build extension/pybindings/portable_lib (#332)

Summary:
Build Pybind11 shared library and load it into python environment. This way we can use `Module` python API to run end2end tests.

Add `test_end2end.py` in `pytest.ini` so that we have e2e test coverage.

Pull Request resolved: #332

Reviewed By: tarun292

Differential Revision: D49254089

Pulled By: larryliu0820

fbshipit-source-id: cc3698bdd6620176d64e14027a37ed5b9897704a

Author: larryliu0820

.github/workflows/_unittest.yml

@@ -36,6 +36,22 @@ jobs:
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-linux.sh "${BUILD_TOOL}"
 
+        # Build //extension/pybindings:portable_lib. Example path:
+        # buck-out/v2/gen/root/524f8da68ea2a374/extension/pybindings/__portable_lib__/portable_lib.so
+        SO_LIB_DIR=$(buck2 build //extension/pybindings:portable_lib --show-output | cut -d' ' -f2 | xargs dirname)
+
+        # Let LD_LIBRARY_PATH include libtorch_python directory
+        PYTHON_LIB_DIR=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+        export LD_LIBRARY_PATH="${PYTHON_LIB_DIR}/torch/lib"
+
+        # Let PYTHONPATH include the output directory so that portable_lib.so can be loaded into Python.
+        export PYTHONPATH="${SO_LIB_DIR}:${PYTHON_LIB_DIR}/torch/lib"
+
+        # Generate a shim file in extension/pybindings/portable_lib/
+        SHIM_PY="${PYTHON_LIB_DIR}/executorch/extension/pybindings/portable_lib.py"
+        touch ${SHIM_PY}
+        echo "from portable_lib import _load_for_executorch_from_buffer,_load_bundled_program_from_buffer,_load_for_executorch_from_bundled_program" > ${SHIM_PY}
+
         # Run pytest with coverage
         pytest -n auto --cov=./ --cov-report=xml
 
@@ -59,6 +75,22 @@ jobs:
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python bash .ci/scripts/setup-macos.sh "${BUILD_TOOL}"
 
+        # Build //extension/pybindings:portable_lib. Example path:
+        # buck-out/v2/gen/root/524f8da68ea2a374/extension/pybindings/__portable_lib__/portable_lib.dylib
+        SO_LIB_DIR=$(buck2 build //extension/pybindings:portable_lib --show-output | cut -d' ' -f2 | xargs dirname)
+
+        # Let LD_LIBRARY_PATH include libtorch_python directory
+        PYTHON_LIB_DIR=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')
+        export LD_LIBRARY_PATH="${PYTHON_LIB_DIR}/torch/lib"
+
+        # Let PYTHONPATH include the output directory so that portable_lib.dylib can be loaded into Python.
+        export PYTHONPATH="${SO_LIB_DIR}:${PYTHON_LIB_DIR}/torch/lib"
+
+        # Generate a shim file in extension/pybindings/portable_lib/
+        SHIM_PY="${PYTHON_LIB_DIR}/executorch/extension/pybindings/portable_lib.py"
+        touch ${SHIM_PY}
+        echo "from portable_lib import _load_for_executorch_from_buffer,_load_bundled_program_from_buffer,_load_for_executorch_from_bundled_program" > ${SHIM_PY}
+
         # Run pytest with coverage
         pytest -n auto --cov=./ --cov-report=xml
         # Run gtest

backends/transforms/TARGETS

@@ -1,7 +1,6 @@
-#load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
-load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
-python_library(
+runtime.python_library(
     name = "lib",
     srcs = [
         "__init__.py",
@@ -14,9 +13,12 @@ python_library(
     ],
 )
 
-python_library(
+runtime.python_library(
     name = "addmm_mm_to_linear",
     srcs = ["addmm_mm_to_linear.py"],
+    visibility = [
+        "//executorch/backends/...",
+    ],
     deps = [
         "//caffe2:torch",
         "//executorch/exir:pass_base",

backends/xnnpack/test/TARGETS

@@ -1,9 +1,9 @@
-load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 load(":targets.bzl", "define_common_targets")
 
 define_common_targets()
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack",
     srcs = [
         "test_xnnpack.py",
@@ -16,7 +16,6 @@ python_unittest(
     ],
     tags = ["long_running"],
     deps = [
-        "//caffe2:torch",
         "//executorch/backends/xnnpack:xnnpack_backend",
         "//executorch/backends/xnnpack:xnnpack_preprocess",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
@@ -30,9 +29,12 @@ python_unittest(
         "//executorch/sdk/bundled_program:core",
         "//executorch/sdk/bundled_program/serialize:lib",
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack_quantized",
     srcs = [
         "test_xnnpack_quantized.py",
@@ -44,7 +46,6 @@ python_unittest(
     ],
     tags = ["long_running"],
     deps = [
-        "//caffe2:torch",
         "//executorch/backends/xnnpack:xnnpack_backend",
         "//executorch/backends/xnnpack:xnnpack_preprocess",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
@@ -59,9 +60,12 @@ python_unittest(
         "//executorch/sdk/bundled_program:core",
         "//executorch/sdk/bundled_program/serialize:lib",
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack_quantized_models",
     srcs = [
         "test_xnnpack_quantized_models.py",
@@ -73,7 +77,6 @@ python_unittest(
     ],
     tags = ["long_running"],
     deps = [
-        "//caffe2:torch",
         "//executorch/backends/xnnpack:xnnpack_backend",
         "//executorch/backends/xnnpack:xnnpack_preprocess",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
@@ -88,9 +91,12 @@ python_unittest(
         "//executorch/sdk/bundled_program/serialize:lib",
         "//pytorch/vision:torchvision",
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack_passes",
     srcs = glob([
         "passes/*.py",
@@ -99,7 +105,6 @@ python_unittest(
         "test_xnnpack_utils_classes.py",
     ],
     deps = [
-        "//caffe2:torch",
         "//executorch/backends/xnnpack/passes:xnnpack_passes",
         "//executorch/backends/xnnpack/test/tester:tester",
         "//executorch/backends/xnnpack/utils:xnnpack_utils",
@@ -108,21 +113,26 @@ python_unittest(
         "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib",
         "//executorch/exir/dialects:lib",
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack_ops",
     srcs = glob([
         "ops/*.py",
     ]),
     deps = [
-        "//caffe2:torch",
         "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
         "//executorch/backends/xnnpack/test/tester:tester",
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )
 
-python_unittest(
+runtime.python_test(
     name = "test_xnnpack_models",
     srcs = glob([
         "models/*.py",
@@ -131,8 +141,10 @@ python_unittest(
     deps = [
         "fbsource//third-party/pypi/timm:timm",
         "fbsource//third-party/pypi/torchsr:torchsr",  # @manual
-        "//caffe2:torch",
         "//executorch/backends/xnnpack/test/tester:tester",
         "//pytorch/vision:torchvision",  # @manual
     ],
+    external_deps = [
+        "libtorch",
+    ],
 )

extension/aten_util/targets.bzl

@@ -29,13 +29,7 @@ def define_common_targets():
             "//executorch/runtime/core:core",
             "//executorch/runtime/core/exec_aten:lib",
         ],
-        fbcode_deps = [
-            "//caffe2:ATen-core",
-            "//caffe2:ATen-cpu",
-            "//caffe2/c10:c10",
-        ],
-        xplat_deps = [
-            "//xplat/caffe2:torch_mobile_core",
-            "//xplat/caffe2/c10:c10",
+        external_deps = [
+            "torch-core-cpp",
         ],
     )

extension/pybindings/TARGETS

@@ -2,8 +2,8 @@
 # Any targets that should be shared between fbcode and xplat must be defined in
 # targets.bzl. This file can contain fbcode-only targets.
 
-load("@fbcode//executorch/extension/pybindings:pybindings.bzl", "ATEN_MODULE_DEPS", "MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB", "MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB", "PORTABLE_MODULE_DEPS", "executorch_pybindings")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+load("@fbsource//xplat/executorch/extension/pybindings:pybindings.bzl", "ATEN_MODULE_DEPS", "MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB", "MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB", "PORTABLE_MODULE_DEPS", "executorch_pybindings")
 
 # Export this so the internal fb/ subdir can create pybindings with custom internal deps
 # without forking the pybinding source.
@@ -36,19 +36,22 @@ runtime.genrule(
 )
 
 executorch_pybindings(
+    compiler_flags = ["-std=c++17"],
     cppdeps = PORTABLE_MODULE_DEPS,
     python_module_name = "core",
     visibility = ["PUBLIC"],
 )
 
 executorch_pybindings(
+    compiler_flags = ["-std=c++17"],
     cppdeps = PORTABLE_MODULE_DEPS + MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB,
     python_module_name = "portable_lib",
     types = ["//executorch/extension/pybindings:pybindings_types_gen[portable_lib.pyi]"],
     visibility = ["PUBLIC"],
 )
 
 executorch_pybindings(
+    compiler_flags = ["-std=c++17"],
     cppdeps = ATEN_MODULE_DEPS + MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB,
     python_module_name = "aten_lib",
     types = ["//executorch/extension/pybindings:pybindings_types_gen[aten_lib.pyi]"],

pytest.ini

@@ -36,6 +36,8 @@ addopts =
     kernels/test/test_case_gen.py
     # backends/arm
     backends/arm/test
+    # test
+    test/end2end/test_end2end.py
 
 # run the same tests multiple times to determine their
 # flakiness status. Default to 50 re-runs

shim/xplat/executorch/backends/backends.bzl

@@ -0,0 +1,20 @@
+def get_all_cpu_backend_targets():
+    """Returns a list of all CPU backend targets.
+
+    For experimenting and testing, not for production, since it will typically
+    include more than necessary for a particular product.
+    """
+    return [
+        "//executorch/backends/xnnpack:xnnpack_backend",
+    ]
+
+def get_all_cpu_aot_and_backend_targets():
+    """Returns a list of all CPU backend targets with aot (ahead of time).
+
+    For experimenting and testing, not for production, since it will typically
+    include more than necessary for a particular product.
+    """
+    return [
+        "//executorch/backends/xnnpack:xnnpack_preprocess",
+        "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
+    ] + get_all_cpu_backend_targets()

shim/xplat/executorch/build/env_interface.bzl

@@ -39,10 +39,11 @@ _EXTERNAL_DEPS = {
     "gtest": "//third-party:gtest",
     "gtest_aten": "//third-party:gtest_aten",
     "libtorch": "//third-party:libtorch",
+    "libtorch_python": "//third-party:libtorch_python",
     "prettytable": "//third-party:prettytable",
-    "pybind11": [],  # TODO(larryliu0820): Add support
+    "pybind11": "//third-party:pybind11",
     # Core C++ PyTorch functionality like Tensor and ScalarType.
-    "torch-core-cpp": [],  # TODO(larryliu0820): Add support
+    "torch-core-cpp": "//third-party:libtorch",
     "torchgen": "//third-party:torchgen",
 }
 
@@ -95,6 +96,13 @@ def _patch_deps(kwargs, dep_type):
 def _patch_platform_build_mode_flags(kwargs):
     return kwargs
 
+def _patch_force_static(kwargs):
+    """For OSS cxx library, force static linkage unless specify otherwise.
+    """
+    if "force_static" not in kwargs:
+        kwargs["force_static"] = True
+    return kwargs
+
 def _remove_platform_specific_args(kwargs):
     """Removes platform specific arguments for BUCK builds
 
@@ -118,6 +126,7 @@ def _remove_platform_specific_args(kwargs):
 def _remove_unsupported_kwargs(kwargs):
     """Removes environment unsupported kwargs
     """
+    kwargs.pop("tags", None)  # tags = ["long_running"] doesn't work in oss
     kwargs.pop("types", None)  # will have to find a different way to handle .pyi files in oss
     return kwargs
 
@@ -204,6 +213,7 @@ env = struct(
     patch_deps = _patch_deps,
     patch_cxx_compiler_flags = _patch_cxx_compiler_flags,
     patch_executorch_genrule_cmd = _patch_executorch_genrule_cmd,
+    patch_force_static = _patch_force_static,
     patch_headers = _patch_headers,
     patch_platform_build_mode_flags = _patch_platform_build_mode_flags,
     patch_platforms = _patch_platforms,

shim/xplat/executorch/build/runtime_wrapper.bzl

@@ -230,6 +230,7 @@ def _cxx_library_common(*args, **kwargs):
     env.patch_headers(kwargs)
     env.patch_pp_flags(kwargs)
     env.patch_cxx_compiler_flags(kwargs)
+    env.patch_force_static(kwargs)
 
     env.cxx_library(*args, **kwargs)
 

shim/xplat/executorch/codegen/codegen.bzl

@@ -315,6 +315,7 @@ def exir_custom_ops_aot_lib(
             supports_python_dlopen = True,
             platforms = platforms,
             compiler_flags = compiler_flags,
+            force_static = False,
         )
 
 def executorch_generated_lib(

shim/xplat/executorch/extension/pybindings/pybindings.bzl

@@ -7,9 +7,6 @@ MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB = [
 ]
 
 PORTABLE_MODULE_DEPS = [
-    "//caffe2:ATen",
-    "//caffe2:torch",
-    "//caffe2:torch_extension",
     "//executorch/runtime/kernel:operator_registry",
     "//executorch/runtime/executor:program",
     "//executorch/sdk/bundled_program/schema:bundled_program_schema_fbs",
@@ -32,9 +29,6 @@ ATEN_MODULE_DEPS = [
     "//executorch/extension/memory_allocator:malloc_memory_allocator",
     "//executorch/util:read_file",
     "//executorch/sdk/bundled_program:runtime_aten",
-    "//caffe2:torch",
-    "//caffe2:torch_extension",
-    "//caffe2:ATen",
     "//executorch/runtime/executor/test:test_backend_compiler_lib_aten",
 ]
 
@@ -44,14 +38,15 @@ MODELS_ATEN_OPS_ATEN_MODE_GENERATED_LIB = [
     "//executorch/kernels/aten:generated_lib_aten",
 ]
 
-def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibility = ["//executorch/..."], types = []):
+def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibility = ["//executorch/..."], types = [], compiler_flags = []):
     runtime.cxx_python_extension(
         name = python_module_name,
         srcs = [
             "//executorch/extension/pybindings:pybindings.cpp",
         ] + srcs,
         types = types,
         base_module = "executorch.extension.pybindings",
+        compiler_flags = compiler_flags,
         preprocessor_flags = [
             "-DEXECUTORCH_PYTHON_MODULE_NAME={}".format(python_module_name),
         ],
@@ -61,6 +56,7 @@ def executorch_pybindings(python_module_name, srcs = [], cppdeps = [], visibilit
         ] + cppdeps,
         external_deps = [
             "pybind11",
+            "libtorch_python",
         ],
         use_static_deps = True,
         _is_external_target = bool(visibility != ["//executorch/..."]),

test/end2end/TARGETS

@@ -55,7 +55,6 @@ python_unittest(
         "//executorch/exir/tests:transformer",
         "//executorch/extension/pybindings:aten_lib",
         "//executorch/extension/pytree:pybindings",
-        "//executorch/kernels/portable:custom_ops_generated_lib",
         "//executorch/sdk/bundled_program:config",
         "//executorch/sdk/bundled_program:core",
         "//executorch/sdk/bundled_program/serialize:lib",
@@ -87,7 +86,6 @@ python_unittest(
         "//executorch/exir/tests:transformer",
         "//executorch/extension/pybindings:portable_lib",
         "//executorch/extension/pytree:pybindings",
-        "//executorch/kernels/portable:custom_ops_generated_lib",
         "//executorch/sdk/bundled_program:config",
         "//executorch/sdk/bundled_program:core",
         "//executorch/sdk/bundled_program/serialize:lib",