Extend PyBundledModule with extension.BundledModule (pytorch#12839)

Summary: Pull Request resolved: pytorch#12839

Differential Revision: D78938344

Author: Gasoonjia

CMakeLists.txt

@@ -561,6 +561,12 @@ install(FILES tools/cmake/executorch-config.cmake
         DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/ExecuTorch
 )
 
+# Build devtools first if needed - some backends depend on protobuf from
+# devtools
+if(EXECUTORCH_BUILD_DEVTOOLS)
+  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools)
+endif()
+
 if(EXECUTORCH_BUILD_ARM_BAREMETAL)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/arm)
   list(APPEND _executorch_backends executorch_delegate_ethos_u)
@@ -609,10 +615,6 @@ if(EXECUTORCH_BUILD_CORTEX_M)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m)
 endif()
 
-if(EXECUTORCH_BUILD_DEVTOOLS)
-  add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools)
-endif()
-
 if(EXECUTORCH_BUILD_EXTENSION_APPLE)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/apple)
   list(APPEND _executorch_extensions apple_extension)
@@ -718,6 +720,30 @@ if(EXECUTORCH_BUILD_PYBIND)
     add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/devtools)
   endif()
 
+  # Create bundled_module target only for pybindings when bundled_program exists
+  # This target has hard dependencies on devtools generated headers
+  if(TARGET bundled_program)
+    add_library(
+      bundled_module STATIC
+      ${CMAKE_CURRENT_SOURCE_DIR}/extension/module/bundled_module.cpp
+    )
+
+    # Ensure bundled_module waits for bundled_program's generated headers
+    add_dependencies(bundled_module bundled_program)
+
+    target_link_libraries(bundled_module PRIVATE extension_data_loader)
+    target_link_libraries(
+      bundled_module PUBLIC extension_module_static bundled_program
+    )
+
+    target_include_directories(
+      bundled_module PUBLIC ${_common_include_directories}
+    )
+    target_compile_options(
+      bundled_module PUBLIC -Wno-deprecated-declarations -fPIC
+    )
+  endif()
+
   # find pytorch lib, to allow pybind to take at::Tensor as input/output
   find_package_torch()
   find_library(
@@ -735,6 +761,16 @@ 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)
+    # Add bundled_module if available
+    if(TARGET bundled_module)
+      list(APPEND _dep_libs bundled_module)
+    endif()
+  endif()
+
   if(EXECUTORCH_BUILD_TESTS)
     list(APPEND _dep_libs test_backend_compiler_lib)
   endif()

backends/apple/coreml/CMakeLists.txt

@@ -205,19 +205,34 @@ if(APPLE)
   )
 
   if(EXECUTORCH_BUILD_DEVTOOLS)
-    target_sources(coremldelegate PRIVATE ${SDK_SOURCES} ${PROTOBUF_SOURCES})
+    # Always add SDK sources
+    target_sources(coremldelegate PRIVATE ${SDK_SOURCES})
     target_include_directories(
       coremldelegate
       PRIVATE
         ${CMAKE_CURRENT_SOURCE_DIR}/runtime/sdk
-        ${CMAKE_CURRENT_SOURCE_DIR}/third-party/coremltools/deps/protobuf/src
-    )
-    add_subdirectory(
-      ${CMAKE_CURRENT_SOURCE_DIR}/third-party/coremltools/deps/protobuf/cmake
     )
 
-    executorch_target_link_options_shared_lib(libprotobuf-lite)
-    target_link_libraries(coremldelegate PRIVATE libprotobuf-lite)
+    # Check if CoreML's protobuf sources exist, otherwise skip protobuf features
+    set(COREML_PROTOBUF_CMAKE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/third-party/coremltools/deps/protobuf/cmake")
+    if(EXISTS "${COREML_PROTOBUF_CMAKE_DIR}")
+      message(STATUS "CoreML: Using vendored protobuf from coremltools")
+      target_include_directories(
+        coremldelegate
+        PRIVATE
+          ${CMAKE_CURRENT_SOURCE_DIR}/third-party/coremltools/deps/protobuf/src
+      )
+      add_subdirectory("${COREML_PROTOBUF_CMAKE_DIR}")
+      executorch_target_link_options_shared_lib(libprotobuf-lite)
+      target_link_libraries(coremldelegate PRIVATE libprotobuf-lite)
+
+      # Only add protobuf sources after protobuf has been configured
+      target_sources(coremldelegate PRIVATE ${PROTOBUF_SOURCES})
+      set(COREML_HAS_PROTOBUF TRUE)
+    else()
+      message(WARNING "CoreML: Protobuf sources not found at ${COREML_PROTOBUF_CMAKE_DIR}. Some CoreML debugging features will be disabled.")
+      set(COREML_HAS_PROTOBUF FALSE)
+    endif()
   endif()
 
   target_link_libraries(
@@ -243,7 +258,10 @@ if(APPLE)
   if(EXECUTORCH_BUILD_DEVTOOLS)
     target_compile_options(executorch_core PUBLIC -DET_EVENT_TRACER_ENABLED)
     target_compile_options(coremldelegate PRIVATE "-frtti")
-    target_compile_options(libprotobuf-lite PRIVATE "-frtti")
+    # Only set compile options on protobuf if it was successfully configured
+    if(COREML_HAS_PROTOBUF AND TARGET libprotobuf-lite)
+      target_compile_options(libprotobuf-lite PRIVATE "-frtti")
+    endif()
   else()
     target_compile_options(coremldelegate PRIVATE "-fno-rtti")
   endif()

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/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.