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()