Extend PyBundledModule with extension.BundledModule (#12839)

Summary: Pull Request resolved: #12839

Differential Revision: D78938344

Author: Gasoonjia

CMakeLists.txt

@@ -731,6 +731,11 @@ if(EXECUTORCH_BUILD_PYBIND)
       torch
   )
 
+  if(EXECUTORCH_BUILD_EXTENSION_MODULE)
+    # Always use static linking for pybindings to avoid runtime symbol resolution issues
+    list(APPEND _dep_libs extension_module_static)
+  endif()
+
   if(EXECUTORCH_BUILD_TESTS)
     list(APPEND _dep_libs test_backend_compiler_lib)
   endif()
@@ -788,6 +793,17 @@ if(EXECUTORCH_BUILD_PYBIND)
   target_compile_options(portable_lib PUBLIC ${_pybind_compile_options})
   target_link_libraries(portable_lib PRIVATE ${_dep_libs})
 
+  # Set RPATH for portable_lib to find libextension_module.so
+  if(APPLE)
+    set_target_properties(portable_lib PROPERTIES
+                          BUILD_RPATH "@loader_path"
+                          INSTALL_RPATH "@loader_path")
+  else()
+    set_target_properties(portable_lib PROPERTIES
+                          BUILD_RPATH "$ORIGIN"
+                          INSTALL_RPATH "$ORIGIN")
+  endif()
+
   install(
     TARGETS portable_lib
     EXPORT ExecuTorchTargets

devtools/bundled_program/test/test_end2end.py

@@ -5,21 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 # flake8: noqa: F401
-import functools
-import inspect
-import os
-import random
 import unittest
-from typing import Callable, Dict, Optional, Tuple, Type
-
-import executorch.exir as exir
-
-import executorch.exir.control_flow as control_flow
-
-# @manual=//executorch/extension/pytree:pybindings
-import executorch.extension.pytree as pytree
-
-import torch
 
 from executorch.devtools.bundled_program.core import BundledProgram
 from executorch.devtools.bundled_program.serialize import (
@@ -35,8 +21,6 @@
 try:
     from executorch.extension.pybindings.portable_lib import (
         _load_bundled_program_from_buffer,
-        _load_for_executorch_from_buffer,
-        _load_for_executorch_from_bundled_program,
     )
 
     kernel_mode = "lean"
@@ -47,8 +31,6 @@
 try:
     from executorch.extension.pybindings.aten_lib import (  # @manual=//executorch/extension/pybindings:aten_lib
         _load_bundled_program_from_buffer,
-        _load_for_executorch_from_buffer,
-        _load_for_executorch_from_bundled_program,
     )
 
     assert kernel_mode is None
@@ -75,19 +57,8 @@ def test_sample_model_e2e(self):
             bundled_program_buffer
         )
 
-        executorch_module = _load_for_executorch_from_bundled_program(
-            executorch_bundled_program
-        )
-
         for method_name in eager_model.method_names:
-            executorch_module.load_bundled_input(
-                executorch_bundled_program,
-                method_name,
-                0,
-            )
-            executorch_module.plan_execute(method_name)
-            executorch_module.verify_result_with_bundled_expected_output(
-                executorch_bundled_program,
+            executorch_bundled_program.verify_result_with_bundled_expected_output(
                 method_name,
                 0,
             )

extension/module/CMakeLists.txt

@@ -17,6 +17,21 @@ if(NOT EXECUTORCH_ROOT)
 endif()
 
 list(TRANSFORM _extension_module__srcs PREPEND "${EXECUTORCH_ROOT}/")
+
+# Only add bundled_module.cpp if devtools is being built (required for bundled_program schema)
+if(EXECUTORCH_BUILD_DEVTOOLS)
+  list(APPEND _extension_module__srcs "${EXECUTORCH_ROOT}/extension/module/bundled_module.cpp")
+  message(STATUS "extension_module: Including bundled_module.cpp (devtools enabled)")
+else()
+  message(STATUS "extension_module: Excluding bundled_module.cpp (devtools disabled)")
+endif()
+
+# Also add bundled_module sources if they exist (in case there are additional files)
+if(DEFINED _extension_bundled_module__srcs)
+  list(TRANSFORM _extension_bundled_module__srcs PREPEND "${EXECUTORCH_ROOT}/")
+  list(APPEND _extension_module__srcs ${_extension_bundled_module__srcs})
+endif()
+
 if(CMAKE_TOOLCHAIN_IOS
    OR CMAKE_TOOLCHAIN_ANDROID
    OR APPLE
@@ -27,31 +42,75 @@ if(CMAKE_TOOLCHAIN_IOS
 else()
   add_library(extension_module SHARED ${_extension_module__srcs})
 endif()
-target_link_libraries(
-  extension_module PRIVATE executorch_core extension_data_loader
-                           extension_flat_tensor
-)
+# Link libraries - only link bundled_program if devtools is enabled
+if(EXECUTORCH_BUILD_DEVTOOLS)
+  target_link_libraries(
+    extension_module PRIVATE executorch_core extension_data_loader
+                             extension_flat_tensor bundled_program
+  )
+  target_compile_definitions(extension_module PRIVATE EXECUTORCH_BUILD_DEVTOOLS)
+else()
+  target_link_libraries(
+    extension_module PRIVATE executorch_core extension_data_loader
+                             extension_flat_tensor
+  )
+endif()
 target_include_directories(
   extension_module PUBLIC ${_common_include_directories}
 )
 target_compile_options(
   extension_module PUBLIC -Wno-deprecated-declarations -fPIC
 )
 
+# Only set USE_ATEN_LIB if PyTorch is actually available
+# This prevents build failures in OSS CMake CI where ATen headers are not available
+if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
+  # Try to find PyTorch headers to check if ATen is available
+  find_package(Torch CONFIG QUIET)
+  if(TARGET torch)
+    target_compile_definitions(extension_module PRIVATE USE_ATEN_LIB)
+    message(STATUS "extension_module: Using ATen library (USE_ATEN_LIB enabled)")
+  else()
+    message(STATUS "extension_module: PyTorch not found, building without ATen support")
+  endif()
+endif()
+
 # Module extension built as a static library. TODO(gjcomer) Remove this target
 # after cleaning up CMake targets.
 add_library(extension_module_static STATIC ${_extension_module__srcs})
-target_link_libraries(
-  extension_module_static PRIVATE executorch_core extension_data_loader
-                                  extension_flat_tensor
-)
+# Link libraries - only link bundled_program if devtools is enabled
+if(EXECUTORCH_BUILD_DEVTOOLS)
+  target_link_libraries(
+    extension_module_static PRIVATE executorch_core extension_data_loader
+                                    extension_flat_tensor bundled_program
+  )
+  target_compile_definitions(extension_module_static PRIVATE EXECUTORCH_BUILD_DEVTOOLS)
+else()
+  target_link_libraries(
+    extension_module_static PRIVATE executorch_core extension_data_loader
+                                    extension_flat_tensor
+  )
+endif()
 target_include_directories(
   extension_module_static PUBLIC ${_common_include_directories}
 )
 target_compile_options(
   extension_module_static PUBLIC -Wno-deprecated-declarations -fPIC
 )
 
+# Only set USE_ATEN_LIB if PyTorch is actually available
+# This prevents build failures in OSS CMake CI where ATen headers are not available
+if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
+  # Try to find PyTorch headers to check if ATen is available
+  find_package(Torch CONFIG QUIET)
+  if(TARGET torch)
+    target_compile_definitions(extension_module_static PRIVATE USE_ATEN_LIB)
+    message(STATUS "extension_module_static: Using ATen library (USE_ATEN_LIB enabled)")
+  else()
+    message(STATUS "extension_module_static: PyTorch not found, building without ATen support")
+  endif()
+endif()
+
 # Install libraries
 install(
   TARGETS extension_module extension_module_static

extension/module/bundled_module.h

@@ -21,6 +21,9 @@ namespace extension {
 
 using ET_MODULE_NAMESPACE::Module;
 
+// Only expose BundledModule when devtools is available
+#ifdef EXECUTORCH_BUILD_DEVTOOLS
+
 namespace ET_BUNDLED_MODULE_NAMESPACE {
 
 /**
@@ -130,5 +133,8 @@ class BundledModule : public Module {
 };
 
 } // namespace ET_BUNDLED_MODULE_NAMESPACE
+
+#endif // EXECUTORCH_BUILD_DEVTOOLS
+
 } // namespace extension
 } // namespace executorch

extension/pybindings/README.md

@@ -27,8 +27,6 @@ CMAKE_ARGS="-DEXECUTORCH_BUILD_MPS=ON" ./install_executorch.sh
 - `_reset_profile_results()`: Reset profile results.
 ## Classes
 ### ExecuTorchModule
-- `load_bundled_input()`: Load bundled input.
-- `verify_result_with_bundled_expected_output(bundle: str, method_name: str, testset_idx: int, rtol: float = 1e-5, atol: float = 1e-8)`: Verify result with bundled expected output.
 - `plan_execute()`: Plan and execute.
 - `run_method()`: Run method.
 - `forward()`: Forward. This takes a pytree-flattend PyTorch-tensor-based input.
@@ -37,5 +35,6 @@ CMAKE_ARGS="-DEXECUTORCH_BUILD_MPS=ON" ./install_executorch.sh
 - `__call__()`: Call method.
 ### BundledModule
 This class is currently empty and serves as a placeholder for future methods and attributes.
+- `verify_result_with_bundled_expected_output(method_name: str, testset_idx: int, rtol: float = 1e-5, atol: float = 1e-8)`: Verify result with bundled expected output.
 ## Note
 All functions and methods are guarded by a call guard that redirects `cout` and `cerr` to the Python environment.