[Executorch][to_backend] Introduce preprocess_multimethod

CMakeLists.txt

@@ -810,6 +810,10 @@ if(EXECUTORCH_BUILD_PYBIND)
       torch
   )
 
+  if(EXECUTORCH_BUILD_TESTS)
+    list(APPEND _dep_libs test_backend_compiler_lib)
+  endif()
+
   if(EXECUTORCH_BUILD_KERNELS_OPTIMIZED)
     list(APPEND _dep_libs optimized_native_cpu_ops_lib)
   else()

devtools/inspector/tests/inspector_utils_test.py

@@ -205,7 +205,7 @@ def test_compare_results(self):
         self.assertAlmostEqual(calculate_cosine_similarity([a], [b])[0], 1.0)
 
     def test_compare_results_uint8(self):
-        a = torch.randint(0, 255, (4, 4), dtype=torch.uint8)
+        a = torch.randint(1, 255, (4, 4), dtype=torch.uint8)
 
         # Create tensor b which has very close value to tensor a
         b = a.clone()

exir/backend/backend_details.py

@@ -50,15 +50,6 @@ class BackendDetails(ABC):
     the decorators, this interface will be static, abstract and all inheritances are
     enforced to implement this method.
 
-    Args:
-        edge_program: The original exported program. It will not be modified in place.
-        compile_specs: List of values needed for compilation
-
-    Returns:
-        PreprocessResult: It wraps the following information:
-            processed_bytes -> bytes: A compiled blob - a binary that can run the desired program in the backend.
-            debug_handle_map (Optional[Dict[int, Tuple[int]]]): For profiling purposes, a map from the node_id in the final graph (either EXIR or the user's self-defined IR)
-            to debug handle id attached in the original exported program.
     """
 
     @staticmethod
@@ -70,6 +61,69 @@ def preprocess(
         edge_program: ExportedProgram,
         compile_specs: List[CompileSpec],
     ) -> PreprocessResult:
+        """
+        Preprocesses an edge program and returns the preprocess result fo the given backend
+
+        Args:
+            edge_program: The original exported program. It will not be modified in place.
+            compile_specs: List of values needed for compilation
+
+        Returns:
+            PreprocessResult: It wraps the following information:
+                processed_bytes -> bytes: A compiled blob - a binary that can run the desired
+                program in the backend.
+                debug_handle_map (Optional[Dict[int, Tuple[int]]]): For profiling purposes, a
+                map from the node_id  in the final graph (either EXIR or the user's self-defined
+                IR) to debug handle id attached in the original exported program.
+        """
         # Users should return a compiled blob - a binary that can run the desired
         # program in the backend.
         pass
+
+    @classmethod
+    def preprocess_multimethod(
+        cls,
+        edge_programs: Dict[str, List[ExportedProgram]],
+        compile_specs: Dict[str, List[List[CompileSpec]]],
+    ) -> Dict[str, list[PreprocessResult]]:
+        """
+        Runs preprocess on all partitioned Edge Programs across multiple methods. This allows
+        backends to share information across partitioned graphs. Backend can serialize shared
+        data by putting the shared data into the data_store_output of the preprocess results.
+        This will record the shared data used by that specific partition.
+
+        Default implementation is running the existing preprocess implementation on all
+
+        Args:
+            edge_programs: Dictionary mapping the method name to a list of all the partitioned
+                           edge_programs from that method to be lowered.
+            compile_specs: Dictionary mapping the method name to a list of compile_specs. The
+                           list of compile specs maps directly to the list of edge_programs for the
+                           same given method name i.e. edge_program[method_name][i] --> compile_specs[method_name][i]
+
+        Returns:
+            Dictionary mapping the method name to a list of PreprocessResults. The list of
+            PreprocessResults maps directly to the list of edge_programs for the same given
+            method name. i.e. edge_program[method_name][i] --> result[method_name][i]
+
+
+        """
+        preprocess_results = {}
+        for method_name, programs in edge_programs.items():
+            assert (
+                method_name in compile_specs
+            ), f"Error: missing compile specs for {method_name}"
+            compile_specs_for_method = compile_specs[method_name]
+            assert len(compile_specs_for_method) == len(
+                programs
+            ), f"Error: method {method_name} has {len(programs)} partitions but only {len(compile_specs_for_method)}"
+            results_for_method = []
+            for program, compile_spec_for_program in zip(
+                programs, compile_specs_for_method
+            ):
+                preprocess_result = cls.preprocess(program, compile_spec_for_program)
+                results_for_method.append(preprocess_result)
+
+            preprocess_results[method_name] = results_for_method
+
+        return preprocess_results

exir/backend/test/test_lowered_backend_module.py

@@ -22,7 +22,6 @@
 from executorch.extension.pybindings.portable_lib import (  # @manual
     _load_for_executorch_from_buffer,
 )
-from hypothesis import given, settings, strategies as st
 from torch.export import export
 
 
@@ -65,7 +64,6 @@ def forward(self, *args):
             .executorch_program
         )
 
-    @settings(deadline=500000)
     def test_emit_lowered_backend_module_end_to_end(self):
         class SinModule(torch.nn.Module):
             def __init__(self):
@@ -109,11 +107,7 @@ def forward(self, x):
             torch.allclose(model_outputs[0], expected_res, atol=1e-03, rtol=1e-03)
         )
 
-    @given(
-        unlift=st.booleans(),  # verify both lifted and unlifted graph
-    )
-    @settings(deadline=500000)
-    def test_emit_lowered_backend_module(self, unlift):
+    def test_emit_lowered_backend_module(self):
         module_list = [
             models.Emformer(),
             models.Repeat(),
@@ -166,11 +160,7 @@ def test_emit_lowered_backend_module(self, unlift):
             _ = lowered_model.buffer()
             self.validate_lowered_module_program(program)
 
-    @given(
-        unlift=st.booleans(),  # verify both lifted and unlifted graph
-    )
-    @settings(deadline=500000)
-    def test_emit_nested_lowered_backend_module(self, unlift):
+    def test_emit_nested_lowered_backend_module(self):
         module_list = [
             models.Emformer(),
             models.Repeat(),

pytest.ini

@@ -63,8 +63,6 @@ addopts =
     --ignore=exir/backend/test/demos
     --ignore=exir/backend/test/test_backends.py
     --ignore=exir/backend/test/test_backends_lifted.py
-    --ignore=exir/backend/test/test_compatibility.py
-    --ignore=exir/backend/test/test_lowered_backend_module.py
     --ignore=exir/backend/test/test_partitioner.py
     --ignore=exir/tests/test_common.py
     --ignore=exir/tests/test_memory_format_ops_pass_aten.py

runtime/executor/test/CMakeLists.txt

@@ -152,3 +152,23 @@ target_include_directories(
   PRIVATE "${CMAKE_INSTALL_PREFIX}/schema/include"
           "${EXECUTORCH_ROOT}/third-party/flatbuffers/include"
 )
+
+list(TRANSFORM _test_backend_compiler_lib__srcs PREPEND "${EXECUTORCH_ROOT}/")
+add_library(
+  test_backend_compiler_lib
+  STATIC
+  ${_test_backend_compiler_lib__srcs}
+)
+
+target_link_libraries(
+  test_backend_compiler_lib
+  PUBLIC
+  executorch_core
+)
+
+target_link_options_shared_lib(test_backend_compiler_lib)
+
+install(
+  TARGETS test_backend_compiler_lib
+  DESTINATION lib
+)

setup.py

@@ -718,6 +718,7 @@ def run(self):
             # enabled. TODO(dbort): Remove this override once this option is
             # managed by cmake itself.
             "-DEXECUTORCH_SEPARATE_FLATCC_HOST_PROJECT=OFF",
+            "-DEXECUTORCH_BUILD_TESTS=ON",
         ]
 
         build_args = [f"-j{self.parallel}"]

tools/cmake/cmake_deps.toml

@@ -150,6 +150,20 @@ deps = [
   "optimized_cpublas",
   "portable_kernels",
 ]
+
+[targets.test_backend_compiler_lib]
+buck_targets = [
+  "//runtime/executor/test:test_backend_compiler_lib",
+]
+filters = [
+  ".cpp$",
+]
+excludes = [
+]
+deps = [
+  "executorch",
+  "executorch_core",
+]
 # ---------------------------------- core end ----------------------------------
 # ---------------------------------- extension start ----------------------------------
 [targets.extension_data_loader]