Collect Named Data Store at construction

Summary: EdgeProgramManager should never take in a NamedDataStore, it should always be collected from the EdgePrograms that it recieves.

Differential Revision: D71407899

Author: mcr229

exir/program/_program.py

@@ -1098,6 +1098,33 @@ def _gen_edge_manager_for_partitioners(
     return edge_manager
 
 
+def collect_named_data_store_from_exported_program(
+    exported_program: ExportedProgram,
+    named_data_store: NamedDataStore,
+) -> None:
+    """
+    Collects all the named data store outputs found within the exported program
+    and adds them to named_data_store.
+    """
+
+    # collected all the named data into the named data store for deduplication
+    def collect_named_data_store_outputs(
+        graph_module: torch.fx.GraphModule,
+    ) -> None:
+        for node in graph_module.graph.nodes:
+            if node.target == executorch_call_delegate:
+                lbm = getattr(graph_module, node.args[0].name)
+                assert is_lowered_module(lbm)
+                data_store_output = lbm.named_data_store_output
+                if data_store_output is not None:
+                    named_data_store.merge_named_data_store(data_store_output)
+
+        for _, submod, _ in get_control_flow_submodules(graph_module):
+            collect_named_data_store_outputs(submod)
+
+    collect_named_data_store_outputs(exported_program.graph_module)
+
+
 @et_logger("to_edge_transform_and_lower")
 def to_edge_transform_and_lower(
     programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
@@ -1307,7 +1334,6 @@ def __init__(
         constant_methods: Optional[Dict[str, Any]] = None,
         compile_config: Optional[EdgeCompileConfig] = None,
         ops_set_to_not_decompose: Optional[List[torch._ops.OpOverload]] = None,
-        named_data_store: Optional[NamedDataStore] = None,
     ):
         """
         Should not be called directly by users. User should use :func:'to_edge' instead.
@@ -1331,7 +1357,11 @@ def __init__(
         self._edge_programs: Dict[str, ExportedProgram] = edge_programs
         self._config_methods = constant_methods
 
-        self._named_data_store = named_data_store or NamedDataStore()
+        self._named_data_store = NamedDataStore()
+        for _, program in self._edge_programs.items():
+            collect_named_data_store_from_exported_program(
+                program, self._named_data_store
+            )
 
     @property
     def methods(self) -> Set[str]:
@@ -1441,30 +1471,11 @@ def to_backend(
             for name, program in self._edge_programs.items():
                 new_edge_programs[name] = to_backend(program, partitioner)
 
-        # collected all the named data into the named data store for deduplication
-        def collect_named_data_store_outputs(
-            graph_module: torch.fx.GraphModule,
-        ) -> None:
-            for node in graph_module.graph.nodes:
-                if node.target == executorch_call_delegate:
-                    lbm = getattr(graph_module, node.args[0].name)
-                    assert is_lowered_module(lbm)
-                    data_store_output = lbm.named_data_store_output
-                    if data_store_output is not None:
-                        self._named_data_store.merge_named_data_store(data_store_output)
-
-            for _, submod, _ in get_control_flow_submodules(graph_module):
-                collect_named_data_store_outputs(submod)
-
-        for _, program in new_edge_programs.items():
-            collect_named_data_store_outputs(program.graph_module)
-
         config = EdgeCompileConfig(_check_ir_validity=False)
         return EdgeProgramManager(
             new_edge_programs,
             copy.deepcopy(self._config_methods),
             config,
-            named_data_store=self._named_data_store,
         )
 
     @et_logger("to_executorch")