Ensure that lifted tensor constants don't show up as inputs in emitted program (#1897)

Summary:

Currently lifted tensor constants are showing up as inputs to the emitted program. This shouldn't be the case as they're embedded inside the program as constants and the user will not be passing these in as inputs.

Reviewed By: angelayi

Differential Revision: D53584903

Author: tarun292

exir/emit/_emitter.py

@@ -1308,18 +1308,33 @@ def placeholder(
         if isinstance(target, str) and (
             target in self.exported_program.graph_signature.inputs_to_parameters
             or target in self.exported_program.graph_signature.inputs_to_buffers
+            or target
+            in self.exported_program.graph_signature.inputs_to_lifted_tensor_constants
         ):
-
-            fqn = (
-                self.exported_program.graph_signature.inputs_to_parameters[target]
-                if target in self.exported_program.graph_signature.inputs_to_parameters
-                else self.exported_program.graph_signature.inputs_to_buffers[target]
-            )
+            if (
+                target
+                in self.exported_program.graph_signature.inputs_to_lifted_tensor_constants
+            ):
+                fqn = self.exported_program.graph_signature.inputs_to_lifted_tensor_constants[
+                    target
+                ]
+            elif target in self.exported_program.graph_signature.inputs_to_buffers:
+                fqn = self.exported_program.graph_signature.inputs_to_buffers[target]
+            else:
+                fqn = self.exported_program.graph_signature.inputs_to_parameters[target]
             if fqn in self.exported_program.state_dict:
                 spec = TensorSpec.from_tensor(
                     self.exported_program.state_dict[fqn], const=True
                 )
                 const_tensor = True
+            elif fqn in self.exported_program.constants:
+                print(
+                    f"Found constant {self.exported_program.constants[fqn]} in constants dict"
+                )
+                spec = TensorSpec.from_tensor(
+                    self.exported_program.constants[fqn], const=True
+                )
+                const_tensor = True
             else:
                 buffers = self.exported_program.named_buffers()
                 buf = next((x[1] for x in buffers if x[0] == fqn), None)

exir/emit/test/test_emit.py

@@ -1335,3 +1335,28 @@ def forward(self, x):
         # confirm that the buffer was emitted
         self.assertEqual(len(program.constant_buffer), 2)
         self.assertEqual(len(program.constant_buffer[1].storage), 8)
+
+    def test_emit_lifted_tensor_constant(self):
+        class LiftedConstants(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                x = x * torch.tensor([[4, 3], [1, 2], [5, 6]], dtype=torch.float)
+                return x
+
+        model = LiftedConstants()
+
+        program = to_edge(
+            export(
+                model,
+                (torch.ones(3, 2),),
+            )
+        ).to_executorch()
+
+        program = program._emitter_output.program
+        exec_plan = program.execution_plan[0]
+        # There should only be 1 input to this model.
+        self.assertEqual(len(exec_plan.inputs), 1)
+        self.assertEqual(len(program.constant_buffer), 2)
+        self.assertEqual(len(program.constant_buffer[1].storage), 24)

exir/lowered_backend_module.py

@@ -465,6 +465,21 @@ def _get_new_signature(
                     ]
                 else:
                     new_constants[buffer_name] = original_program.constants[buffer_name]
+            elif node.name in old_signature.inputs_to_lifted_tensor_constants:
+                constant_name = old_signature.inputs_to_lifted_tensor_constants[
+                    node.name
+                ]
+                # add constant to graph signature
+                input_specs.append(
+                    InputSpec(
+                        kind=InputKind.CONSTANT_TENSOR,
+                        arg=TensorArgument(name=node.name),
+                        target=constant_name,
+                    )
+                )
+
+                # add constant to new_constants
+                new_constants[constant_name] = original_program.constants[constant_name]
             else:
                 # not param or buffer then user input
                 input_specs.append(