feat(backend): calibrate the a5 (Ascend950) L0 GEMM cost model from a5-sim#2001
feat(backend): calibrate the a5 (Ascend950) L0 GEMM cost model from a5-sim#2001tonibohnlein wants to merge 10 commits into
Conversation
…line) Make Backend950::GetL0CostModel() return an EXPLICIT L0CostModel instead of the inherited default, so each a5 roofline constant becomes an editable slot for a5-sim calibration. Every value currently equals the a2a3 default, so this is a behavioural NO-OP -- a5 tile picks are unchanged (verified: a5 differs from a2a3 only via its 256 KB L0C, which is already correct and grounded in pto-isa buffer_limits.hpp). The a5 L0 capacities (L0A/L0B 64 KB, L0C 256 KB, Mat 512 KB) match the pto-isa hardware reference and are what already make a5 tile differently. Only the cost CONSTANTS (bandwidths, drain, mad_head) remain a2a3-derived; the model FORM (per-M-row max(floor,throughput) + oddPart misalignment) is arch-general. The override carries the calibration recipe (per-pipe a5-sim sweep + transfer the a2a3 op-sim->device correction) for the follow-up. hw-native-sys#1912 a5-calibration follow-up.
device_scripts/a5_cost_model_calib.py: runs the compiled chooser with the a5 L0 capacities under A2A3_BASELINE (what a5 inherits today) vs an editable A5_FITTED constant set, over a shape grid, and prints every tile pick that changes. This is the sim-free validation of an a5 calibration's EFFECT on picks. Ships with the A5_FITTED slot equal to the baseline (0 diffs) until an a5-sim sweep fills it. Companion to the Backend950 override scaffold.
The cube MAD cost derived cpr = bytes_a/2 (fp32 -> 2 passes/K-fractal), baking in a2a3's fp32 decomposition. The a5-sim shows a5 prices full fp32 MMAD ~4x a2a3 per fractal (MMAD(128,128,64) = 1045 a2a3 vs 4121 a5 = 21 + 8*512), i.e. a5 does a fuller fp32 mantissa decomposition. Per pto-isa the fractal WIDTH (K0=C0/E, C0=32) is arch-invariant, so the pass COUNT is the only arch-varying fp32 cube term. Add L0CostModel::mad_fp32_passes (default 2 = a2a3, byte-identical) and use it for the fp32 branch of cpr; Backend950 sets 8 (a5-sim measured). Wired through L0TileConfig + binding + stub + the Python oracle. a2a3 picks unchanged (427 UTs pass); a5 fp32 picks rebalance (17/180, all fp32 -- bf16 untouched). Since pypto emits full fp32 (no tf32_mode), this 4x is the real a5 shipping cost. Part of the a5 cost-model calibration (hw-native-sys#1912 follow-up); unblocks the magnitude refit which a5's cube form otherwise invalidated.
…rmed a5-sim re-run confirmed the cube calibration (the primary, high-confidence result): - mad_fp32_passes=8 CONFIRMED — fp32 mmad = 25 + 512*k_fractal (4x a2a3's 128) at m=128 (4121) and m=256 (8217); - bf16 stays cpr=1 (bf16 mmad=281=25+256, unaffected by the 4x) — resolves the open "is bf16 also 4x" question: the 8x is fp32-only, exactly as mad_fp32_passes intends; - mad_head_cycles=25 (a5 intercept of both fp32 and bf16 k-sweeps; a2a3=21). a2a3 unchanged (a5-only override; 385 UTs pass). The bw + 4 drain constants remain a2a3-inherited: a5-sim was too slow to fit them this pass (LOAD source is LOAD_2Dv2/MTE1; drain n-sweep timed out) -- documented as inheritance, no worse than status quo, refit pending (see a5_cost_model_device_task.md). Cube fix drives the 17/180 fp32 pick changes in the audit.
…ansfer) The a5-sim work-cycle extraction (excluding WAIT/BAR stalls) reproduces the shipped a2a3 *device* model within ~15%, so raw a5-sim work-cycle numbers are fit directly -- the earlier "transfer the op-sim->device BW/~1.54 correction" note is superseded (it was for raw pipe sums). Also drop the stale "behavioural no-op" claim (cube is now a5-calibrated: mad_head=25, mad_fp32_passes=8). No a5 device -> a5-sim is the ground truth; validate via a5-sim OLD-vs-NEW wall. Full spec in a5_cost_model_device_task.md.
Complete the a5 (Ascend950) roofline calibration from a5-sim work-cycle sweeps
(no a5 device -> a5-sim is ground truth; raw fit, fits the sweeps to <1%):
load: bw_l0a 129.7->206.3, bw_l0b 85.4->223.8 (a5 L0 ports faster; via LOAD_2Dv2)
drain: bw_drain 118->30, drain_fixed 164->343.4, drain_row 4.45->4.59,
drain_penalty 2.6->0.26 (a5 misalignment much milder; oddPart form re-confirmed)
cube (already landed): mad_head 25, mad_fp32_passes 8, bf16 cpr 1
a2a3 unchanged (a5-only override; 385 UTs pass). Full a5 pick-diff vs the
a2a3-inherited baseline is now 59/180 (was 17 cube-only), shifting to wider-N
(drain_fixed doubled). Single-tile roofline ranking validated on a5-sim
(Spearman ~0.9).
KNOWN CAVEAT (documented at the constant): bw_drain=30 carries a ~4x sim-vs-device
gap -- a5-sim FIXPIPE drain isn't device-anchored (a2a3-sim drain is likewise ~4x
a2a3-device). It over-predicts wide-N single-tile walls ~34% because a5-sim
pipelines the drain under compute; if it mis-picks tiles the fix is the drain-hidden
max(compute,drain) path, not a constant edit. Multi-tile a5-sim wall tooling is
broken (span under-counts), so validation is single-tile / work-cycle-lane only.
device_scripts/a5_validate_picks.py: per shape, enumerate EVERY legal (m,n,k,stationarity) tile candidate, score each with the a5-calibrated model (reusing the validated chooser-UT oracle), rank, and mark the chooser's pick. The device agent measures each candidate's full-matmul a5-sim wall and confirms the model's hw-native-sys#1 is the simulator's hw-native-sys#1 -- the argmin-match validation (stronger than the single-tile ranking correlation). Top model candidates are within ~0.4%, so the sim measurement is what discriminates. Blocked on the multi-tile a5-sim wall tooling fix (see follow-ups).
a5-sim validation (Exp A/B/C) resolved the bw_drain caveat: - Exp A: the chooser's dbC=1 picks are the sim's hw-native-sys#2-hw-native-sys#3, within ~1-2% of the measured-lane best -- the ~34% absolute over-prediction preserves the ranking, so it does NOT mis-pick. - Exp B: making dbC=2 the a5 default is REFUTED -- its full-K constraint forces tiny fp32 tiles -> more drains -> ~40% fp32 LOSS (bf16 only ~1.25x). So the expensive a5 drain is NOT fixed by a schedule-default flip. - Exp C: emit is tight (bubbles <10%, operand DB, K-reuse); low cube util is a5's drain-bound reality (FIXP 44-52% of the wall), not an emit defect. Update the comment (was: "the fix is dbC=2") to reflect this.
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds calibrated Ascend950 L0 cost-model parameters, introduces FP32 MAD-pass handling across backend and tile selection, updates the scoring oracle, and adds scripts to audit and validate tile choices. ChangesAscend950 L0 Cost Model
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Ascend950Handler
participant AnalyzeMatmul
participant ChooseL0Tile
Ascend950Handler->>AnalyzeMatmul: provide calibrated L0 cost model
AnalyzeMatmul->>ChooseL0Tile: pass mad_fp32_passes in L0TileConfig
ChooseL0Tile-->>AnalyzeMatmul: select FP32-aware L0 tile
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Code Review
This pull request implements a fully calibrated L0 cost model for the Ascend950 (a5) architecture, introducing arch-specific FP32 cube passes (mad_fp32_passes) across the C++ tile chooser, Python bindings, and test suites, alongside new calibration and validation scripts. The review feedback suggests guarding mad_fp32_passes with a minimum value of 1 in both the C++ implementation and Python test oracle to ensure robustness against non-positive configurations.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| const int64_t cpr = (cfg.bytes_a == 4) ? static_cast<int64_t>(cfg.mad_fp32_passes) | ||
| : std::max<int64_t>(1, static_cast<int64_t>(cfg.bytes_a) / 2); |
There was a problem hiding this comment.
If cfg.mad_fp32_passes is misconfigured to be non-positive (e.g., <= 0), it could lead to incorrect cost calculations (such as zero or negative cycles) in the cost model. It is safer to guard this value with std::max to ensure cpr is at least 1.
| const int64_t cpr = (cfg.bytes_a == 4) ? static_cast<int64_t>(cfg.mad_fp32_passes) | |
| : std::max<int64_t>(1, static_cast<int64_t>(cfg.bytes_a) / 2); | |
| const int64_t cpr = (cfg.bytes_a == 4) ? std::max<int64_t>(1, static_cast<int64_t>(cfg.mad_fp32_passes)) | |
| : std::max<int64_t>(1, static_cast<int64_t>(cfg.bytes_a) / 2); |
There was a problem hiding this comment.
Fixed in 237af01 — guarded with std::max<int64_t>(1, cfg.mad_fp32_passes) in MadCycles and the matching max(1, cfg.mad_fp32_passes) in the Python oracle, so a non-positive misconfig can't yield cpr < 1. No-op for valid configs (a2a3=2, a5=8).
| # Cube passes per K-fractal: arch-specific mad_fp32_passes for full fp32, else 1 (bf16). | ||
| cpr = cfg.mad_fp32_passes if cfg.bytes_a == 4 else max(1, cfg.bytes_a // 2) |
There was a problem hiding this comment.
To maintain perfect consistency with the C++ implementation's cost model calculations, guard cfg.mad_fp32_passes with max(1, ...) here as well.
| # Cube passes per K-fractal: arch-specific mad_fp32_passes for full fp32, else 1 (bf16). | |
| cpr = cfg.mad_fp32_passes if cfg.bytes_a == 4 else max(1, cfg.bytes_a // 2) | |
| # Cube passes per K-fractal: arch-specific mad_fp32_passes for full fp32, else 1 (bf16). | |
| cpr = max(1, cfg.mad_fp32_passes) if cfg.bytes_a == 4 else max(1, cfg.bytes_a // 2) |
There was a problem hiding this comment.
Fixed in 237af01 — the oracle mirror now uses max(1, cfg.mad_fp32_passes) too, kept in lock-step with the C++ MadCycles.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bbf833df94
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| print(f"\n=== {dn} {M}x{K}x{N} | {len(cands)} legal candidates | chooser pick ({r.m},{r.n},{r.k}) ===") | ||
| print(f" {'rank':>4} {'m,n,k':>14} {'stat':>4} {'model_wall':>11} (agent: measure a5-sim wall)") | ||
| for i, (w, m, n, k, st) in enumerate(cands[:TOPK]): | ||
| mark = " <-- CHOOSER PICK" if (m, n, k) == pick else "" |
There was a problem hiding this comment.
Include stationarity when marking chooser picks
When a full-K tile shape is legal under multiple regimes, the same (m, n, k) can appear once as OS and again as A/B stationary, but this comparison ignores r.stationarity. In those cases the validation harness can mark the wrong loop order as the chooser pick (or mark multiple rows), causing the device-agent measurements to validate a candidate that the emitted matmul would not actually use; compare the stationarity token as well when setting mark.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Resolved by removing the file. device_scripts/a5_validate_picks.py (and a5_cost_model_calib.py) were local a5-sim dev/validation harnesses — infrastructure, not shippable code — so they've been dropped from this PR in 237af01. The a5-sim validation numbers stand (produced by those local harnesses); only the shipped cost-model constants + mad_fp32_passes term remain in the PR.
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
include/pypto/backend/common/backend_handler.h (1)
94-99: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNew
mad_fp32_passesfield: correct and well documented, but leaves stale doc above.The field addition itself is sound and matches the a5 override (
mad_fp32_passes = 8) and default (2) used elsewhere. However, the class-level docstring a few lines above (68-71) still states ``cpr(1 for 2-byte inputs, 2 for 4-byte) is derived from the operand byte width'' — that is no longer accurate now that `cpr` for fp32 is arch-specific via `mad_fp32_passes` (e.g. 8 on a5, not `bytes_a/2`). Worth updating that earlier paragraph to reference `mad_fp32_passes` so the doc doesn't contradict the new field it directly governs.📝 Suggested doc fix
* The MAD term mirrors the cube's per-TMATMUL cost * `mad_head_cycles + cpr * ceil(m/16) * ceil(k/kt) * ceil(n/16)`, where - * `kt = mad_k_fractal_bytes / bytes_a` and `cpr` (1 for 2-byte inputs, 2 for - * 4-byte) is derived from the operand byte width in the chooser. + * `kt = mad_k_fractal_bytes / bytes_a` and `cpr` is 1 for <=2-byte inputs, or + * `mad_fp32_passes` (arch-specific) for full 4-byte fp32 inputs. */🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/pypto/backend/common/backend_handler.h` around lines 94 - 99, Update the class-level documentation near the cpr description to remove the claim that cpr is derived solely from operand byte width; explain that fp32 cpr is controlled by the architecture-specific mad_fp32_passes field, while other input types retain the existing byte-width behavior.include/pypto/ir/transforms/utils/l0_tile_chooser.h (1)
247-249: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStale
cpr = bytes_a/2reference in theChooseL0Tiledoc.This line in the algorithm docstring still states
cpr = bytes_a/2, which no longer holds for fp32 (bytes_a == 4) now thatcprismad_fp32_passesthere (arch-specific, e.g. 8 on a5). Same staleness as theL0CostModeldoc inbackend_handler.h.📝 Suggested doc fix
- with kt = mad_k_fractal_bytes/bytes_a, cpr = bytes_a/2. Ties break by lex + with kt = mad_k_fractal_bytes/bytes_a, cpr = 1 for <=2-byte inputs or + mad_fp32_passes (arch-specific) for fp32. Ties break by lex🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/pypto/ir/transforms/utils/l0_tile_chooser.h` around lines 247 - 249, Update the ChooseL0Tile algorithm documentation to remove the stale “cpr = bytes_a/2” claim. Describe cpr as the architecture- and datatype-dependent MAD pass count, including that fp32 uses mad_fp32_passes, and keep the surrounding tie-breaker description accurate.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@device_scripts/a5_validate_picks.py`:
- Around line 95-104: Update the chooser-pick comparison in the
candidate-printing loop to include normalized stationarity and the OS hold
direction, not just (m, n, k). Build the chooser identity from all corresponding
fields of r and compare it against each candidate’s full identity so duplicate
tile shapes with different OS/A/B variants are marked correctly.
- Around line 65-85: Preserve the complete chooser ordering in candidates by
storing the full result of O._wall_key(...) rather than only index [0]. Sort
candidates using the full key, while retaining key[0] as the reported model wall
value; update downstream tuple unpacking or output logic accordingly.
---
Nitpick comments:
In `@include/pypto/backend/common/backend_handler.h`:
- Around line 94-99: Update the class-level documentation near the cpr
description to remove the claim that cpr is derived solely from operand byte
width; explain that fp32 cpr is controlled by the architecture-specific
mad_fp32_passes field, while other input types retain the existing byte-width
behavior.
In `@include/pypto/ir/transforms/utils/l0_tile_chooser.h`:
- Around line 247-249: Update the ChooseL0Tile algorithm documentation to remove
the stale “cpr = bytes_a/2” claim. Describe cpr as the architecture- and
datatype-dependent MAD pass count, including that fp32 uses mad_fp32_passes, and
keep the surrounding tie-breaker description accurate.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: e22cdd45-0779-4429-b0ff-6d9dd2761c2f
📒 Files selected for processing (10)
device_scripts/a5_cost_model_calib.pydevice_scripts/a5_validate_picks.pyinclude/pypto/backend/950/backend_950_handler.hinclude/pypto/backend/common/backend_handler.hinclude/pypto/ir/transforms/utils/l0_tile_chooser.hpython/bindings/modules/passes.cpppython/pypto/pypto_core/passes.pyisrc/ir/transforms/auto_tile_matmul_l0_pass.cppsrc/ir/transforms/utils/l0_tile_chooser.cpptests/ut/ir/transforms/test_l0_tile_chooser.py
…_passes
- Remove device_scripts/{a5_cost_model_calib,a5_validate_picks}.py — these are
local a5-sim dev/validation harnesses (infrastructure), not shippable code.
They caused the pre-commit pyright failure (runtime-only compiled-submodule +
sys.path-injected test imports) and drew review nits on the harness itself.
Drop the file reference from the Backend950 comment too.
- Guard `mad_fp32_passes` with max(1, ...) in MadCycles (l0_tile_chooser.cpp) and
its Python oracle mirror, so a non-positive misconfig can't yield cpr < 1
(per review). No-op for valid configs (a2a3=2, a5=8).
|
CI note: the single
Could a maintainer re-run the |
Summary
Closes the a5-calibration follow-up from #1912 —
Backend950inherited thea2a3-calibrated
L0CostModel, so a5 tiled with a2a3's roofline constants (only its 256 KBL0C differed). This calibrates a5's cost model from an a5-sim sweep (there is no a5
device, so the a5-sim is the ground truth), and adds one arch-general model term the
calibration required.
Two parts:
mad_fp32_passes— a per-arch fp32 cube-pass term (shared cost model). The cube MADcost derived
cpr = bytes_a/2(fp32 → 2 passes/K-fractal), baking in a2a3's fp32decomposition. a5-sim shows a5 prices full fp32 MMAD ~4× a2a3 per fractal (a fuller
mantissa decomposition; the fractal width
K0=C0/Eis arch-invariant per pto-isa, sothe pass count is the only arch-varying fp32 cube term). New
L0CostModel::mad_fp32_passes(default 2 = a2a3, byte-identical; a5 = 8). a2a3 is unchanged (a5-only override).
The 7 a5
L0CostModelconstants inBackend950::GetL0CostModel(), fit from a5-sim.Calibrated a5 constants (raw a5-sim work-cycle fit)
mad_head_cyclesmad_fp32_passesmmad = 25 + 8·fractals(R²=1); bf16 cpr stays 1bw_l0aLOAD_2Dv2, analytic-bytes fitbw_l0bbw_drainFIX_L0C_TO_DST(a5 drain ~4× costlier)drain_fixed_cyclesdrain_row_cyclesdrain_penalty_cyclesoddPartform re-confirmeddrain_c0_bytes,mad_k_fractal_bytes= 32 (ISA, arch-invariant). Fits the a5-sim sweepsto <1%. L0 capacities (L0C 256 KB, L0A/L0B 64 KB, Mat 512 KB) were already correct
(grounded in pto-isa
buffer_limits.hpp) and are untouched.Method
Work-cycle extraction by instruction source (
MMAD/FIX_L0C_TO_DST/LOAD_2Dv2),excluding
WAIT_FLAG/BARstalls. This reproduces the shipped a2a3 device model within~15% on the load/cube lanes, so a5 values are fit raw (no transfer correction). a5-sim
cannot execute a multi-tile tile loop (it collapses the K/tile loops), so full-matmul walls
were reconstructed from single-tile lanes × tile-count × the
WallCyclesschedule.Validation (a5-sim)
a5-sim wall: the chooser's tile is the sim's feat: add CHECK and INTERNAL_CHECK macros for runtime validation #2–refactor(ir): add null checks to visitor and mutator base classes #3, within ~1–2% of the measured best
on all tested shapes (model refactor(ir): Extract Expr base class from ScalarExpr for extensibility #1 = sim refactor(ir): Extract Expr base class from ScalarExpr for extensibility #1 for divisor-k). The ~34% absolute over-prediction
from
bw_drain(below) preserves the ranking, so it does not mis-pick.forces tiny fp32 tiles → more drains → ~40% fp32 loss (bf16 only ~1.25×). a5's expensive
drain is not fixed by a schedule-default flip.
hit pto-isa targets; low cube utilization is a5's drain-bound reality (FIXP 44–52% of the
wall), not an emit defect.
Known caveat (documented at the constant)
bw_drain=30carries a ~4× a5-sim-vs-device gap (a2a3-sim drain is likewise ~4× a2a3-device's118). It over-predicts absolute a5 walls ~34% but is ranking-safe (validation above). The
real absolute-accuracy lever is a5-device drain data (none exists) or a tail-accurate
drain-hidden pricing — not a
bw_drainhand-edit.Scope / testing
mad_fp32_passesdefaults to 2 (= oldcpr); a5 override is a5-only.Chooser + backend UTs pass (a2a3 identical). a5 default-planner pick-diff vs the
a2a3-inherited baseline: 59/180 (all the shifts explained above).
sweep); those scripts are dev infrastructure and are not included in this PR.
Follow-ups (tracked, not blocking)
bw_drainabsolute gap.KNOWN_ISSUES).Builds on the
L0CostModelfields added by #1959 (now merged); rebased onto currentmain.a2a3-only, independent of the in-flight PyPTO-planner dbC=2 work.