[ExecuTorch][#10447] Extend PyBundledModule with extension.BundledModule

CMakeLists.txt

@@ -579,6 +579,14 @@ if(EXECUTORCH_BUILD_PYBIND)
       torch
   )
 
+  if(EXECUTORCH_BUILD_EXTENSION_MODULE)
+    if(CMAKE_TOOLCHAIN_IOS OR CMAKE_TOOLCHAIN_ANDROID OR APPLE)
+      list(APPEND _dep_libs extension_module_static)
+    else()
+      list(APPEND _dep_libs extension_module)
+    endif()
+  endif()
+
   if(EXECUTORCH_BUILD_TESTS)
     list(APPEND _dep_libs test_backend_compiler_lib)
   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/module/bundled_module.cpp

@@ -27,6 +27,8 @@ std::unique_ptr<BufferDataLoader> program_data_loader(
 }
 } // namespace
 
+namespace ET_BUNDLED_MODULE_NAMESPACE {
+
 BundledModule::BundledModule(
     const void* bundled_program_ptr,
     std::unique_ptr<runtime::MemoryAllocator> memory_allocator,
@@ -107,6 +109,6 @@ runtime::Error BundledModule::verify_method_outputs(
   return executorch::BUNDLED_PROGRAM_NAMESPACE::verify_method_outputs(
       *method, bundled_program_ptr_, testset_idx, rtol, atol);
 }
-
+} // namespace ET_BUNDLED_MODULE_NAMESPACE
 } // namespace extension
 } // namespace executorch

extension/module/bundled_module.h

@@ -10,9 +10,19 @@
 
 #include <executorch/extension/module/module.h>
 
+#ifdef USE_ATEN_LIB
+#define ET_BUNDLED_MODULE_NAMESPACE bundled_module::aten
+#else // !USE_ATEN_LIB
+#define ET_BUNDLED_MODULE_NAMESPACE bundled_module
+#endif // USE_ATEN_LIB
+
 namespace executorch {
 namespace extension {
 
+using ET_MODULE_NAMESPACE::Module;
+
+namespace ET_BUNDLED_MODULE_NAMESPACE {
+
 /**
  * A facade class for loading bundled programs and executing methods within
  * them.
@@ -119,5 +129,6 @@ class BundledModule : public Module {
   bool is_loaded_from_file_ = false;
 };
 
+} // namespace ET_BUNDLED_MODULE_NAMESPACE
 } // namespace extension
 } // namespace executorch

extension/module/module.cpp

@@ -36,6 +36,7 @@
 
 namespace executorch {
 namespace extension {
+namespace ET_MODULE_NAMESPACE {
 
 using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::Program;
@@ -308,5 +309,6 @@ runtime::Error Module::set_output(
       output_tensor.mutable_data_ptr(), output_tensor.nbytes(), output_index);
 }
 
+} // namespace ET_MODULE_NAMESPACE
 } // namespace extension
 } // namespace executorch

extension/module/module.h

@@ -16,6 +16,12 @@
 
 #include <executorch/runtime/executor/program.h>
 
+#ifdef USE_ATEN_LIB
+#define ET_MODULE_NAMESPACE module::aten
+#else // !USE_ATEN_LIB
+#define ET_MODULE_NAMESPACE module
+#endif // USE_ATEN_LIB
+
 namespace executorch {
 namespace extension {
 
@@ -24,6 +30,9 @@ using ET_RUNTIME_NAMESPACE::MethodMeta;
 using ET_RUNTIME_NAMESPACE::NamedDataMap;
 using ET_RUNTIME_NAMESPACE::Program;
 
+class ExecuTorchJni;
+
+namespace ET_MODULE_NAMESPACE {
 /**
  * A facade class for loading programs and executing methods within them.
  */
@@ -493,16 +502,24 @@ class Module {
  protected:
   std::unordered_map<std::string, MethodHolder> methods_;
 
-  friend class ExecuTorchJni;
+  friend class executorch::extension::ExecuTorchJni;
 };
 
+} // namespace ET_MODULE_NAMESPACE
 } // namespace extension
 } // namespace executorch
 
 namespace torch {
 namespace executor {
 // TODO(T197294990): Remove these deprecated aliases once all users have moved
 // to the new `::executorch` namespaces.
-using ::executorch::extension::Module;
+using ::executorch::extension::ET_MODULE_NAMESPACE::Module;
 } // namespace executor
 } // namespace torch
+
+namespace executorch {
+namespace extension {
+// backward compatible namespace alias
+using ::executorch::extension::ET_MODULE_NAMESPACE::Module;
+} // namespace extension
+} // namespace executorch

extension/module/test/bundled_module_test.cpp

@@ -9,7 +9,7 @@
 #include <executorch/extension/module/bundled_module.h>
 #include <gtest/gtest.h>
 
-using namespace ::executorch::extension;
+using namespace ::executorch::extension::ET_BUNDLED_MODULE_NAMESPACE;
 using namespace ::executorch::runtime;
 
 class BundledModuleTest : public ::testing::Test {

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.