Add support for outputting constants

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.

Differential Revision: D53256808

Author: tarun292

exir/emit/_emitter.py

@@ -1342,9 +1342,12 @@ 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)):
+                    arg = self._emit_evalue(self._constant_to_evalue(arg, None))
+                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

@@ -187,6 +187,23 @@ def forward(
             "T2#1#0(T1#1($),D0())",
         )
 
+    def test_int_list_output(self):
+        class M(torch.nn.Module):
+            def forward(self, x):
+                return [((1, 3, 1.2), [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), 5)
+        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
+        )
+
     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

@@ -453,7 +453,7 @@ 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))]
 
 
 @register_algo