Fix lift_constant_tensor_pass to make sure constants are not inserted in the state dict

Differential Revision: D55175415

Author: tarun292

exir/passes/constant_prop_pass.py

@@ -5,7 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from torch._export.utils import get_buffer, get_param, is_buffer, is_param
+from torch._export.utils import (
+    get_buffer,
+    get_param,
+    is_buffer,
+    is_lifted_tensor_constant,
+    is_param,
+)
 from torch._guards import detect_fake_mode
 from torch.export import ExportedProgram
 from torch.export.exported_program import InputKind, InputSpec, TensorArgument
@@ -21,6 +27,7 @@ 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
@@ -34,6 +41,8 @@ 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/program/_program.py

@@ -186,7 +186,7 @@ def lift_constant_tensor_pass(ep):
             if not isinstance(constant_tensor, torch.Tensor):
                 continue
 
-            constant_tensor_fqn = f"_lifted_tensor_constant{len(buffers)}"
+            constant_tensor_fqn = f"_lifted_tensor_constant{len(ep.constants)}"
 
             with ep.graph.inserting_before(first_user_input):
                 # Insert the constant node before the first user input
@@ -206,14 +206,14 @@ def lift_constant_tensor_pass(ep):
                 # Add the constant as a buffer to the graph signature
                 lifted_constants.append(
                     InputSpec(
-                        kind=InputKind.BUFFER,
+                        kind=InputKind.CONSTANT_TENSOR,
                         arg=TensorArgument(name=const_placeholder_node.name),
                         target=constant_tensor_fqn,
-                        persistent=True,
+                        persistent=None,
                     )
                 )
                 buffers.append(constant_tensor_fqn)
-                ep.state_dict[constant_tensor_fqn] = constant_tensor
+                ep.constants[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_constant1"
+            "_prop_tensor_constant0"
         ).check_not("executorch_exir_dialects_edge__ops_aten__to_copy_default").run(
             new_ep.graph_module.code
         )