[ExecuTorch][Core AI][5/x] Enumerated input-shape support

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)

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

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,

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