end2end test for bundled program (#395)

Summary:
Pull Request resolved: #395

Update pybinding functions to support bundled program (bp) APIs with method name. Split the end2end test for bp into a single test file.

The reasonsfor spliting are:
1. Current overall end2end test does not support multiple methods, but we need to test bp on multiple methods.
2. Bp does not need to be tested on several different kinds of program.
3. Decompose the test for core ET and Bp.

Differential Revision: D49336485

fbshipit-source-id: 54eb03103a31cdd2ae4999c7471eb6e6f0378eb0

Author: Gasoonjia

bundled_program/tests/TARGETS

@@ -1,3 +1,5 @@
+# @noautodeps
+
 load("@fbcode_macros//build_defs:python_library.bzl", "python_library")
 load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest")
 
@@ -58,3 +60,33 @@ python_unittest(
         "//executorch/extension/pytree:pylib",
     ],
 )
+
+python_unittest(
+    name = "end2end",
+    srcs = [
+        "test_end2end.py",
+    ],
+    deps = [
+        ":lib",
+        "//caffe2:torch",
+        "//executorch/bundled_program:config",
+        "//executorch/bundled_program:core",
+        "//executorch/bundled_program/serialize:lib",
+        "//executorch/exir:dynamic_shape",
+        "//executorch/exir:lib",
+        "//executorch/exir:memory",
+        "//executorch/exir:pass_manager",
+        "//executorch/exir:print_program",
+        "//executorch/exir:tensor",
+        "//executorch/exir/_serialize:lib",
+        "//executorch/exir/emit:lib",
+        "//executorch/exir/passes:lib",
+        "//executorch/exir/tests:control_flow_models",
+        "//executorch/exir/tests:dynamic_shape_models",
+        "//executorch/exir/tests:models",
+        "//executorch/exir/tests:transformer",
+        "//executorch/extension/pybindings:portable_lib",
+        "//executorch/extension/pytree:pybindings",
+        "//executorch/kernels/portable:custom_ops_generated_lib",
+    ],
+)

bundled_program/tests/test_end2end.py

@@ -0,0 +1,93 @@
+# 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.
+
+# 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.bundled_program.config import BundledConfig
+
+from executorch.bundled_program.core import create_bundled_program
+from executorch.bundled_program.serialize import (
+    serialize_from_bundled_program_to_flatbuffer,
+)
+
+from executorch.bundled_program.tests.common import get_common_program, SampleModel
+
+kernel_mode = None  # either aten mode or lean mode
+try:
+    # pyre-ignore[21]
+    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"
+except ImportError as e:
+    print(e)
+    pass
+
+try:
+    # pyre-ignore[21]
+    from executorch.extension.pybindings.aten_lib import (
+        _load_bundled_program_from_buffer,
+        _load_for_executorch_from_buffer,
+        _load_for_executorch_from_bundled_program,
+    )
+
+    assert kernel_mode is None
+    kernel_mode = "aten"
+except ImportError as e:
+    print(e)
+    pass
+
+assert kernel_mode is not None
+
+
+class BundledProgramE2ETest(unittest.TestCase):
+    def test_sample_model_e2e(self):
+        program, bundled_config = get_common_program()
+        eager_model = SampleModel()
+
+        bundled_program = create_bundled_program(program, bundled_config)
+
+        bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
+            bundled_program
+        )
+
+        executorch_bundled_program = _load_bundled_program_from_buffer(
+            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,
+                method_name,
+                0,
+            )

extension/pybindings/pybindings.cpp

@@ -175,6 +175,14 @@ class Module final {
     return result;
   }
 
+  Method& get_method(const std::string& method_name) {
+    if (methods_.count(method_name) == 0) {
+      THROW_IF_ERROR(
+          Error(), "no such method in program: %s", method_name.c_str());
+    }
+    return *methods_[method_name].get();
+  }
+
  private:
   /// A wrapper/util class for executorch memory allocations/manager.
   class Memory {
@@ -451,6 +459,48 @@ struct PyModule final {
     return run_method("forward", inputs);
   }
 
+  void load_bundled_input(
+      PyBundledModule& m,
+      const string method_name,
+      size_t testset_idx) {
+    const void* bundled_program_ptr = m.get_bundled_program_ptr();
+    Error status = util::LoadBundledInput(
+        module_->get_method(method_name),
+        bundled_program_ptr,
+        &m.get_bundled_input_allocator(),
+        method_name.c_str(),
+        testset_idx);
+    ET_CHECK_MSG(
+        status == Error::Ok,
+        "LoadBundledInput failed with status %" PRIu32,
+        status);
+  }
+
+  void verify_result_with_bundled_expected_output(
+      PyBundledModule& m,
+      const string method_name,
+      size_t testset_idx) {
+    const void* bundled_program_ptr = m.get_bundled_program_ptr();
+    Error status = util::VerifyResultWithBundledExpectedOutput(
+        module_->get_method(method_name),
+        bundled_program_ptr,
+        &m.get_bundled_input_allocator(),
+        method_name.c_str(),
+        testset_idx);
+    ET_CHECK_MSG(
+        status == Error::Ok,
+        "Result verification failed with status %" PRIu32,
+        status);
+  }
+
+  void plan_execute(const string method_name) {
+    auto status = module_->get_method(method_name).execute();
+    THROW_IF_ERROR(
+        status,
+        "executing execution plan for method 'forward' failed with error: 0x%" PRIx32,
+        status);
+  }
+
  private:
   std::unique_ptr<Module> module_;
 };
