From 912d80da7263534a2c5d452605220d89d968f852 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Thu, 17 Sep 2026 14:04:18 -0700 Subject: [PATCH] Pool immutable integer constants during EXIR emission Repeated literal operator arguments allocate separate EValue slots, and integer lists also duplicate persistent boxed-list storage. Pool read-only integer literals and lists made entirely from literal integers within each execution plan, using the argument schema to exclude out and aliased values. Preserve existing abstract values, method inputs and outputs, symbolic values, and opaque delegate arguments. Construct references with their final pooled IDs during emission so no serialization-time remapping is needed. Add tests for type distinctions, signed 64-bit limits, list references, mutation and aliasing boundaries, plan scope, and repeated control flow. Update the keyword-argument test to follow instruction argument IDs. Validation: 78 emitter tests and 14 subtests passed, with two existing skips. All four trained MLPerf Tiny models produced exact before/after outputs on 67,029 Cortex-M55 FVP validation inputs. Operator calls and planned memory were unchanged, and independent re-exports were byte-identical. DS-CNN persistent method metadata decreased from 8,008 to 5,056 measured bytes. Authored with OpenAI Codex. --- exir/emit/_emitter.py | 57 ++++++- exir/emit/test/BUCK | 2 + exir/emit/test/test_emit.py | 8 +- exir/emit/test/test_emit_constants.py | 232 ++++++++++++++++++++++++++ 4 files changed, 289 insertions(+), 10 deletions(-) create mode 100644 exir/emit/test/test_emit_constants.py diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py index 4e3baf26592..cd46ba52c60 100644 --- a/exir/emit/_emitter.py +++ b/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)) diff --git a/exir/emit/test/BUCK b/exir/emit/test/BUCK index 79f2134d191..06f21991cdb 100644 --- a/exir/emit/test/BUCK +++ b/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", diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index d6273b01d90..78480d3b55c 100644 --- a/exir/emit/test/test_emit.py +++ b/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 diff --git a/exir/emit/test/test_emit_constants.py b/exir/emit/test/test_emit_constants.py new file mode 100644 index 00000000000..123272565e3 --- /dev/null +++ b/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()