Pool immutable integer constants during EXIR emission

exir/emit/_emitter.py

@@ -150,6 +150,8 @@ class _EmitterState:
     emit_mutable_buffer_names: bool
 
     spec2id_dict: Dict[TensorSpec, int] = field(default_factory=dict)
+    # Only literal arguments with no schema aliasing enter this per-plan pool.
+    constant_cache: Dict[Union[int, Tuple[int, ...]], int] = field(default_factory=dict)
 
     def spec2id(self, spec: TensorSpec) -> int:
         """Map a TensorSpec to value index in the values array."""
@@ -270,7 +272,9 @@ def _internal_assert_emitter(
         if not pred:
             raise InternalError(self._emit_node_specific_error(node, assert_msg))
 
-    def _emit_int_list(self, val: List[_Argument]) -> EValue:
+    def _emit_int_list(
+        self, val: List[_Argument], *, immutable: bool = False
+    ) -> EValue:
         """Emits a list of integers as a collection of EValues.
 
         For every argument in 'val':
@@ -286,7 +290,7 @@ def _emit_int_list(self, val: List[_Argument]) -> EValue:
                 boxed_list.append(item.id)
             elif isinstance(item, int):
                 boxed_list.append(
-                    self._emit_evalue(self._constant_to_evalue(item, None)).id
+                    self._emit_argument(item, None, immutable=immutable).id
                 )
             else:
                 self._internal_assert_emitter(
@@ -295,7 +299,13 @@ def _emit_int_list(self, val: List[_Argument]) -> EValue:
 
         return EValue(IntList(boxed_list))
 
-    def _emit_list(self, val: List[_Argument], val_type: _SchemaType) -> EValue:
+    def _emit_list(
+        self,
+        val: List[_Argument],
+        val_type: _SchemaType,
+        *,
+        immutable: bool = False,
+    ) -> EValue:
         """Emits a list type.
 
         Emits the list stored in val. If the list is of Tensors, Optionals, or Ints the emitted list
@@ -308,7 +318,7 @@ def _emit_list(self, val: List[_Argument], val_type: _SchemaType) -> EValue:
             return EValue(BoolList(typing.cast(List[bool], val)))
 
         if isinstance(val_type, torch.IntType):
-            return self._emit_int_list(val)
+            return self._emit_int_list(val, immutable=immutable)
 
         if isinstance(val_type, torch.FloatType):
             return EValue(DoubleList(typing.cast(List[float], val)))
@@ -541,6 +551,8 @@ def _constant_to_evalue(  # noqa: C901
         self,
         val: _Argument,
         val_type: Optional[_SchemaType],
+        *,
+        immutable: bool = False,
     ) -> EValue:
         """Converts a constant value to an EValue.
 
@@ -564,6 +576,7 @@ def _constant_to_evalue(  # noqa: C901
             return self._emit_list(
                 typing.cast(List[_Argument], val),
                 typing.cast(_SchemaType, val_type.getElementType()),
+                immutable=immutable,
             )
 
         if isinstance(val, float):
@@ -1353,13 +1366,37 @@ def _add_delegate_map(
         }
 
     def _emit_argument(
-        self, arg: _Argument, arg_type: Optional[_SchemaType]
+        self,
+        arg: _Argument,
+        arg_type: Optional[_SchemaType],
+        *,
+        immutable: bool = False,
     ) -> _AbstractValue:
         """Emit an argument to an operator or delegate if it had not already been emitted otherwise
         return the previously emitted location"""
         if isinstance(arg, _AbstractValue):
             return arg
-        return self._emit_evalue(self._constant_to_evalue(arg, arg_type))
+        value = self._constant_to_evalue(arg, arg_type, immutable=immutable)
+        key: Optional[Union[int, Tuple[int, ...]]] = None
+        if immutable:
+            if isinstance(value.val, Int):
+                key = value.val.int_val
+            elif (
+                isinstance(value.val, IntList)
+                and isinstance(arg, (list, tuple))
+                and all(type(item) is int for item in arg)
+            ):
+                # Boxed lists can reference mutable SymInts. Only literal lists
+                # may share their unboxed buffer; their elements are already pooled.
+                key = tuple(value.val.items)
+        if key is not None:
+            index = self.emitter_state.constant_cache.get(key)
+            if index is not None:
+                return _AbstractValue(index, None)
+        result = self._emit_evalue(value)
+        if key is not None:
+            self.emitter_state.constant_cache[key] = result.id
+        return result
 
     def _get_sym_ret(
         self,
@@ -1537,7 +1574,13 @@ def _get_empty_tensor_evalue() -> EValue:
             if kernel_arg is None and isinstance(schema_arg.type, torch.TensorType):
                 kernel_arg = self._emit_evalue(_get_empty_tensor_evalue())
 
-            kernel_args.append(self._emit_argument(kernel_arg, schema_arg.type).id)
+            kernel_args.append(
+                self._emit_argument(
+                    kernel_arg,
+                    schema_arg.type,
+                    immutable=not schema_arg.is_out and schema_arg.alias_info is None,
+                ).id
+            )
 
             if schema_arg.is_out:
                 out_args.append((schema_arg.name, kernel_arg))

exir/emit/test/BUCK

@@ -15,6 +15,7 @@ fbcode_target(_kind = runtime.python_test,
     name = "emit",
     srcs = [
         "test_emit.py",
+        "test_emit_constants.py",
     ],
     deps = [
         "fbsource//third-party/pypi/pytest:pytest",
@@ -24,6 +25,7 @@ fbcode_target(_kind = runtime.python_test,
         "//executorch/exir:lib",
         "//executorch/exir:print_program",
         "//executorch/exir:schema",
+        "//executorch/exir:tensor",
         "//executorch/exir/backend/test/demos/rpc:executor_backend_partitioner",
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/backend:compile_spec_schema",

exir/emit/test/test_emit.py

@@ -427,9 +427,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             .to_executorch()
             .executorch_program
         )
-        # The value for beta should appear before alpha
-        self.assertEqual(program.execution_plan[0].values[12].val, Int(3))
-        self.assertEqual(program.execution_plan[0].values[13].val, Int(2))
+        plan = program.execution_plan[0]
+        call = plan.chains[0].instructions[-1].instr_args
+        self.assertEqual(plan.operators[call.op_index].name, "aten::addbmm")
+        self.assertEqual(plan.values[call.args[3]].val, Int(3))
+        self.assertEqual(plan.values[call.args[4]].val, Int(2))
 
     def test_kwargs2(self) -> None:
         """Tests that the kwargs are placed in the order specified by

exir/emit/test/test_emit_constants.py

@@ -0,0 +1,232 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+import unittest
+
+import torch
+from executorch.exir import to_edge
+from executorch.exir._serialize._program import deserialize_pte_binary
+from executorch.exir.emit._emitter import _Emitter, _EmitterState, _ProgramState
+from executorch.exir.schema import Bool, Double, EValue, Int, IntList
+from executorch.exir.tensor import TensorSpec
+from executorch.extension.pybindings.portable_lib import (
+    _load_for_executorch_from_buffer,
+)
+from torch._higher_order_ops import cond, map as torch_map
+from torch.export import export
+
+
+class TestEmitConstants(unittest.TestCase):
+    def make_emitter(self, state=None):
+        graph = torch.fx.Graph()
+        node = graph.placeholder("x")
+        graph.output(node)
+        module = torch.fx.GraphModule({}, graph)
+        module.meta["non_const_buffer_sizes"] = [0, 0]
+        if state is None:
+            state = _EmitterState([], [], [], {}, False, False)
+        emitter = _Emitter(module, state, _ProgramState())
+        emitter.node = node
+        return emitter
+
+    def test_scalar_types_and_signed_int64(self):
+        emitter = self.make_emitter()
+        for value in (-(2**63), -128, 0, 1, 2**63 - 1):
+            with self.subTest(value=value):
+                first = emitter._emit_argument(value, None, immutable=True)
+                second = emitter._emit_argument(value, None, immutable=True)
+                self.assertEqual(first.id, second.id)
+                self.assertEqual(emitter.emitter_state.values[first.id].val, Int(value))
+
+        integer = emitter._emit_argument(1, None, immutable=True)
+        boolean = emitter._emit_argument(True, None, immutable=True)
+        double = emitter._emit_argument(1.0, None, immutable=True)
+        self.assertEqual(len({integer.id, boolean.id, double.id}), 3)
+        self.assertEqual(emitter.emitter_state.values[boolean.id].val, Bool(True))
+        self.assertEqual(emitter.emitter_state.values[double.id].val, Double(1.0))
+
+    def test_literal_lists_and_element_references(self):
+        emitter = self.make_emitter()
+        list_type = torch.ListType.ofInts()
+        values = emitter.emitter_state.values
+        scalar = emitter._emit_argument(1, None, immutable=True)
+        pair = emitter._emit_argument([1, 1], list_type, immutable=True)
+        repeated = emitter._emit_argument((1, 1), list_type, immutable=True)
+        self.assertEqual(pair.id, repeated.id)
+        self.assertEqual(values[pair.id].val, IntList([scalar.id, scalar.id]))
+        for items in ([0, 1], [1, 0], [1], []):
+            with self.subTest(items=items):
+                value = emitter._emit_argument(items, list_type, immutable=True)
+                optional = emitter._emit_argument(
+                    items, torch.OptionalType(list_type), immutable=True
+                )
+                self.assertEqual(value.id, optional.id)
+                self.assertNotEqual(value.id, pair.id)
+                self.assertEqual(
+                    [values[index].val.int_val for index in values[value.id].val.items],
+                    items,
+                )
+        self.assertEqual(sum(isinstance(value.val, Int) for value in values), 2)
+        self.assertEqual(sum(isinstance(value.val, IntList) for value in values), 5)
+
+    def test_serialized_signed_int64(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                return (
+                    torch.clamp(x, min=-(2**63), max=2**63 - 1),
+                    torch.clamp(x.flip(0), min=-(2**63), max=2**63 - 1),
+                )
+
+        model = Model()
+        inputs = (torch.tensor([-(2**63), 0, 2**63 - 1]),)
+        program = to_edge(export(model, inputs, strict=True)).to_executorch()
+        plan = deserialize_pte_binary(program.buffer).program.execution_plan[0]
+        for limit in (-(2**63), 2**63 - 1):
+            self.assertEqual(sum(value.val == Int(limit) for value in plan.values), 1)
+        runtime = _load_for_executorch_from_buffer(program.buffer)
+        for actual, expected in zip(runtime.forward(inputs), model(*inputs)):
+            torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    def test_dynamic_lists_preserve_mutable_elements(self):
+        emitter = self.make_emitter()
+        values = emitter.emitter_state.values
+        symbol = emitter._emit_evalue(EValue(Int(1)))
+        literal = emitter._emit_argument(1, None, immutable=True)
+        self.assertNotEqual(symbol.id, literal.id)
+        self.assertEqual(
+            emitter._emit_argument(symbol, None, immutable=True).id, symbol.id
+        )
+        lists = [
+            emitter._emit_argument([symbol, 1], torch.ListType.ofInts(), immutable=True)
+            for _ in range(2)
+        ]
+        self.assertNotEqual(lists[0].id, lists[1].id)
+        for value in lists:
+            self.assertEqual(values[value.id].val.items, [symbol.id, literal.id])
+        values[symbol.id] = EValue(Int(7))
+        self.assertEqual(values[literal.id].val, Int(1))
+        for value in lists:
+            self.assertEqual(
+                [values[index].val.int_val for index in values[value.id].val.items],
+                [7, 1],
+            )
+
+    def test_opaque_arguments_do_not_enter_pool(self):
+        emitter = self.make_emitter()
+        values = emitter.emitter_state.values
+        for arg, arg_type in ((1, None), ([1, 1], torch.ListType.ofInts())):
+            with self.subTest(arg=arg):
+                first = emitter._emit_argument(arg, arg_type)
+                second = emitter._emit_argument(arg, arg_type)
+                pooled = emitter._emit_argument(arg, arg_type, immutable=True)
+                self.assertEqual(len({first.id, second.id, pooled.id}), 3)
+                if isinstance(arg, list):
+                    self.assertTrue(
+                        set(values[first.id].val.items).isdisjoint(
+                            values[pooled.id].val.items
+                        )
+                    )
+
+    def test_operator_alias_and_mutation_boundaries(self):
+        emitter = self.make_emitter()
+        emitter.node.meta["spec"] = TensorSpec.from_tensor(torch.ones(2))
+        tensor = emitter._emit_spec(emitter.node.meta["spec"])
+        with torch.library._scoped_library("emit_constant_test", "FRAGMENT") as library:
+            for name, scalar, items in (
+                ("read", "int", "int[]"),
+                ("alias", "int(a)", "int[](b)"),
+                ("mutate", "int(a!)", "int[](b!)"),
+            ):
+                library.define(
+                    f"{name}.out(Tensor x, {scalar} scalar, {items} items, "
+                    "*, Tensor(c!) out) -> Tensor(c!)"
+                )
+            for op, should_pool in (
+                (torch.ops.emit_constant_test.read.out, True),
+                (torch.ops.emit_constant_test.alias.out, False),
+                (torch.ops.emit_constant_test.mutate.out, False),
+            ):
+                with self.subTest(op=op):
+                    for _ in range(2):
+                        emitter._emit_operator(op, (tensor, 1, [1, 1]), {"out": tensor})
+                    first, second = [
+                        instruction.instr_args.args
+                        for instruction in emitter.chain.instructions[-2:]
+                    ]
+                    for index in (1, 2):
+                        self.assertEqual(first[index] == second[index], should_pool)
+                    if not should_pool:
+                        items = emitter.emitter_state.values[first[2]].val.items
+                        self.assertNotIn(first[1], items)
+                        self.assertTrue(
+                            set(items).isdisjoint(
+                                emitter.emitter_state.values[second[2]].val.items
+                            )
+                        )
+
+    def test_input_output_and_nested_containers(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, number, flag):
+                return {"tensors": [x + 1, x - 1], "constants": (1, [number, flag])}
+
+        inputs = (torch.ones(2), 1, True)
+        program = to_edge(export(Model(), inputs, strict=True)).to_executorch()
+        plan = deserialize_pte_binary(program.buffer).program.execution_plan[0]
+        input_int = plan.inputs[1]
+        self.assertEqual(plan.values[input_int].val, Int(1))
+        self.assertEqual(plan.values[plan.inputs[2]].val, Bool(True))
+        literal_ids = [
+            index
+            for instruction in plan.chains[0].instructions
+            for index in instruction.instr_args.args
+            if plan.values[index].val == Int(1)
+        ]
+        self.assertGreaterEqual(len(literal_ids), 2)
+        self.assertEqual(len(set(literal_ids)), 1)
+        self.assertNotIn(input_int, literal_ids)
+        self.assertTrue(set(plan.outputs).isdisjoint(literal_ids))
+        runtime = _load_for_executorch_from_buffer(program.buffer)
+        result = runtime.forward(inputs)
+        torch.testing.assert_close(result[0], inputs[0] + 1, rtol=0, atol=0)
+        torch.testing.assert_close(result[1], inputs[0] - 1, rtol=0, atol=0)
+        self.assertEqual(result[2:], [1, 1, True])
+
+    def test_pool_scope(self):
+        emitter = self.make_emitter()
+        first = emitter._emit_argument(7, None, immutable=True)
+        subgraph = self.make_emitter(emitter.emitter_state)
+        self.assertEqual(subgraph._emit_argument(7, None, immutable=True).id, first.id)
+        other_method = self.make_emitter()
+        other_method._emit_evalue(EValue(Int(0)))
+        other = other_method._emit_argument(7, None, immutable=True)
+        self.assertNotEqual(other.id, first.id)
+        self.assertEqual(other_method.emitter_state.values[other.id].val, Int(7))
+
+    def test_repeated_control_flow_execution(self):
+        class Model(torch.nn.Module):
+            def forward(self, pred, xs):
+                def body(x):
+                    return cond(pred, lambda x: x + 1, lambda x: x - 1, (x,))
+
+                return torch_map(body, xs)
+
+        model = Model()
+        xs = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+        program = to_edge(
+            export(model, (torch.tensor(True), xs), strict=True)
+        ).to_executorch()
+        runtime = _load_for_executorch_from_buffer(program.buffer)
+        for pred in (True, False, True):
+            inputs = (torch.tensor(pred), xs)
+            torch.testing.assert_close(
+                runtime.forward(inputs)[0], model(*inputs), rtol=0, atol=0
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()