Skip to content

Commit 2d5c81d

Browse files
cccclaipytorchbot
authored andcommitted
add duplicate constant node pass (#2570)
Summary: Pull Request resolved: #2570 This diff is the follow up for #2424 In the case like ``` consant_0 (tag_10) ----> op_b (tag_10) |-------------> op_a (tag_11) ``` `op_b` and `op_a` are in two delegated payload and `constant_0` have two options: In this diff, we're making the default behavior as allowing copying, meaning it will become ``` consant_0 (tag_10)------------------> op_b (tag_10) consant_0_copy (tag_11) -------------> op_a (tag_11) ``` The backend can tag the node with `no_copy` to allert users in cases like constants are too large or etc. In this case, a better approach can be ``` consant_0 (tag_10) ----> op_b (tag_10) |-----(output consant_0) --------> op_a (tag_11) ``` Reviewed By: angelayi Differential Revision: D55113232 fbshipit-source-id: 49559a8cc8225b866a399b4db04ef59fc6860cc1 (cherry picked from commit 7f96f5a)
1 parent 46566b5 commit 2d5c81d

File tree

2 files changed

+2
-4
lines changed

2 files changed

+2
-4
lines changed

exir/backend/test/test_passes.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,6 @@ def forward(self, x):
4242
self.assertEqual(len(copied_nodes), 1)
4343

4444
# Check that the new constant node is in the graph
45-
FileCheck().check("b_const_copy_0").run(
45+
FileCheck().check("arg0_1_copy_0").run(
4646
edge.exported_program().graph_module.code
4747
)

exir/backend/utils.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -246,9 +246,7 @@ def _maybe_duplicate_constant_nodes(
246246
copied_nodes = set()
247247
for candidate_node in candidate_nodes:
248248
# Both tagged exported program and the owning program need to go through the same duplication pass
249-
copied_nodes = copied_nodes.union(
250-
duplicate_constant_node(tagged_exported_program, candidate_node)
251-
)
249+
copied_nodes = duplicate_constant_node(tagged_exported_program, candidate_node)
252250
duplicate_constant_node(owning_program, candidate_node)
253251
candidate_node_with_copies = candidate_nodes.union(copied_nodes)
254252
_assign_new_tag(tagged_exported_program, candidate_node_with_copies)

0 commit comments

Comments
 (0)