|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Enumerated input-shape propagation for the Core AI backend. |
| 8 | +
|
| 9 | +Enumerations are declared on the *ET model inputs*, resolved to ``torch.export`` |
| 10 | +symbols (partitioner side, :func:`resolve_input_enumerations`), then substituted |
| 11 | +into each delegate subgraph's boundary symbolic shapes and attached via |
| 12 | +``AIProgram.set_static_shape_config`` (preprocess side, |
| 13 | +:func:`apply_enumerated_shapes`). |
| 14 | +""" |
| 15 | + |
| 16 | +import logging |
| 17 | +from itertools import product |
| 18 | +from typing import Dict, List, Optional, Sequence, Set |
| 19 | + |
| 20 | +import sympy |
| 21 | + |
| 22 | +from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT |
| 23 | +from torch.export.exported_program import ExportedProgram |
| 24 | + |
| 25 | +logger = logging.getLogger(__name__) |
| 26 | + |
| 27 | + |
| 28 | +def resolve_input_enumerations( |
| 29 | + exported_program: ExportedProgram, |
| 30 | + input_enumerations: Optional[Sequence[Optional[Dict[int, Sequence[int]]]]], |
| 31 | +) -> Optional[Dict[str, List[int]]]: |
| 32 | + """Map ET-input ``(index, dim)`` enumerations to ``{symbol_name: [values]}``. |
| 33 | +
|
| 34 | + Reads the export symbol backing each enumerated ET-input dim from the |
| 35 | + exported program; raises if the dim is static (no symbol) or is anything |
| 36 | + other than a bare symbol. |
| 37 | + """ |
| 38 | + if not input_enumerations: |
| 39 | + return None |
| 40 | + placeholders = { |
| 41 | + n.name: n for n in exported_program.graph.nodes if n.op == "placeholder" |
| 42 | + } |
| 43 | + user_inputs = list(exported_program.graph_signature.user_inputs) |
| 44 | + resolved: Dict[str, List[int]] = {} |
| 45 | + for i, dim_map in enumerate(input_enumerations): |
| 46 | + if not dim_map: |
| 47 | + continue |
| 48 | + if i >= len(user_inputs): |
| 49 | + raise ValueError( |
| 50 | + f"input_enumerations has an entry for input {i} but the model " |
| 51 | + f"has {len(user_inputs)} input(s)" |
| 52 | + ) |
| 53 | + val = ( |
| 54 | + placeholders[user_inputs[i]].meta.get("val") |
| 55 | + if (isinstance(user_inputs[i], str) and user_inputs[i] in placeholders) |
| 56 | + else None |
| 57 | + ) |
| 58 | + if not hasattr(val, "shape"): |
| 59 | + raise ValueError( |
| 60 | + f"input {i} is not a tensor input; cannot enumerate its shape" |
| 61 | + ) |
| 62 | + for dim_index, values in dim_map.items(): |
| 63 | + if not 0 <= dim_index < len(val.shape): |
| 64 | + raise ValueError( |
| 65 | + f"input {i} has {len(val.shape)} dim(s); cannot enumerate " |
| 66 | + f"dim {dim_index}" |
| 67 | + ) |
| 68 | + values = list(values) |
| 69 | + if not values: |
| 70 | + raise ValueError( |
| 71 | + f"input {i} dim {dim_index} has no enumerated values; drop " |
| 72 | + "the entry instead of passing an empty list" |
| 73 | + ) |
| 74 | + dim = val.shape[dim_index] |
| 75 | + if isinstance(dim, int): |
| 76 | + raise ValueError( |
| 77 | + f"input {i} dim {dim_index} is static ({dim}); export it " |
| 78 | + "with dynamic_shapes to enumerate it" |
| 79 | + ) |
| 80 | + # Only a bare symbol can be enumerated: apply_enumerated_shapes |
| 81 | + # substitutes these values for the symbol itself, so a derived dim |
| 82 | + # (2*batch) would pin 2*value rather than value, and the range |
| 83 | + # check below would bound the wrong quantity. |
| 84 | + expr = dim.node.expr |
| 85 | + if not isinstance(expr, sympy.Symbol): |
| 86 | + raise ValueError( |
| 87 | + f"input {i} dim {dim_index} shape {expr} is a derived or " |
| 88 | + "compound expression, not a plain Dim; enumerate the dim(s) " |
| 89 | + "it is derived from instead" |
| 90 | + ) |
| 91 | + symbol = expr |
| 92 | + _check_within_range(exported_program, symbol, values, i, dim_index) |
| 93 | + name = str(symbol) |
| 94 | + if name in resolved and resolved[name] != values: |
| 95 | + raise ValueError( |
| 96 | + f"input {i} dim {dim_index} shares symbol {name} with an " |
| 97 | + f"earlier input but enumerates different values " |
| 98 | + f"{values} vs {resolved[name]}; dims tied to one Dim must " |
| 99 | + "agree" |
| 100 | + ) |
| 101 | + resolved[name] = values |
| 102 | + return resolved |
| 103 | + |
| 104 | + |
| 105 | +def _check_within_range(exported_program, symbol, values, input_index, dim_index): |
| 106 | + """Reject values the exported dim cannot take. |
| 107 | +
|
| 108 | + Out-of-range values are accepted by the shape config and only fail later, |
| 109 | + when the compiler or runtime rejects a specialization it cannot map back to |
| 110 | + any input. |
| 111 | + """ |
| 112 | + bad = [v for v in values if v <= 0] |
| 113 | + if bad: |
| 114 | + raise ValueError( |
| 115 | + f"input {input_index} dim {dim_index} enumerates non-positive " |
| 116 | + f"value(s) {bad}" |
| 117 | + ) |
| 118 | + constraint = exported_program.range_constraints.get(symbol) |
| 119 | + if constraint is None: |
| 120 | + return |
| 121 | + low, high = constraint.lower, constraint.upper |
| 122 | + outside = [v for v in values if v < low or v > high] |
| 123 | + if outside: |
| 124 | + raise ValueError( |
| 125 | + f"input {input_index} dim {dim_index} enumerates {outside}, " |
| 126 | + f"outside the exported range [{low}, {high}]" |
| 127 | + ) |
| 128 | + |
| 129 | + |
| 130 | +def _shape_symbols(shape) -> Set[str]: |
| 131 | + """Names of the export symbols appearing in a (possibly symbolic) shape.""" |
| 132 | + symbols = set() |
| 133 | + for dim in shape: |
| 134 | + if not isinstance(dim, int): |
| 135 | + symbols |= {str(s) for s in dim.node.expr.free_symbols} |
| 136 | + return symbols |
| 137 | + |
| 138 | + |
| 139 | +def graph_input_names(program, entrypoint: str = MAIN_ENTRYPOINT) -> List[str]: |
| 140 | + """Names of the converted coreai graph's inputs (in argument order).""" |
| 141 | + graph = program.get_graph(entrypoint) |
| 142 | + names = [] |
| 143 | + for attr in graph.arg_attrs: |
| 144 | + if "coreai.name" in attr: |
| 145 | + names.append(attr["coreai.name"].value) |
| 146 | + return names |
| 147 | + |
| 148 | + |
| 149 | +def _eval_dim(dim, subs: Dict[str, int], uname: str, axis: int) -> int: |
| 150 | + """Resolve a possibly-symbolic shape dim to a concrete int under ``subs``. |
| 151 | +
|
| 152 | + A dim can mix an enumerated symbol with one that was left dynamic (batch |
| 153 | + enumerated, sequence not). Substitution leaves that symbolic, and ``int()`` |
| 154 | + would raise ``Cannot convert symbols to int`` naming neither the input nor |
| 155 | + the dim, so report it here instead. |
| 156 | + """ |
| 157 | + if isinstance(dim, int): |
| 158 | + return dim |
| 159 | + expr = dim.node.expr |
| 160 | + mapping = {s: subs[str(s)] for s in expr.free_symbols if str(s) in subs} |
| 161 | + resolved = expr.subs(mapping) |
| 162 | + if resolved.free_symbols: |
| 163 | + raise ValueError( |
| 164 | + f"input '{uname}' dim {axis} is {expr}, which still depends on " |
| 165 | + f"{sorted(str(s) for s in resolved.free_symbols)} after " |
| 166 | + f"substitution; enumerate those dims too, or leave this input " |
| 167 | + f"fully dynamic" |
| 168 | + ) |
| 169 | + return int(resolved) |
| 170 | + |
| 171 | + |
| 172 | +def apply_enumerated_shapes( |
| 173 | + program, |
| 174 | + edge_program: ExportedProgram, |
| 175 | + enumerations: Dict[str, Sequence[int]], |
| 176 | +) -> None: |
| 177 | + """Attach enumerated shapes to the coreai program (all delivery modes). |
| 178 | +
|
| 179 | + ``enumerations`` is ``{symbol_name: [value, ...]}`` (from |
| 180 | + :func:`resolve_input_enumerations`). For each of this subgraph's user inputs |
| 181 | + we read its symbolic shape from ``edge_program`` and substitute every |
| 182 | + combination of the enumerated symbol values, then attach the resulting shapes |
| 183 | + via ``set_static_shape_config``. Inputs are matched to coreai graph inputs by |
| 184 | + name (the converter names each coreai input after its edge placeholder). |
| 185 | + """ |
| 186 | + if not enumerations: |
| 187 | + return |
| 188 | + |
| 189 | + placeholders = { |
| 190 | + n.name: n for n in edge_program.graph.nodes if n.op == "placeholder" |
| 191 | + } |
| 192 | + user_inputs = list(edge_program.graph_signature.user_inputs) |
| 193 | + coreai_names = set(graph_input_names(program)) |
| 194 | + |
| 195 | + # Every entry has to name every boundary input, since |
| 196 | + # set_static_shape_config treats a key as a whole-graph shape: an input left |
| 197 | + # out is unconstrained in that specialization, and inputs sharing a symbol |
| 198 | + # (a batch dim, say) would not be held equal. |
| 199 | + boundary_inputs = [] |
| 200 | + active_symbols = set() |
| 201 | + for uname in user_inputs: |
| 202 | + # Match edge user inputs to coreai graph inputs by name; skip anything |
| 203 | + # that isn't a coreai tensor input (e.g. a taken-over constant). |
| 204 | + if not isinstance(uname, str) or uname not in coreai_names: |
| 205 | + continue |
| 206 | + node = placeholders.get(uname) |
| 207 | + val = node.meta.get("val") if node is not None else None |
| 208 | + if not hasattr(val, "shape"): |
| 209 | + continue |
| 210 | + dims = list(val.shape) |
| 211 | + boundary_inputs.append((uname, dims)) |
| 212 | + active_symbols |= _shape_symbols(dims) & set(enumerations) |
| 213 | + |
| 214 | + if not active_symbols: |
| 215 | + # Nothing here is enumerated, so there is nothing to specialize. Routine |
| 216 | + # under multi-partition lowering, and legitimate when the enumerated |
| 217 | + # input is not delegated at all. |
| 218 | + return |
| 219 | + |
| 220 | + # One axis per enumerated symbol, shared across every input that uses it. |
| 221 | + # _eval_dim rejects any boundary dim left symbolic, so a specialization |
| 222 | + # cannot be attached while some other input stays dynamic. |
| 223 | + active = [s for s in enumerations if s in active_symbols] |
| 224 | + shapes_config: Dict[str, Dict[str, tuple]] = {} |
| 225 | + for combo in product(*[enumerations[s] for s in active]): |
| 226 | + subs = dict(zip(active, combo)) |
| 227 | + # Name the specialization after the values it pins. The key becomes an |
| 228 | + # MLIR dictionary attribute name, so it has to stay an identifier: no |
| 229 | + # "=" or other punctuation. |
| 230 | + key = "_".join(f"{s}_{v}" for s, v in subs.items()) |
| 231 | + shapes_config[key] = { |
| 232 | + uname: tuple(_eval_dim(d, subs, uname, axis) for axis, d in enumerate(dims)) |
| 233 | + for uname, dims in boundary_inputs |
| 234 | + } |
| 235 | + |
| 236 | + if shapes_config: |
| 237 | + # Every combination is a separate compile, so the cartesian product can |
| 238 | + # make a build much slower than the enumeration lists suggest. |
| 239 | + logger.info( |
| 240 | + "attaching %d shape specialization(s) over symbol(s) %s", |
| 241 | + len(shapes_config), |
| 242 | + active, |
| 243 | + ) |
| 244 | + program.set_static_shape_config(MAIN_ENTRYPOINT, shapes_config) |
0 commit comments