From 697232f5fd8134685e21f146cda517d99dbd20f9 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 30 Jul 2026 17:22:52 -0700 Subject: [PATCH] [ExecuTorch][Core AI] Enumerated input-shape support --- .../coreai/compiler/enumerated_shapes.py | 141 ++++++++++++++++++ backends/apple/coreai/compiler/preprocess.py | 18 +++ .../apple/coreai/partition/partitioner.py | 33 +++- .../apple/coreai/test/test_shape_config.py | 105 +++++++++++++ 4 files changed, 296 insertions(+), 1 deletion(-) create mode 100644 backends/apple/coreai/compiler/enumerated_shapes.py create mode 100644 backends/apple/coreai/test/test_shape_config.py diff --git a/backends/apple/coreai/compiler/enumerated_shapes.py b/backends/apple/coreai/compiler/enumerated_shapes.py new file mode 100644 index 00000000000..743e193f500 --- /dev/null +++ b/backends/apple/coreai/compiler/enumerated_shapes.py @@ -0,0 +1,141 @@ +# 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. + +"""Enumerated input-shape propagation for the Core AI backend. + +Enumerations are declared on the *ET model inputs*, resolved to ``torch.export`` +symbols (partitioner side, :func:`resolve_input_enumerations`), then substituted +into each delegate subgraph's boundary symbolic shapes and attached via +``AIProgram.set_static_shape_config`` (preprocess side, +:func:`apply_enumerated_shapes`). +""" + +from itertools import product +from typing import Dict, List, Optional, Sequence + +from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT +from torch.export.exported_program import ExportedProgram + + +def resolve_input_enumerations( + exported_program: ExportedProgram, + input_enumerations: Optional[Sequence[Optional[Dict[int, Sequence[int]]]]], +) -> Optional[Dict[str, List[int]]]: + """Map ET-input ``(index, dim)`` enumerations to ``{symbol_name: [values]}``. + + Reads the export symbol backing each enumerated ET-input dim from the + exported program; raises if the dim is static (no symbol) or maps to a + non-single-symbol expression. + """ + if not input_enumerations: + return None + placeholders = { + n.name: n for n in exported_program.graph.nodes if n.op == "placeholder" + } + user_inputs = list(exported_program.graph_signature.user_inputs) + resolved: Dict[str, List[int]] = {} + for i, dim_map in enumerate(input_enumerations): + if not dim_map: + continue + if i >= len(user_inputs): + raise ValueError( + f"input_enumerations has an entry for input {i} but the model " + f"has {len(user_inputs)} input(s)" + ) + val = ( + placeholders[user_inputs[i]].meta.get("val") + if (isinstance(user_inputs[i], str) and user_inputs[i] in placeholders) + else None + ) + if not hasattr(val, "shape"): + raise ValueError( + f"input {i} is not a tensor input; cannot enumerate its shape" + ) + for dim_index, values in dim_map.items(): + dim = val.shape[dim_index] + if isinstance(dim, int): + raise ValueError( + f"input {i} dim {dim_index} is static ({dim}); export it " + "with dynamic_shapes to enumerate it" + ) + symbols = dim.node.expr.free_symbols + if len(symbols) != 1: + raise ValueError( + f"input {i} dim {dim_index} shape {dim.node.expr} is not a " + "single symbol; cannot enumerate" + ) + resolved[str(next(iter(symbols)))] = list(values) + return resolved + + +def graph_input_names(program, entrypoint: str = MAIN_ENTRYPOINT) -> List[str]: + """Names of the converted coreai graph's inputs (in argument order).""" + graph = program.get_graph(entrypoint) + names = [] + for attr in graph.arg_attrs: + if "coreai.name" in attr: + names.append(attr["coreai.name"].value) + return names + + +def _eval_dim(dim, subs: Dict[str, int]) -> int: + """Resolve a possibly-symbolic shape dim to a concrete int under ``subs``.""" + if isinstance(dim, int): + return dim + expr = dim.node.expr + mapping = {s: subs[str(s)] for s in expr.free_symbols if str(s) in subs} + return int(expr.subs(mapping)) + + +def apply_enumerated_shapes( + program, + edge_program: ExportedProgram, + enumerations: Dict[str, Sequence[int]], +) -> None: + """Attach enumerated shapes to the coreai program (all delivery modes). + + ``enumerations`` is ``{symbol_name: [value, ...]}`` (from + :func:`resolve_input_enumerations`). For each of this subgraph's user inputs + we read its symbolic shape from ``edge_program`` and substitute every + combination of the enumerated symbol values, then attach the resulting shapes + via ``set_static_shape_config``. Inputs are matched to coreai graph inputs by + name (the converter names each coreai input after its edge placeholder). + """ + if not enumerations: + return + + placeholders = { + n.name: n for n in edge_program.graph.nodes if n.op == "placeholder" + } + user_inputs = list(edge_program.graph_signature.user_inputs) + coreai_names = set(graph_input_names(program)) + + shapes_config: Dict[str, Dict[str, tuple]] = {} + for uname in user_inputs: + # Match edge user inputs to coreai graph inputs by name; skip anything + # that isn't a coreai tensor input (e.g. a taken-over constant). + if not isinstance(uname, str) or uname not in coreai_names: + continue + node = placeholders.get(uname) + val = node.meta.get("val") if node is not None else None + if not hasattr(val, "shape"): + continue + dims = list(val.shape) + # Enumerated symbols this input's shape actually depends on. + symbols = set() + for d in dims: + if not isinstance(d, int): + symbols |= {str(s) for s in d.node.expr.free_symbols} + active = [s for s in enumerations if s in symbols] + if not active: + continue + for k, combo in enumerate(product(*[enumerations[s] for s in active])): + subs = dict(zip(active, combo)) + concrete = tuple(_eval_dim(d, subs) for d in dims) + shapes_config[f"{uname}_{k}"] = {uname: concrete} + + if shapes_config: + program.set_static_shape_config(MAIN_ENTRYPOINT, shapes_config) diff --git a/backends/apple/coreai/compiler/preprocess.py b/backends/apple/coreai/compiler/preprocess.py index ef890368a9d..0fcc2095816 100644 --- a/backends/apple/coreai/compiler/preprocess.py +++ b/backends/apple/coreai/compiler/preprocess.py @@ -21,6 +21,9 @@ from typing import Any, Dict, final, Iterator, List, Optional, Tuple from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT +from executorch.backends.apple.coreai.compiler.enumerated_shapes import ( + apply_enumerated_shapes, +) from executorch.backends.apple.coreai.compiler.io_compat import assert_io_compatible from executorch.backends.apple.coreai.passes.replace_copy_ops import ( ReplaceCopyOpsWithFunctionalPass, @@ -58,6 +61,14 @@ class COMPILE_SPEC_KEYS(Enum): # applies to every delivery path, AOT or not. Numeric string, e.g. "27.0". MIN_DEPLOYMENT_VERSION = "coreai_min_deployment_version" + # Enumerated input shapes for this delegate, derived by the partitioner from + # ET-input enumerations and propagated to this subgraph via torch.export + # symbols. General (applies to every delivery mode, AOT or not). JSON: + # {symbol_name: [value, ...]} # e.g. {"s31": [4, 16, 32]} + # Preprocess substitutes these into each user-input placeholder's symbolic + # shape and attaches the results via AIProgram.set_static_shape_config. + INPUT_ENUMERATIONS = "coreai_input_enumerations" + # Ahead-of-time compilation (xcrun coreai-build). # Presence implies compiling the .aimodel to per-architecture .aimodelc # bundles at build time (requires macOS + the Metal Toolchain). The value is @@ -435,6 +446,13 @@ def preprocess( program = _convert_to_aiprogram(edge_program) # Fail fast if the .aimodel boundary I/O won't match what ET feeds/reads. assert_io_compatible(program, edge_program) + raw_enum = _get_compile_spec( + compile_specs, COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS + ) + if raw_enum: + apply_enumerated_shapes( + program, edge_program, json.loads(raw_enum.decode()) + ) # min-deployment-version is a single knob applied to whichever artifact # ships: for aot-compiled delivery the .aimodelc (via coreai-build diff --git a/backends/apple/coreai/partition/partitioner.py b/backends/apple/coreai/partition/partitioner.py index af70c22a6c8..d07297e0fd2 100644 --- a/backends/apple/coreai/partition/partitioner.py +++ b/backends/apple/coreai/partition/partitioner.py @@ -5,10 +5,14 @@ # LICENSE file in the root directory of this source tree. import inspect +import json import logging -from typing import Callable, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Sequence, Tuple import torch +from executorch.backends.apple.coreai.compiler.enumerated_shapes import ( + resolve_input_enumerations, +) from executorch.backends.apple.coreai.compiler.preprocess import ( AOTCompileConfig, @@ -196,6 +200,9 @@ def __init__( uses_sidecar: bool = False, aot_compile_config: Optional[AOTCompileConfig] = None, min_deployment_version: Optional[str] = None, + input_enumerations: Optional[ + Sequence[Optional[Dict[int, Sequence[int]]]] + ] = None, take_over_constant_data: bool = True, take_over_mutable_buffer: bool = True, ) -> None: @@ -212,6 +219,11 @@ def __init__( # single JSON AOT_COMPILE_CONFIG spec (presence implies AOT). # min_deployment_version is a general spec (it also sets the portable # .aimodel's OS floor), so it is emitted separately. + # + # input_enumerations declares enumerated shapes on the ET model inputs + # ({dim_index: [values]}, aligned to inputs); it is resolved to export + # symbols in :meth:`partition` and propagated to each subgraph boundary + # in preprocess, so it is not built into a compile spec here. specs = [] if uses_sidecar: specs.append(CompileSpec(COMPILE_SPEC_KEYS.USES_SIDECAR.value, b"1")) @@ -234,6 +246,8 @@ def __init__( backend_id=CoreAIBackend.__name__, compile_specs=specs, ) + self._base_compile_specs = specs + self._input_enumerations = input_enumerations self.take_over_constant_data = take_over_constant_data self.take_over_mutable_buffer = take_over_mutable_buffer @@ -261,7 +275,24 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: ) logger.info("CoreAIPartitioner::partition") partition_tags = {} + + # Resolve ET-input enumerations to export symbols and attach as a compile + # spec; preprocess propagates them to each subgraph's boundary. delegation_spec = self.delegation_spec + resolved = resolve_input_enumerations( + exported_program, self._input_enumerations + ) + if resolved: + delegation_spec = DelegationSpec( + backend_id=CoreAIBackend.__name__, + compile_specs=self._base_compile_specs + + [ + CompileSpec( + COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS.value, + json.dumps(resolved).encode(), + ) + ], + ) capability_partitioner = CapabilityBasedPartitioner( exported_program.graph_module, diff --git a/backends/apple/coreai/test/test_shape_config.py b/backends/apple/coreai/test/test_shape_config.py new file mode 100644 index 00000000000..468135b5f9b --- /dev/null +++ b/backends/apple/coreai/test/test_shape_config.py @@ -0,0 +1,105 @@ +# 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. + +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.apple.coreai.compiler.enumerated_shapes import ( + resolve_input_enumerations, +) +from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import executorch_call_delegate + + +def _model(): + return nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 8)).eval() + + +def _dynamic_ep(): + batch = torch.export.Dim("batch", min=1, max=64) + return torch.export.export( + _model(), (torch.randn(2, 8),), dynamic_shapes={"input": {0: batch}} + ) + + +def _static_ep(): + return torch.export.export(_model(), (torch.randn(2, 8),)) + + +class ResolveInputEnumerationsTest(unittest.TestCase): + def test_maps_dynamic_dim_to_symbol(self): + resolved = resolve_input_enumerations(_dynamic_ep(), [{0: [4, 16, 32]}]) + self.assertEqual(len(resolved), 1) + self.assertEqual(list(resolved.values()), [[4, 16, 32]]) + # key is the export symbol name (e.g. "s31") + (sym,) = resolved.keys() + self.assertTrue(sym.startswith("s")) + + def test_static_dim_raises(self): + with self.assertRaises(ValueError): + resolve_input_enumerations(_static_ep(), [{0: [4, 16]}]) + + def test_no_enumerations_is_none(self): + self.assertIsNone(resolve_input_enumerations(_dynamic_ep(), None)) + + +class InputEnumerationsLoweringTest(unittest.TestCase): + def _lower(self, input_enumerations): + return to_edge_transform_and_lower( + _dynamic_ep(), + partitioner=[CoreAIPartitioner(input_enumerations=input_enumerations)], + ) + + def test_lowers_to_executorch(self): + lowered = self._lower([{0: [4, 16, 32]}]) + gm = lowered.exported_program().graph_module + self.assertTrue( + any( + n.op == "call_function" and n.target is executorch_call_delegate + for n in gm.graph.nodes + ) + ) + self.assertGreater(len(bytes(lowered.to_executorch().buffer)), 0) + + def test_shapes_are_substituted_at_boundary(self): + # Verify the ET-input enumeration propagates + substitutes into the + # delegate boundary's symbolic shape (s31, 8) -> (4,8)/(16,8)/(32,8). + from coreai.authoring import AIProgram + + captured = [] + orig = AIProgram.set_static_shape_config + + def _record(self, entrypoint, shapes_config): + captured.append(shapes_config) + return orig(self, entrypoint, shapes_config) + + AIProgram.set_static_shape_config = _record + try: + self._lower([{0: [4, 16, 32]}]).to_executorch() + finally: + AIProgram.set_static_shape_config = orig + + self.assertTrue(captured, "set_static_shape_config was not called") + shapes = { + tuple(next(iter(entry.values()))) + for cfg in captured + for entry in cfg.values() + } + self.assertEqual(shapes, {(4, 8), (16, 8), (32, 8)}) + + def test_static_dim_fails_lowering(self): + with self.assertRaisesRegex(Exception, "static"): + to_edge_transform_and_lower( + _static_ep(), + partitioner=[CoreAIPartitioner(input_enumerations=[{0: [4, 16]}])], + ) + + +if __name__ == "__main__": + unittest.main()