generate shared state test program

Differential Revision: D82329519

Pull Request resolved: pytorch#14407

Author: JacobSzwejbka

exir/passes/memory_planning_pass.py

@@ -287,7 +287,7 @@ def run(
         return PassResult(graph_module, True)
 
     def run_multimethod(self):
-        "Resolve any memory planning done across entry points"
+        """Resolve any memory planning done across entry points, called after run is called on all entry points."""
         if self.share_mutable_buffers:
             arena: int = 0
 

extension/module/test/targets.bzl

@@ -19,6 +19,7 @@ def define_common_targets(is_fbcode=False):
             "ET_MODULE_ADD_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])",
             "ET_MODULE_ADD_MUL_PROGRAM_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.pte])",
             "ET_MODULE_ADD_MUL_DATA_PATH": "$(location fbcode//executorch/test/models:exported_program_and_data[ModuleAddMul.ptd])",
+            "ET_MODULE_SHARED_STATE": "$(location fbcode//executorch/test/models:exported_programs[ModuleSharedState.pte])",
         }
 
         for aten_mode in get_aten_mode_options():

test/end2end/exported_module.py

@@ -15,6 +15,7 @@
 import executorch.exir as exir
 import torch
 from executorch.exir import ExecutorchBackendConfig, ExecutorchProgramManager, to_edge
+from executorch.exir.capture._capture import patch_forward
 from executorch.exir.dynamic_shape import DynamicMemoryPlanningMode
 from executorch.exir.passes import (
     DebugPass,
@@ -70,6 +71,7 @@ def export(
         export_joint_graph: bool = False,
         external_constants: bool = False,
         export_state_names: bool = False,
+        share_mutable_buffers: bool = False,
     ) -> "ExportedModule":
         """
         Creates a new ExportedModule for the specified module class.
@@ -134,10 +136,13 @@ def return_wrapper():
             # all exported methods must have the same signature so just pick the first one.
             methods[0],
         )
-        trace_inputs: Sequence = get_trace_inputs()
+        inputs = get_trace_inputs()
         method_name_to_args = {}
         for method in methods:
-            method_name_to_args[method] = trace_inputs
+            if hasattr(eager_module, "get_random_inputs_per_method"):
+                # pyre-ignore
+                inputs = eager_module.get_random_inputs_per_method()[method]  # type: ignore[operator]
+            method_name_to_args[method] = inputs
 
         method_name_to_dynamic_shapes = None
         if hasattr(eager_module, "get_dynamic_shapes"):
@@ -149,23 +154,18 @@ def return_wrapper():
                 method_name_to_dynamic_shapes[method] = trace_dynamic_shapes
 
         memory_planning_pass = MemoryPlanningPass(
-            alloc_mutable_buffers=not export_state_names
+            alloc_mutable_buffers=not export_state_names,
+            share_mutable_buffers=share_mutable_buffers,
         )
         if hasattr(eager_module, "get_memory_planning_pass"):
             memory_planning_pass = eager_module.get_memory_planning_pass()  # type: ignore[operator]
 
-        class WrapperModule(nn.Module):
-            def __init__(self, method):
-                super().__init__()
-                self.forward = method
-
         exported_methods = {}
         # These cleanup passes are required to convert the `add` op to its out
         # variant, along with some other transformations.
         for method_name, method_input in method_name_to_args.items():
             # if not isinstance(eager_module, torch.nn.Module):
             if export_joint_graph:
-                # _export was having issues with WrapperModule.
                 assert method_name == "forward"
                 ep = _export(
                     eager_module,
@@ -179,15 +179,16 @@ def __init__(self, method):
                 )
                 exported_methods[method_name] = _export_forward_backward(ep)
             else:
-                exported_methods[method_name] = export(
-                    eager_module,
-                    method_input,  # type: ignore[arg-type]
-                    dynamic_shapes=(
-                        method_name_to_dynamic_shapes[method_name]
-                        if method_name_to_dynamic_shapes
-                        else None
-                    ),
-                )
+                with patch_forward(eager_module, getattr(eager_module, method_name)):
+                    exported_methods[method_name] = export(
+                        eager_module,
+                        method_input,  # type: ignore[arg-type]
+                        dynamic_shapes=(
+                            method_name_to_dynamic_shapes[method_name]
+                            if method_name_to_dynamic_shapes
+                            else None
+                        ),
+                    )
 
         exec_prog = to_edge(
             exported_methods,
@@ -229,6 +230,6 @@ def __init__(self, method):
             methods=methods,
             executorch_program=exec_prog,
             exported_program=exported_program,
-            trace_inputs=trace_inputs,
+            trace_inputs=inputs,
             get_random_inputs_fn=get_random_inputs_fn,
         )

test/models/export_program.py

@@ -262,6 +262,42 @@ def get_random_inputs(self):
         return (torch.randint(100, [1, 3], dtype=torch.long),)
 
 
+class ModuleSharedState(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("state", torch.ones(1))
+
+    def forward(self, x):
+        return self.state.add_(1) + x
+
+    def get_state(self):
+        return self.state
+
+    def set_state(self, x):
+        self.state.copy_(x)
+
+    # Including this is tech debt since we will immediately override it with the per method one.
+    # ExportedModule is really old infra though from before multiple methods were supported. So
+    # its really obnoxious to change.
+    def get_random_inputs(self):
+        return (torch.ones(1),)
+
+    def get_random_inputs_per_method(self):
+        return {
+            "forward": (torch.ones(1),),
+            "get_state": (),
+            "set_state": (torch.ones(1),),
+        }
+
+    @staticmethod
+    def get_method_names_to_export() -> List[str]:
+        return ["forward", "get_state", "set_state"]
+
+    @staticmethod
+    def share_mutable_buffers():
+        return True
+
+
 #
 # Main logic.
 #
@@ -280,21 +316,28 @@ def export_module_to_program(
         export_kwargs = module_class.get_export_kwargs()
     export_joint = False
     export_state_names = False
+    share_mutable_buffers = False
     if hasattr(module_class, "export_joint"):
-        export_joint = module_class.export_joint()  # pyre-ignore
+        # pyre-ignore[16]: pyre just cant figure it out
+        export_joint = module_class.export_joint()
     if hasattr(module_class, "export_state_names"):
+        # pyre-ignore[16]: pyre just cant figure it out
         export_state_names = module_class.export_state_names()
     if hasattr(module_class, "get_method_names_to_export"):
-        # pyre-ignore[16]: pyre doesn't know about get_export_kwargs.
+        # pyre-ignore[16]: pyre just cant figure it out
         methods = module_class.get_method_names_to_export()
     else:
         methods = ["forward"]
+    if hasattr(module_class, "share_mutable_buffers"):
+        # pyre-ignore[16]: pyre just cant figure it out
+        share_mutable_buffers = module_class.share_mutable_buffers()
     module = ExportedModule.export(
         module_class,
         methods,
         export_joint_graph=export_joint,
         external_constants=external_constants,
         export_state_names=export_state_names,
+        share_mutable_buffers=share_mutable_buffers,
         **export_kwargs,
     )
     return module.executorch_program

test/models/targets.bzl

@@ -71,6 +71,7 @@ def define_common_targets():
         "ModuleDynamicCatUnallocatedIO",
         "ModuleSimpleTrain",
         "ModuleStateful",
+        "ModuleSharedState",
     ]
 
     # Generates Executorch .pte program files for various modules at build time.