Add support for outputting constants (#1774)

Summary:

For Seamless model we ran into a case where we have to either return an int/float or a list of int's/float's from the model. We need to add support for this in the emitter and make sure that the memory planning pass ignores these.

Reviewed By: angelayi

Differential Revision: D53256808

Author: tarun292

exir/emit/_emitter.py

@@ -1369,9 +1369,19 @@ def output(
             self.outputs.append(args_tuple.id)
         else:
             for arg in args_tuple:
-                # Every output should already have its value emitted outputs should only be abstract
-                # IDs at this point.
-                assert isinstance(arg, _AbstractValue)
+                if isinstance(arg, (int, float, bool)):
+                    arg = self._emit_evalue(self._constant_to_evalue(arg, None))
+                elif isinstance(arg, (type(None), str)):
+                    raise InternalError(
+                        self._emit_node_specific_error(
+                            self.node,
+                            f"Returning {arg} is not yet supported in the emitter.",
+                        )
+                    )
+                else:
+                    # Every other output should already have its value emitted.
+                    # They should only be abstract IDs at this point.
+                    assert isinstance(arg, _AbstractValue)
                 self.outputs.append(arg.id)
 
     def plan(self) -> ExecutionPlan:

exir/emit/test/test_emit.py

@@ -209,6 +209,39 @@ def forward(
             "T2#1#0(T1#1($),D0())",
         )
 
+    def test_constant_output(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return [((1, 3, 1.2), True, [x + x, x * x])]
+
+        ep = torch.export.export(M(), (torch.ones(2, 3),))
+        res = ep(torch.ones(2, 3))
+        self.assertEqual(res[0][0], (1, 3, 1.2))
+        program = to_edge(ep).to_executorch().executorch_program
+        outputs = program.execution_plan[0].outputs
+        self.assertEqual(len(outputs), 6)
+        self.assertEqual(program.execution_plan[0].values[outputs[0]].val.int_val, 1)
+        self.assertEqual(program.execution_plan[0].values[outputs[1]].val.int_val, 3)
+        self.assertEqual(
+            program.execution_plan[0].values[outputs[2]].val.double_val, 1.2
+        )
+        self.assertEqual(
+            program.execution_plan[0].values[outputs[2]].val.bool_val, True
+        )
+
+    def test_int_list_input(self):
+        class M(torch.nn.Module):
+            def forward(self, x, y, z):
+                return x + y, x + x, x + y + z
+
+        ep = torch.export.export(M(), (torch.ones(2, 3), 2, True))
+        ep(torch.ones(2, 3), 2, True)
+        program = to_edge(ep).to_executorch().executorch_program
+        inputs = program.execution_plan[0].inputs
+        self.assertEqual(len(inputs), 3)
+        self.assertEqual(program.execution_plan[0].values[inputs[1]].val.int_val, 2)
+        self.assertEqual(program.execution_plan[0].values[inputs[2]].val.bool_val, True)
+
     def test_buffers_with_perfect_alignment(self) -> None:
         class Foo(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor:

exir/memory_planning.py

@@ -495,7 +495,11 @@ def get_node_tensor_specs(
     if not isinstance(specs, (list, tuple)):
         return []
     else:
-        return specs
+        return [
+            spec
+            for spec in specs
+            if not isinstance(spec, (int, float, bool, str, type(None)))
+        ]
 
 
 @register_algo