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.

Pull Request resolved: #1897

Reviewed By: chakriu, angelayi

Differential Revision: D53584903

Pulled By: tarun292

fbshipit-source-id: 9349bc20216f9ffe877fc6e32df17409f1131e83

Author: tarun292

exir/emit/_emitter.py

@@ -1308,18 +1308,30 @@ 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:
+                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) -> None:
+        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(