@@ -487,6 +537,11 @@ PYBIND11_MODULE(EXECUTORCH_PYTHON_MODULE_NAME, m) {
   m.def("_reset_profile_results", []() { EXECUTORCH_RESET_PROFILE_RESULTS(); });
 
   py::class_<PyModule>(m, "ExecutorchModule")
+      .def("load_bundled_input", &PyModule::load_bundled_input)
+      .def(
+          "verify_result_with_bundled_expected_output",
+          &PyModule::verify_result_with_bundled_expected_output)
+      .def("plan_execute", &PyModule::plan_execute)
       .def("run_method", &PyModule::run_method)
       .def("forward", &PyModule::forward);
 

test/end2end/test_end2end.py

@@ -20,12 +20,7 @@
 # @manual=//executorch/extension/pytree:pybindings
 import executorch.extension.pytree as pytree
 import torch
-from executorch.bundled_program.config import BundledConfig
 
-from executorch.bundled_program.core import create_bundled_program
-from executorch.bundled_program.serialize import (
-    serialize_from_bundled_program_to_flatbuffer,
-)
 from executorch.exir import (
     CaptureConfig,
     EdgeCompileConfig,
@@ -486,7 +481,6 @@ def maketest(
     allow_non_contiguous_tensor: bool = False,
     method: str = "forward",
     dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND,
-    bundled_io=False,
     capture_config=None,
     verify_graph: Optional[Callable] = None,
 ) -> Callable[[unittest.TestCase], None]:
@@ -510,9 +504,6 @@ def maketest(
             program only contains contiguous tensors.
         method: The name of the module_cls method to trace.
         dynamic_memory_planning_mode: The dynamic memory planning mode to use.
-        bundled_io: If true, will wrap the given program into bundled program
-            format, run and compared the program with bundled input and expected
-            output.
 
     Returns:
         A TestCase method that tests the provided module class and method.
@@ -587,54 +578,6 @@ def wrapper(self: unittest.TestCase) -> None:
                     print(f"actual result: {actual}")
                 self.assertTrue(is_close)
 
-        if bundled_io:
-            print("Being verified by Bundled Program")
-            expected_outputs_list = [
-                [[module.eager_module(*x)] for x in inputs_list],
-            ]
-            bundled_config = BundledConfig(
-                [
-                    method,
-                ],
-                [
-                    inputs_list,
-                ],
-                expected_outputs_list,
-            )
-            bundled_program = create_bundled_program(
-                module.executorch_program.program, bundled_config
-            )
-            bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
-                bundled_program
-            )
-
-            # pyre-fixme[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
-            executorch_bundled_program = _load_bundled_program_from_buffer(
-                bundled_program_buffer
-            )
-
-            # pyre-fixme[16]: Module `executorch.extension.pybindings` has no attribute `portable`.
-            executorch_module = _load_for_executorch_from_bundled_program(
-                executorch_bundled_program
-            )
-
-            default_execution_plan_id = 0
-
-            # TODO(T144329357): check bundled attachment correctness
-            # No load_bundled_input() method
-            # for testset_idx in range(niter):
-            #     executorch_module.load_bundled_input(
-            #         executorch_bundled_program,
-            #         default_execution_plan_id,
-            #         testset_idx,
-            #     )
-            #     executorch_module.plan_execute()
-            #     executorch_module.verify_result_with_bundled_expected_output(
-            #         executorch_bundled_program,
-            #         default_execution_plan_id,
-            #         testset_idx,
-            #     )
-
     return wrapper
 
 
@@ -823,55 +766,3 @@ def test_batch_norm(self):
             # run this on aten mode.
             run_executor=is_aten_mode,
         )(self)
-
-
-class BundledProgramE2ETest(unittest.TestCase):
-    # Using all models supporting executor running in this test.
-
-    def test_mem_planning_toy_model_bundled_program(self):
-        maketest(ToyModelForMemPlanning, bundled_io=True)(self)
-
-    def test_executorch_forward_bundled_program(self):
-        maketest(ModuleAdd, bundled_io=True)(self)
-
-    @skipUnless(RUN_SKIPPED, "TODO(larryliu0820) Fix this in both fbcode and oss")
-    def test_containers_bundled_program(self):
-        maketest(ModuleContainers, do_tree_flatten=True, bundled_io=True)(self)
-
-    # Failed to produce a graph during tracing w/ dynamo because there are no torch ops
-    # test_return_input_bundled_program = maketest(
-    #     ModuleReturnInput, do_tree_flatten=True, bundled_io=True
-    # )
-
-    # can not run this on the executor because missing the following ops:
-    #   aten::select_copy.int_out, aten::eq.Scalar_out
-    # TODO(gasoonjia) re-enable these tests.
-    # test_ifelse_bundled_program = maketest(ModuleIfElse, bundled_io=True)
-    # test_ifelse_with_bool_bundled_program = maketest(ModuleIfElseWithBool, bundled_io=True)
-
-    # fail to trace with functionalization enabled
-    # Fail with error: Missing out variants: {'aten::select', 'aten::_shape_as_tensor', 'aten::tensor_split'}
-    # TODO(gasoonjia) re-enable these tests.
-    # test_while_bundled_program = maketest(ModuleWhile, bundled_io=True)
-
-    # test_while_if_bundled_program = maketest(ModuleWhileIf, bundled_io=True)
-    # test_if_while_bundled_program = maketest(ModuleIfWhile, bundled_io=True)
-
-    def test_input_dynamic_shape_bundled_program(self):
-        maketest(
-            ModuleInputDynamicShape,
-            run_graph_module=False,
-            bundled_io=True,
-            capture_config=exir.CaptureConfig(
-                enable_dynamic_shape=True,
-            ),
-        )(self)
-
-    @skip("revisit when unbacked symint is ready")
-    def test_intermediate_dynamic_shape_bundled_program(self):
-        maketest(
-            ModuleIntermediateDynamicShape,
-            run_graph_module=False,
-            allow_non_contiguous_tensor=True,
-            bundled_io=True,
-        )(self)