Skip to content

Commit 273ecf9

Browse files
authored
[CoreAI] Record ordered I/O names in the delegate manifest (pytorch#23278)
The Core AI runtime binds ExecuTorch inputs and outputs to the .aimodel function's inputs and outputs by name. This PR makes the exporter name the delegate boundary explicitly and record those names, in boundary order, in the delegate manifest, so the runtime never relies on dictionary iteration order. * Explicit boundary names (compiler/preprocess.py): preprocess passes input_0…input_N and output_0…output_M to coreai-torch, one per ExecuTorch user input and output as reported by _edge_io. State arguments for delegate-owned mutated buffers aren't boundary I/O, so the converter keeps its FX placeholder names for them. * Manifest bindings: source (inline) and AOT (aot_compiled_inline) manifests now include: * function: the entrypoint name, "main". * input_names / output_names: read by the new coreai_io_names (compiler/io_compat.py) in graph argument and result order. State arguments, and the results they name, are skipped. * Name check: assert_io_compatible now also compares those names with what ExecuTorch expects. That catches I/O that is reordered but has the same dtype and shape, which the existing dtype and shape checks can't tell apart. * Enumerated shapes: names are read before apply_enumerated_shapes runs, because coreai-core 1.0.0b3's set_static_shape_config replaces main with per-shape main_<key> clones (which keep the same names). apply_enumerated_shapes now takes the converter's names as a required argument, so the shape config matches the renamed inputs. The backend is under construction, so the manifest format carries no version or compatibility guarantees.
1 parent a0ca769 commit 273ecf9

6 files changed

Lines changed: 163 additions & 23 deletions

File tree

‎backends/apple/coreai/compiler/enumerated_shapes.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -173,15 +173,17 @@ def apply_enumerated_shapes(
173173
program,
174174
edge_program: ExportedProgram,
175175
enumerations: Dict[str, Sequence[int]],
176+
*,
177+
input_names: Sequence[str],
176178
) -> None:
177179
"""Attach enumerated shapes to the coreai program (all delivery modes).
178180
179181
``enumerations`` is ``{symbol_name: [value, ...]}`` (from
180182
:func:`resolve_input_enumerations`). For each of this subgraph's user inputs
181183
we read its symbolic shape from ``edge_program`` and substitute every
182184
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+
via ``set_static_shape_config``. ``input_names`` are the names the converter
186+
was given for the user inputs, in boundary order.
185187
"""
186188
if not enumerations:
187189
return
@@ -198,17 +200,17 @@ def apply_enumerated_shapes(
198200
# (a batch dim, say) would not be held equal.
199201
boundary_inputs = []
200202
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:
203+
for uname, converted_name in zip(user_inputs, input_names):
204+
# Skip anything that isn't a coreai tensor input (e.g. a taken-over
205+
# constant).
206+
if not isinstance(uname, str) or converted_name not in coreai_names:
205207
continue
206208
node = placeholders.get(uname)
207209
val = node.meta.get("val") if node is not None else None
208210
if not hasattr(val, "shape"):
209211
continue
210212
dims = list(val.shape)
211-
boundary_inputs.append((uname, dims))
213+
boundary_inputs.append((converted_name, dims))
212214
active_symbols |= _shape_symbols(dims) & set(enumerations)
213215

214216
if not active_symbols:

‎backends/apple/coreai/compiler/io_compat.py‎

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
"""
1414

1515
import logging
16-
from typing import Any, List, Sequence, Tuple
16+
from typing import Any, List, Optional, Sequence, Tuple
1717

1818
import torch
1919
from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT
@@ -180,7 +180,49 @@ def io_mismatches(
180180
return errors
181181

182182

183-
def assert_io_compatible(program, edge_program: ExportedProgram) -> None:
183+
def coreai_io_names(program) -> Tuple[List[str], List[str]]:
184+
"""Boundary I/O names in graph argument/result order, never dictionary order.
185+
186+
State arguments (marked ``MutableBuffers.buffer_mutation``) and the results
187+
they name are skipped: they are not ExecuTorch boundary I/O.
188+
"""
189+
graph = program._get_graph(MAIN_ENTRYPOINT)
190+
arg_attrs = list(graph.arg_attrs)
191+
res_attrs = list(graph.res_attrs)
192+
for kind, attrs, types in (
193+
("input", arg_attrs, graph.function_type.value.inputs),
194+
("output", res_attrs, graph.function_type.value.results),
195+
):
196+
if len(attrs) != len(types) or any("coreai.name" not in a for a in attrs):
197+
raise ValueError(f"Core AI {kind} names must be complete and unique")
198+
state_results = {
199+
a["MutableBuffers.buffer_mutation"].value
200+
for a in arg_attrs
201+
if "MutableBuffers.buffer_mutation" in a
202+
}
203+
inputs = [
204+
a["coreai.name"].value
205+
for a in arg_attrs
206+
if "MutableBuffers.buffer_mutation" not in a
207+
]
208+
outputs = [
209+
a["coreai.name"].value
210+
for a in res_attrs
211+
if a["coreai.name"].value not in state_results
212+
]
213+
for kind, names in (("input", inputs), ("output", outputs)):
214+
if len(set(names)) != len(names):
215+
raise ValueError(f"Core AI {kind} names must be complete and unique")
216+
return inputs, outputs
217+
218+
219+
def assert_io_compatible(
220+
program,
221+
edge_program: ExportedProgram,
222+
*,
223+
input_names: Optional[Sequence[str]] = None,
224+
output_names: Optional[Sequence[str]] = None,
225+
) -> None:
184226
"""Raise if the ``.aimodel`` boundary I/O is incompatible with ExecuTorch."""
185227
coreai_in, coreai_out = _coreai_io(program)
186228
edge_in, edge_out = _edge_io(edge_program)
@@ -198,6 +240,17 @@ def assert_io_compatible(program, edge_program: ExportedProgram) -> None:
198240
errors = io_mismatches(coreai_in, edge_in, "input") + io_mismatches(
199241
coreai_out, edge_out, "output"
200242
)
243+
if input_names is not None or output_names is not None:
244+
actual_inputs, actual_outputs = coreai_io_names(program)
245+
for kind, actual, expected in (
246+
("input", actual_inputs, input_names),
247+
("output", actual_outputs, output_names),
248+
):
249+
if expected is not None and actual != list(expected):
250+
errors.append(
251+
f"{kind} names mismatch: .aimodel has {actual}, "
252+
f"ExecuTorch expects {list(expected)}"
253+
)
201254
if errors:
202255
raise ValueError(
203256
"Core AI delegate boundary is incompatible with ExecuTorch:\n "

‎backends/apple/coreai/compiler/preprocess.py‎

Lines changed: 46 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,11 @@
2424
from executorch.backends.apple.coreai.compiler.enumerated_shapes import (
2525
apply_enumerated_shapes,
2626
)
27-
from executorch.backends.apple.coreai.compiler.io_compat import assert_io_compatible
27+
from executorch.backends.apple.coreai.compiler.io_compat import (
28+
_edge_io,
29+
assert_io_compatible,
30+
coreai_io_names,
31+
)
2832
from executorch.backends.apple.coreai.passes.replace_copy_ops import (
2933
ReplaceCopyOpsWithFunctionalPass,
3034
)
@@ -319,12 +323,22 @@ def _prepare_program_for_conversion(edge_program: ExportedProgram) -> ExportedPr
319323
return ep
320324

321325

322-
def _convert_to_aiprogram(edge_program: ExportedProgram):
326+
def _convert_to_aiprogram(
327+
edge_program: ExportedProgram,
328+
*,
329+
input_names: Optional[List[str]] = None,
330+
output_names: Optional[List[str]] = None,
331+
):
323332
from coreai_torch import TorchConverter
324333

325334
aten_program = _prepare_program_for_conversion(edge_program)
326335
converter = TorchConverter()
327-
converter.add_exported_program(aten_program)
336+
converter.add_exported_program(
337+
aten_program,
338+
input_names=input_names,
339+
output_names=output_names,
340+
entrypoint_name=MAIN_ENTRYPOINT,
341+
)
328342
return converter.to_coreai()
329343

330344

@@ -558,16 +572,33 @@ def preprocess(
558572
"build-time output directory (set it via coreai_sidecar_dir)"
559573
)
560574

561-
program = _convert_to_aiprogram(edge_program)
575+
edge_inputs, edge_outputs = _edge_io(edge_program)
576+
input_names = [f"input_{i}" for i in range(len(edge_inputs))]
577+
output_names = [f"output_{i}" for i in range(len(edge_outputs))]
578+
program = _convert_to_aiprogram(
579+
edge_program, input_names=input_names, output_names=output_names
580+
)
562581
# Fail fast if the .aimodel boundary I/O won't match what ET feeds/reads.
563-
assert_io_compatible(program, edge_program)
582+
assert_io_compatible(
583+
program, edge_program, input_names=input_names, output_names=output_names
584+
)
585+
# Enumerated shapes replace ``main`` with renamed clones, so read first.
586+
bound_inputs, bound_outputs = coreai_io_names(program)
564587
raw_enum = _get_compile_spec(
565588
compile_specs, COMPILE_SPEC_KEYS.INPUT_ENUMERATIONS
566589
)
567590
if raw_enum:
568591
apply_enumerated_shapes(
569-
program, edge_program, json.loads(raw_enum.decode())
592+
program,
593+
edge_program,
594+
json.loads(raw_enum.decode()),
595+
input_names=input_names,
570596
)
597+
bindings = {
598+
"function": MAIN_ENTRYPOINT,
599+
"input_names": bound_inputs,
600+
"output_names": bound_outputs,
601+
}
571602

572603
# min-deployment-version is a single knob applied to whichever artifact
573604
# ships: for aot-compiled delivery the .aimodelc (via coreai-build
@@ -579,18 +610,19 @@ def preprocess(
579610
_maybe_warn_sidecar_env_ignored()
580611
if aot_compiled:
581612
return CoreAIBackend._preprocess_aot_compiled(
582-
program, _aot_compile_options(compile_specs), sidecar_dir
613+
program, _aot_compile_options(compile_specs), sidecar_dir, bindings
583614
)
584615
return CoreAIBackend._preprocess_portable(
585616
program,
586617
_min_os_version(compile_specs),
587618
sidecar_dir,
619+
bindings,
588620
)
589621

590622
# Portable .aimodel delivery.
591623
@staticmethod
592624
def _preprocess_portable(
593-
program, min_os, sidecar_dir: Optional[str]
625+
program, min_os, sidecar_dir: Optional[str], bindings: Dict[str, Any]
594626
) -> PreprocessResult:
595627
with TemporaryDirectory() as tmp:
596628
model_hash = _save_and_hash(program, Path(tmp) / "model.aimodel", min_os)
@@ -599,6 +631,7 @@ def _preprocess_portable(
599631
model_hash,
600632
AssetPackaging.SIDECAR if sidecar_dir else AssetPackaging.INLINE,
601633
{
634+
**bindings,
602635
# relative path the runtime resolves against its base
603636
"path": f"{model_hash}/model.aimodel",
604637
"min_deployment_version": _os_version_text(min_os),
@@ -609,7 +642,10 @@ def _preprocess_portable(
609642
# AOT-compiled .aimodelc delivery (per architecture).
610643
@staticmethod
611644
def _preprocess_aot_compiled(
612-
program, opts: Dict[str, Any], sidecar_dir: Optional[str]
645+
program,
646+
opts: Dict[str, Any],
647+
sidecar_dir: Optional[str],
648+
bindings: Dict[str, Any],
613649
) -> PreprocessResult:
614650
with TemporaryDirectory() as tmp:
615651
model_hash, out = _compile_aot(program, opts, Path(tmp))
@@ -622,6 +658,7 @@ def _preprocess_aot_compiled(
622658
else AssetPackaging.AOT_COMPILED_INLINE
623659
),
624660
{
661+
**bindings,
625662
"platform": opts["platform"],
626663
"min_deployment_version": opts["min_deployment_version"],
627664
"archs": {

‎backends/apple/coreai/test/test_io_compat.py‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,13 @@
2222
get_default_compile_config,
2323
get_default_passes,
2424
)
25-
from executorch.backends.apple.coreai.compiler.io_compat import io_mismatches
25+
from executorch.backends.apple.coreai.compiler.io_compat import (
26+
assert_io_compatible,
27+
io_mismatches,
28+
)
29+
from executorch.backends.apple.coreai.compiler.preprocess import _convert_to_aiprogram
2630
from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner
27-
from executorch.exir import to_edge_transform_and_lower
31+
from executorch.exir import to_edge, to_edge_transform_and_lower
2832
from executorch.exir.lowered_backend_module import executorch_call_delegate
2933
from torch.export import Dim
3034

@@ -33,6 +37,28 @@ class _Sym:
3337
"""Stand-in for a symbolic (dynamic) edge dim."""
3438

3539

40+
class OrderedNamesTest(unittest.TestCase):
41+
def test_same_shaped_reordered_names_are_rejected(self):
42+
"""Dtype and shape checks can't tell same-shaped I/O apart; names can."""
43+
44+
class Subtract(nn.Module):
45+
def forward(self, z, a):
46+
return z - a, a - z
47+
48+
edge = to_edge(
49+
torch.export.export(Subtract(), (torch.ones(2, 4), torch.zeros(2, 4)))
50+
).exported_program()
51+
inputs = ["input_0", "input_1"]
52+
outputs = ["output_0", "output_1"]
53+
program = _convert_to_aiprogram(
54+
edge, input_names=inputs[::-1], output_names=outputs
55+
)
56+
with self.assertRaisesRegex(ValueError, "input names mismatch"):
57+
assert_io_compatible(
58+
program, edge, input_names=inputs, output_names=outputs
59+
)
60+
61+
3662
class IoMismatchesTest(unittest.TestCase):
3763
def test_io_mismatches(self):
3864
# (name, coreai_io, edge_io, expected substring or None if compatible)

‎backends/apple/coreai/test/test_preprocess.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,21 @@ def test_end_to_end_sidecar(self):
238238
self.assertTrue((bundles[0] / "main.mlirb").exists())
239239

240240

241+
class ManifestBindingsTest(unittest.TestCase):
242+
def test_names_follow_boundary_order_past_single_digits(self):
243+
class ManyInputs(nn.Module):
244+
def forward(self, *args):
245+
return tuple(x - float(i) for i, x in enumerate(args))
246+
247+
edge = to_edge(
248+
torch.export.export(ManyInputs(), tuple(torch.ones(2) for _ in range(12)))
249+
).exported_program()
250+
manifest = json.loads(CoreAIBackend.preprocess(edge, []).processed_bytes)
251+
self.assertEqual(manifest["function"], "main")
252+
self.assertEqual(manifest["input_names"], [f"input_{i}" for i in range(12)])
253+
self.assertEqual(manifest["output_names"], [f"output_{i}" for i in range(12)])
254+
255+
241256
def _lower_with_break(partitioner):
242257
"""Lower a model whose middle op is tagged, forcing two delegates."""
243258
from executorch.backends.apple.coreai import get_default_passes

‎backends/apple/coreai/test/test_shape_config.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -294,7 +294,9 @@ def test_inputs_sharing_a_symbol_are_enumerated_together(self):
294294
for key, shapes in entries.items():
295295
with self.subTest(key):
296296
self.assertEqual(
297-
set(shapes), {"x", "y"}, f"{key} leaves an input unconstrained"
297+
set(shapes),
298+
{"input_0", "input_1"},
299+
f"{key} leaves an input unconstrained",
298300
)
299301
self.assertEqual(len({tuple(s) for s in shapes.values()}), 1)
300302

@@ -329,7 +331,7 @@ def test_every_boundary_input_is_named(self):
329331
for cfg in captured:
330332
for key, shapes in cfg.items():
331333
with self.subTest(key):
332-
self.assertEqual(set(shapes), {"x", "y"})
334+
self.assertEqual(set(shapes), {"input_0", "input_1"})
333335

334336

335337
class SubgraphBoundaryTest(unittest.TestCase):
@@ -356,7 +358,12 @@ def _lower_split(self):
356358
len(captured), 2, "expected a graph break into exactly 2 delegates"
357359
)
358360
# Identify by content rather than order, which preprocess does not fix.
359-
model_input = next(c for c in captured if "x" in next(iter(c.values())))
361+
model_input = next(
362+
c
363+
for c in captured
364+
if sorted(tuple(s) for v in c.values() for s in v.values())
365+
== [(4, 8), (16, 8)]
366+
)
360367
intermediate = next(c for c in captured if c is not model_input)
361368
return model_input, intermediate
362369

0 commit comments

Comments
 (0)