From c518d3fe0dae6a87c57aa5ce40ec49168a90acf0 Mon Sep 17 00:00:00 2001 From: Max Ren Date: Tue, 18 Mar 2025 15:06:43 -0700 Subject: [PATCH 1/2] Collect Named Data Store at construction (#9370) Summary: EdgeProgramManager should never take in a NamedDataStore, it should always be collected from the EdgePrograms that it recieves. Reviewed By: tarun292 Differential Revision: D71407899 --- .../test/test_backend_with_named_data_map.py | 4 +- exir/program/_program.py | 53 +++++++++++-------- 2 files changed, 34 insertions(+), 23 deletions(-) diff --git a/exir/backend/test/test_backend_with_named_data_map.py b/exir/backend/test/test_backend_with_named_data_map.py index cc7aad641f0..c9e458a1878 100644 --- a/exir/backend/test/test_backend_with_named_data_map.py +++ b/exir/backend/test/test_backend_with_named_data_map.py @@ -45,7 +45,7 @@ def forward(self, x): return y - y ep = to_edge(torch.export.export(M(), (torch.randn(1, 2),))) - ep.to_backend(BackendWithNDMPartitioner()) + ep = ep.to_backend(BackendWithNDMPartitioner()) ndm_output = ep._named_data_store.get_named_data_store_output() buffer_entries = ndm_output.buffers @@ -71,7 +71,7 @@ def forward(self, x, y): return z - z ep = to_edge(torch.export.export(M(), (torch.randn(1, 2), torch.randn(1, 2)))) - ep.to_backend(BackendWithNDMPartitioner()) + ep = ep.to_backend(BackendWithNDMPartitioner()) ndm_output = ep._named_data_store.get_named_data_store_output() buffer_entries = ndm_output.buffers diff --git a/exir/program/_program.py b/exir/program/_program.py index d8d07aefe85..1a30c739414 100644 --- a/exir/program/_program.py +++ b/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") From c7f62d5a4065723906f8176ac4e4ea6af85824a6 Mon Sep 17 00:00:00 2001 From: Max Ren Date: Wed, 19 Mar 2025 17:13:14 -0700 Subject: [PATCH 2/2] Fix up issues from named data collection --- exir/lowered_backend_module.py | 1 + exir/program/_program.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index ed155555ef5..6bcc1b2f3d8 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -107,6 +107,7 @@ def __deepcopy__(self, memo: Optional[Dict[int, Any]]) -> "LoweredBackendModule" backend_id=self._backend_id, processed_bytes=self._processed_bytes, compile_specs=copy.deepcopy(self._compile_specs, memo), + named_data_store_output=self._named_data_store_output, ) # pyre-fixme[16]: `LoweredBackendModule` has no attribute `meta`. res.meta = copy.copy(getattr(self, "meta", {})) diff --git a/exir/program/_program.py b/exir/program/_program.py index 1a30c739414..d3cc0578068 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1113,7 +1113,7 @@ def collect_named_data_store_outputs( ) -> None: for node in graph_module.graph.nodes: if node.target == executorch_call_delegate: - lbm = getattr(graph_module, node.args[0].name) + lbm = getattr(graph_module, node.args[0].target) assert is_lowered_module(lbm) data_store_output = lbm.named_data_store_output if data_store_output is not None: