Add a pure python wrapper fo pybindings.{aten,portable}_lib

Summary:
When installed as a pip wheel, we must import `torch` before trying to import the pybindings shared library extension. This will load libtorch.so and related libs, ensuring that the pybindings lib can resolve those runtime dependencies.

So, add a pure python wrapper that lets us do this when users say `import executorch.extension.pybindings.portable_lib`

We only need this for OSS, so don't bother doing this for other pybindings targets.

Differential Revision: D56317150

Author: dbort

CMakeLists.txt

@@ -581,15 +581,16 @@ if(EXECUTORCH_BUILD_PYBIND)
   target_compile_options(util PUBLIC ${_pybind_compile_options})
   target_link_libraries(util PRIVATE torch c10 executorch)
 
-  # pybind portable_lib
-  pybind11_add_module(portable_lib extension/pybindings/pybindings.cpp)
-  target_compile_definitions(portable_lib
-                             PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=portable_lib)
-  target_include_directories(portable_lib PRIVATE ${TORCH_INCLUDE_DIRS})
-  target_compile_options(portable_lib PUBLIC ${_pybind_compile_options})
-  target_link_libraries(portable_lib PUBLIC ${_dep_libs})
-
-  install(TARGETS portable_lib
+  # pybind _portable_lib. Note that the leading underscore is important, because
+  # it needs to coexist with portable_lib.py in the same python package.
+  pybind11_add_module(_portable_lib extension/pybindings/pybindings.cpp)
+  target_compile_definitions(_portable_lib
+                             PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=_portable_lib)
+  target_include_directories(_portable_lib PRIVATE ${TORCH_INCLUDE_DIRS})
+  target_compile_options(_portable_lib PUBLIC ${_pybind_compile_options})
+  target_link_libraries(_portable_lib PUBLIC ${_dep_libs})
+
+  install(TARGETS _portable_lib
           LIBRARY DESTINATION executorch/extension/pybindings)
 endif()
 

extension/pybindings/TARGETS

@@ -29,10 +29,10 @@ runtime.genrule(
     name = "pybindings_types_gen",
     srcs = [":pybinding_types"],
     outs = {
-        "aten_lib.pyi": ["aten_lib.pyi"],
-        "portable_lib.pyi": ["portable_lib.pyi"],
+        "_aten_lib.pyi": ["_aten_lib.pyi"],
+        "_portable_lib.pyi": ["_portable_lib.pyi"],
     },
-    cmd = "cp $(location :pybinding_types)/* $OUT/portable_lib.pyi && cp $(location :pybinding_types)/* $OUT/aten_lib.pyi",
+    cmd = "cp $(location :pybinding_types)/* $OUT/_portable_lib.pyi && cp $(location :pybinding_types)/* $OUT/_aten_lib.pyi",
     visibility = ["//executorch/extension/pybindings/..."],
 )
 
@@ -46,15 +46,29 @@ executorch_pybindings(
 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]"],
+    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]"],
+    python_module_name = "_aten_lib",
+    types = ["//executorch/extension/pybindings:pybindings_types_gen[_aten_lib.pyi]"],
     visibility = ["PUBLIC"],
 )
+
+runtime.python_library(
+    name = "portable_lib",
+    srcs = ["portable_lib.py"],
+    visibility = ["@EXECUTORCH_CLIENTS"],
+    deps = [":_portable_lib"],
+)
+
+runtime.python_library(
+    name = "aten_lib",
+    srcs = ["aten_lib.py"],
+    visibility = ["@EXECUTORCH_CLIENTS"],
+    deps = [":_aten_lib"],
+)

extension/pybindings/aten_lib.py

@@ -0,0 +1,22 @@
+# 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.
+
+# pyre-strict
+
+# When installed as a pip wheel, we must import `torch` before trying to import
+# the pybindings shared library extension. This will load libtorch.so and
+# related libs, ensuring that the pybindings lib can resolve those runtime
+# dependencies.
+import torch
+
+# Import everything from the shared library extension into our namespace.
+from executorch.extension.pybindings._aten_lib import *
+
+# Let users import everything from _aten_lib as if our current file defined
+# them.
+from executorch.extension.pybindings import _aten_lib
+__all__ = [name for name in dir(_aten_lib) if not name.startswith('_')]
+del _aten_lib

extension/pybindings/portable_lib.py

@@ -0,0 +1,22 @@
+# 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.
+
+# pyre-strict
+
+# When installed as a pip wheel, we must import `torch` before trying to import
+# the pybindings shared library extension. This will load libtorch.so and
+# related libs, ensuring that the pybindings lib can resolve those runtime
+# dependencies.
+import torch
+
+# Import everything from the shared library extension into our namespace.
+from executorch.extension.pybindings._portable_lib import *
+
+# Let users import everything from _portable_lib as if our current file defined
+# them.
+from executorch.extension.pybindings import _portable_lib
+__all__ = [name for name in dir(_portable_lib) if not name.startswith('_')]
+del _portable_lib

setup.py

@@ -376,13 +376,13 @@ def run(self):
             cmake_args += [
                 "-DEXECUTORCH_BUILD_PYBIND=ON",
             ]
-            build_args += ["--target", "portable_lib"]
+            build_args += ["--target", "_portable_lib"]
             if ShouldBuild.xnnpack:
                 cmake_args += [
                     "-DEXECUTORCH_BUILD_XNNPACK=ON",
                 ]
                 # No target needed; the cmake arg will link xnnpack
-                # into the portable_lib target.
+                # into the _portable_lib target.
             # TODO(dbort): Add MPS/CoreML backends when building on macos.
 
         if ShouldBuild.llama_custom_ops:
@@ -444,7 +444,7 @@ def get_ext_modules() -> list[Extension]:
             # portable kernels, and a selection of backends. This lets users
             # load and execute .pte files from python.
             BuiltExtension(
-                "portable_lib.*", "executorch.extension.pybindings.portable_lib"
+                "_portable_lib.*", "executorch.extension.pybindings._portable_lib"
             )
         )
     if ShouldBuild.llama_custom_ops: