Fix output spec + insert clone for constant_prop_pass.

Differential Revision: D75473310

Pull Request resolved: #11209

Author: hsharma35

exir/passes/constant_prop_pass.py

@@ -295,6 +295,37 @@ def create_constant_nodes_and_return_specs(
     return name_to_spec_dict
 
 
+def _update_output_node_and_specs(exported_program: ExportedProgram) -> None:
+    """
+    Update the output node and output specs in the exported program.
+    In case a constant node is used as output, we replace it with a clone of the constant node.
+    """
+    # Dict [node.name -> InputSpec]
+    updated_constant_placeholders = get_constant_placeholder_dict(exported_program)
+    output = exported_program.graph.find_nodes(op="output")[0]
+    output_nodes = cast(list[torch.fx.Node], list(output.args[0]))
+    output_specs = exported_program.graph_signature.output_specs
+    assert len(output_nodes) == len(output_specs)
+
+    for i in range(len(output_specs)):
+        out_node = output_nodes[i]
+        if out_node not in updated_constant_placeholders:
+            continue
+
+        with exported_program.graph.inserting_after(out_node):
+            new_node = exported_program.graph.call_function(
+                exir_ops.edge.aten.clone.default, (out_node,)
+            )
+        assert "val" in out_node.meta
+        new_node.meta["val"] = out_node.meta["val"]
+        output_nodes[i] = new_node
+
+        # Update the constant-propagated output node.
+        output_specs[i].arg = TensorArgument(name=output_nodes[i].name)
+
+    output.args = (output_nodes,)
+
+
 def constant_prop_pass(
     exported_program: ExportedProgram,
     custom_skip_targets: Optional[set[EdgeOpOverload]] = None,
@@ -341,12 +372,12 @@ def constant_prop_pass(
 
     # Generate new input spec.
     new_input_specs = []
-    for node in exported_program.graph.nodes:
-        if node.op != "placeholder":
-            continue
+    for node in exported_program.graph.find_nodes(op="placeholder"):
         new_input_specs.append(name_to_spec_dict[node.name])
     exported_program.graph_signature.input_specs = new_input_specs
 
+    _update_output_node_and_specs(exported_program)
+
     # Cleanup the graph.
     exported_program.graph.eliminate_dead_code()
     exported_program.graph_module.recompile()

exir/tests/test_passes.py

@@ -1026,6 +1026,34 @@ def forward(self, x):
             "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"
         ).run(gm.code)
 
+    def test_constant_prop_for_output(self) -> None:
+        class Add(torch.nn.Module):
+            def forward(self) -> torch.Tensor:
+                return torch.add(torch.tensor(3), torch.tensor(5))
+
+        add = Add()
+
+        edge = to_edge(
+            export(add, (), strict=True),
+            compile_config=EdgeCompileConfig(_skip_dim_order=False),
+        )
+        # Check there is a lifted tensor followed by a to_copy node
+        FileCheck().check("c_lifted_tensor_0").check("c_lifted_tensor_1").run(
+            edge.exported_program().graph_module.code
+        )
+
+        edge._edge_programs["forward"] = constant_prop_pass(
+            edge.exported_program("forward")
+        )
+
+        # Check (c_lifted_tensor_*) nodes are all replaced by _prop_tensor_constant.
+        FileCheck().check_not("c_lifted_tensor_").check("_prop_tensor_constant").run(
+            edge.exported_program().graph_module.code
+        )
+        # Validate that the program successfully passes validation to executorch:
+        edge.exported_program()._validate()
+        edge.to_executorch()
+
     def test_constant_prop_pass_for_add(self) -> None:
         class Add(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor: