Revert D55175415: Multisect successfully blamed "D55175415: Fix lift_constant_tensor_pass to make sure constants are not inserted in the state dict" for one test failure (#2767)

Summary:
Pull Request resolved: #2767

This diff reverts D55175415
D55175415: Fix lift_constant_tensor_pass to make sure constants are not inserted in the state dict by tarun292 causes the following test failure:

Tests affected:
- [cogwheel:cogwheel_gpu_lowering_cws_sampled_lowering_replay_test#main](https://www.internalfb.com/intern/test/844425001247266/)

Here's the Multisect link:
https://www.internalfb.com/multisect/4734742
Here are the tasks that are relevant to this breakage:

The backout may land if someone accepts it.

If this diff has been generated in error, you can Commandeer and Abandon it.

Reviewed By: dbort

Differential Revision: D55521379

fbshipit-source-id: 7521019d07b244876358c6906d1be607a010e055

Author: tarun292

exir/lowered_backend_module.py

@@ -203,7 +203,10 @@ def program(self, emit_stacktrace: bool = False) -> Program:
             for node in lowered_exported_program.graph.nodes
             if (
                 node.op == "placeholder"
-                and node.name in lowered_exported_program.graph_signature.user_inputs
+                and node.name
+                not in lowered_exported_program.graph_signature.inputs_to_buffers
+                and node.name
+                not in lowered_exported_program.graph_signature.inputs_to_parameters
             )
         ]
 
@@ -227,8 +230,6 @@ def program(self, emit_stacktrace: bool = False) -> Program:
                 node.name in lowered_exported_program.graph_signature.inputs_to_buffers
                 or node.name
                 in lowered_exported_program.graph_signature.inputs_to_parameters
-                or node.name
-                in lowered_exported_program.graph_signature.inputs_to_lifted_tensor_constants
             ):
                 lowered_exported_program.graph.erase_node(node)
 

exir/passes/constant_prop_pass.py

@@ -5,13 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torch._export.utils import (
-    get_buffer,
-    get_param,
-    is_buffer,
-    is_lifted_tensor_constant,
-    is_param,
-)
+from torch._export.utils import get_buffer, get_param, is_buffer, is_param
 from torch._guards import detect_fake_mode
 from torch.export import ExportedProgram
 from torch.export.exported_program import InputKind, InputSpec, TensorArgument
@@ -27,7 +21,6 @@ def is_const(arg, exported_program, const_data_list) -> bool:
     elif (
         is_param(exported_program, arg)
         or is_buffer(exported_program, arg)
-        or is_lifted_tensor_constant(exported_program, arg)
         or arg.name in const_data_list
     ):
         return True
@@ -41,8 +34,6 @@ def get_data(exported_program, arg):
         return get_param(exported_program, arg)
     elif is_buffer(exported_program, arg):
         return get_buffer(exported_program, arg)
-    elif arg.name in exported_program.constants:
-        return exported_program.constants[arg.name]
     return None
 
 

exir/passes/spec_prop_pass.py

@@ -78,8 +78,6 @@ def update_placeholder_tensor_specs(
                     node.target in exported_program.graph_signature.inputs_to_buffers
                     and not _is_mutable_buffer(node, exported_program.graph_signature)
                 )
-                or node.target
-                in exported_program.graph_signature.inputs_to_lifted_tensor_constants
             ):
                 spec.const = True
 

exir/program/_program.py

@@ -189,7 +189,7 @@ def lift_constant_tensor_pass(ep):
             if not isinstance(constant_tensor, torch.Tensor):
                 continue
 
-            constant_tensor_fqn = f"_lifted_tensor_constant{len(ep.constants)}"
+            constant_tensor_fqn = f"_lifted_tensor_constant{len(buffers)}"
 
             with ep.graph.inserting_before(first_user_input):
                 # Insert the constant node before the first user input
@@ -209,14 +209,14 @@ def lift_constant_tensor_pass(ep):
                 # Add the constant as a buffer to the graph signature
                 lifted_constants.append(
                     InputSpec(
-                        kind=InputKind.CONSTANT_TENSOR,
+                        kind=InputKind.BUFFER,
                         arg=TensorArgument(name=const_placeholder_node.name),
                         target=constant_tensor_fqn,
-                        persistent=None,
+                        persistent=True,
                     )
                 )
                 buffers.append(constant_tensor_fqn)
-                ep.constants[constant_tensor_fqn] = constant_tensor
+                ep.state_dict[constant_tensor_fqn] = constant_tensor
 
     new_input_specs = []
     for s in graph_signature.input_specs:

exir/tests/test_passes.py

@@ -1131,7 +1131,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         # Check (_lifted_tensor_constant + to_copy) node is replaced by prop tensor
         FileCheck().check_not("_lifted_tensor_constant").check(
-            "_prop_tensor_constant0"
+            "_prop_tensor_constant1"
         ).check_not("executorch_exir_dialects_edge__ops_aten__to_copy_default").run(
             new_ep.graph_module.code
         )