A4: shared cost-aware MMA selector (lift ROCm's model to all targets) - #309
Conversation
Promote ROCm's per-arch matrix-core model into one lane-count-parameterized
selector so NVIDIA / Apple gain the cost-aware MMA selection they lacked, while
ROCm stays the byte-identical reference (COMPILER_REFACTOR_PLAN A4).
compiler/mma_selector.py:
- MmaIsa: arch-neutral ISA record (target/arch/mma_class, cooperative +
lane_count, legal shape table, dtype->K).
- accumulator_regs_per_lane(shape, lane_count) = M*N // lane_count — the
promoted footprint, delegating to rocm_target.mfma_accumulator_regs so the
arithmetic + validation stay single-sourced with the lead.
- rank_shapes_by_footprint / cheapest_shape / select_mma -> MmaSelection (the
target-neutral chosen-MMA descriptor + operands + per-lane footprint).
- Per-arch ISA records: ROCm built FROM the live rocm_target/rocm_mma tables
(feature gates inherited by construction — no FP8 WMMA on gfx1151, no fp32
WMMA on RDNA — never re-encoded); NVIDIA mma.sync m16n8k{8,16,32,64} (warp 32);
Apple simdgroup_matrix<8,8> (simdgroup 32); x86 AMX as tile-register (honestly
cooperative=False / lane_count=None, per-lane footprint N/A).
Lead-safe (Theory rule #1): no emit path changes — hardware-free selector
(Decision #19). test_mma_selector.py (19 tests) gates that the shared footprint
ranking on a ROCm ISA is byte-identical to ROCm's own rank_mfma_shapes_by_footprint,
and that the ROCm ISA's dtype set matches exactly what rocm_mma.select_mma accepts.
Cross-target selection verified: nvidia bf16/fp16->m16n8k16 (4 regs/lane), fp8->k32,
fp4->k64; apple fp16->8x8x8 (2 regs); x86 amx bf16->16x16x32 (None). ruff + mypy
clean; rocm_mma consumers unaffected (59 tests); generated-doc drift gate green.
Follow-on (own PR — touches drift-gated dashboards): wire get_isa/select_mma as
backend_manifest's cross-target MMA-metadata source + the D1 arbiter cost_model.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eba55474bb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| shapes=((16, 16, 32),), | ||
| k_by_dtype={"bf16": 32, "int8": 64}, |
There was a problem hiding this comment.
For get_isa("x86", "amx"), this advertises int8 in isa.dtypes with K=64, but the only AMX shape listed has K=32. As a result, select_mma(M.get_isa("x86", "amx"), "int8") always raises MmaSelectorError from cheapest_shape even though the ISA record says int8 is supported; either include the legal int8 shape or remove the dtype mapping.
Useful? React with 👍 / 👎.
| if prefer_shape not in isa.shapes: | ||
| raise MmaSelectorError( | ||
| f"prefer_shape {prefer_shape} is not legal on {isa.target}:" | ||
| f"{isa.arch} (legal: {sorted(isa.shapes)})") | ||
| shape = prefer_shape |
There was a problem hiding this comment.
When callers pass prefer_shape, this only checks that the tuple exists somewhere on the ISA, not that its K matches isa.k_by_dtype[dtype]. On multi-K ISAs this can return impossible descriptors, e.g. NVIDIA bf16 with prefer_shape=(16, 8, 32) produces a bf16 selection at K=32 even though bf16 is mapped to K=16, so downstream metadata/lowering can see an invalid dtype/shape pairing instead of the promised MmaSelectorError.
Useful? React with 👍 / 👎.
1 participant
Workstream A4 (
COMPILER_REFACTOR_PLAN): the one place a lead abstraction lifts upward. ROCm already models matrix-core selection as a cost-aware choice — a per-arch shape table ranked by theM×N // lanesaccumulator footprint. NVIDIA, Apple, and x86 had no cost-aware MMA selector. This promotes ROCm's model into onelane_count-parameterized selector so every cooperative-matrix target gets the same footprint-ranked selection, keyed by(target, arch, dtype)— theshape_table/cost_modelthe D1 arbiter parks here (Decision #28).compiler/mma_selector.pyMmaIsa— arch-neutral ISA record:mma_class,cooperative+lane_count, the legal(M,N,K)shape table, dtype→K.accumulator_regs_per_lane(shape, lane_count)=M*N // lane_count— the promoted footprint, delegating torocm_target.mfma_accumulator_regsso the arithmetic + validation stay single-sourced with the lead.rank_shapes_by_footprint/cheapest_shape/select_mma → MmaSelection— the cost model + the target-neutral chosen-MMA descriptor (shape anchor + derived nt operands + per-lane footprint).rocm_target/rocm_mmatables; NVIDIAmma.syncm16n8k{8,16,32,64} (warp 32); Applesimdgroup_matrix<8,8>(simdgroup 32); x86 AMX tile-register (honestlycooperative=False/lane_count=None).Lead-safety (Theory rule #1)
ROCm stays the reference and cannot be silently perturbed:
rocm_mma.select_mmaitself, so every feature gate — no FP8 WMMA on gfx1151, no fp32 WMMA on RDNA, fp4/xf32 gating — is inherited by construction, never re-encoded.test_mma_selector.pygates thatrank_shapes_by_footprinton a ROCm ISA is byte-identical to ROCm's ownrank_mfma_shapes_by_footprint, and that the ISA's dtype set is exactly whatrocm_mma.select_mmaaccepts.The lift (verified)
Verification
test_mma_selector.py— 19 tests (ROCm equivalence, cross-target selection, feature gates, error paths).rocm_mmaconsumers unaffected (59 tests); generated-doc drift gate green (no dashboard changes — fully additive).Follow-on (own PR — it touches drift-gated dashboards): wire
get_isa/select_mmaasbackend_manifest's cross-target MMA-metadata source (today ROCm-only via_rocm_mma_descriptor_for) and as the D1 arbiter'scost_model.🤖 Generated with Claude Code