Use shapes for MX physical TLoad box; re-enable host-prequant matmul_mx ST - #2620
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (4)
💤 Files with no reviewable changes (1)
Included review availability: Your plan provides up to 2 included reviews per hour; 1 remains after this review. 📝 WalkthroughWalkthroughMX-layout tile loads now accept dynamic partition offsets. Static offsets remain validated for minimum and alignment requirements. Backend codegen emits runtime division for dynamic physical offsets. Tests cover generated MLIR and previously skipped MX matmul tests. ChangesDynamic MX offsets
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: 🟡 Moderate · up to This change allows runtime-provided MX scale offsets to select memory partitions, but invalid or misaligned values could read an unintended or out-of-range scale region. The PR is not merge-ready until that runtime invariant is enforced or explicitly accepted by the owner. Sequence Diagram(s)sequenceDiagram
participant TileLoad
participant DeduceTileLoadType
participant EmitMxPhysicalView
participant MLIR
TileLoad->>DeduceTileLoadType: provide dynamic block and group offsets
DeduceTileLoadType-->>TileLoad: accept offsets
TileLoad->>EmitMxPhysicalView: generate MX physical view
EmitMxPhysicalView->>MLIR: emit arith.maxsi and arith.divui
EmitMxPhysicalView->>MLIR: emit mx5d_view and pto.tload
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Title checkExplanation The title identifies two real changes: MX physical TLoad shape handling and re-enabling host-prequantized matmul_mx system tests. It omits dynamic MX offset support but remains related and sufficiently specific. Full details: Description checkExplanation The description covers MX physical views, MX offsets, test re-enablement, and testing status. It conflicts with the changeset by stating that dynamic-offset support is excluded, although the code and objectives include dynamic offsets.
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 |
![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: f7c9e31908
ℹ️ 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".
f7c9e31 to
015a412
Compare
YunjiQin
left a comment
YunjiQin
left a comment
There was a problem hiding this comment.
Review — dynamic MX scale load offsets
Reviewed at 015a412. Two independent changes are bundled here. The valid_shape → shapes half is correct and I would take it as-is; the dynamic-offset half cannot work on PTOAS as written, and the mechanism it needs already exists in the tree.
Blocking — cf.assert is not emittable by PTOAS, at any version
EmitMxPhysicalOffsetCode emits the first cf dialect op in the codebase — a scan of src/codegen + src/backend shows we only ever emit pto / arith / scf / func. PTOAS parses it (so the dialect is registered) but the C++ emitter has no printer for it:
$ ptoas --pto-arch=a5 --pto-level=level3 -o /dev/null mx_assert.mlir
loc("mx_assert.mlir":17:3): error: 'cf.assert' op unable to find printer for op
Error: Failed to emit C++.
$ echo $?
1The same file with only the cf.assert line deleted compiles clean (exit 0). I swept every installed version — 0.54, 0.55, 0.57, 0.58, 0.59, 0.60 all fail identically. This is not the v0.57 pin: PTOAS's EmitC backend has never had a printer for cf.assert. Reproduce with /usr/local/ptoas/<ver>/bin/ptoas; I can share the two .mlir files.
Upstream cf.assert is a runtime assertion (ConvertControlFlowToLLVM lowers it to cond_br + abort), but PTOAS never takes that path, so the intended semantics can't be reached.
This collapses the PR's premise. The summary says codegen emits remui + cf.assert + divui "so unaligned runtime offsets cannot silently floor to the wrong scale block". Remove the unemittable op and remui/cmpi have no consumer; what remains is maxsi + divui — precisely the silent-floor behaviour the PR set out to prevent. There is no degraded-but-useful subset here.
One question on the checklist: Local A5 … matmul_mx path PASS (aligned dynamic offsets). Since IR containing the assert cannot compile, were those offsets constant-folded to ConstInt by Simplify before codegen — i.e. did that run take the ci->value_ / divisor branch and never exercise the new path?
The mechanism we need is already in the tree
include/pypto/ir/transforms/utils/tensor_view_semantics.h solves this exact problem at compile time, for NZ:
IsProvableMultipleOf(expr, divisor, facts, budget)— proves a symbolic offset is a multiple of the fractal factor (constants, loop vars whose start and step are both multiples, and sums/differences/constant multiples of those), walking the SSA def chain under a visit budget so the pass stays O(N).IsProvableNonNegative(...)— its doc comment describes our failure mode verbatim: "FloorDiv(n0, 16)is-1; codegen then clamps a negativepto.partition_viewoffset to 0 rather than failing, so the load silently reads fractal 0 instead of reporting anything. That is the same silent-wrong-data shape #2543 fixed for row indices."BlockNzOffsets— maps logical[…, r0, c0off]into blocked[…, c0off/c0, r0/16, 0, 0], folding constants with a user-facingCHECK_SPANand dividing the result (MakeFloorDiv) rather than re-associating the arithmetic, with the wraparound-soundness argument written out.
Structurally identical to MX_A_ZZ [M, G] → [1, M/16, G/2, 16, 2]. All compile-time, zero runtime cost, no PTOAS feature required, and failures are user-level errors that name the provable forms.
The consumer is BlockNzTensorViews (pass 14), currently gated on TensorLayout::NZ. Its header states why the blocking lives in a pass rather than the backend:
After this pass no logical-shaped NZ TensorType survives, so nothing downstream — including codegen, which derives every rank from
TensorType::shape_— needs to know NZ is special. That is the reason the blocking lives here rather than in the backend:EmitMakeTensorViews,GetTensorViewTypeStringand thetile.loadpartition_viewemitter each read the rank independently and must agree.
That is the argument against EmitMxPhysicalView, which hand-builds the rank-5 view and consequently has to hand-build the partition type and bypass the generic path ("skip the generic parameter view — see EmitMakeTensorViews").
For contrast, here is an NZ load with a symbolic loop-var offset today:
%x_view = pto.make_tensor_view %arg0, shape = [%c4_index, %c8_index, %c16_index, %c16_index], ... {layout = #pto.layout<nz>}
scf.for %k__idx_v0 = %c0_index to %c64_index step %c16_index {
%0 = arith.divsi %k__idx_v0, %c16_index : index // pass 14's MakeFloorDiv, generic scalar codegen
%1 = arith.maxsi %0, %c0_index : index // existing GetIndexOffsetCodes UB guard
%x_pview = pto.partition_view %x_view, offsets = [%c0_index, %1, %c0_index, %c0_index], ...A generic view, two arith ops, no runtime check — alignment and non-negativity were proven in the pass. And MX fits pass 14's stated Milestone 1 scope point for point: read-only, target_memory=Mat (which MX already mandates), whole-byte dtype, static shapes, fractal-aligned shapes. Only the factors differ (16×2 vs 16×c0). Moving MX there is not a rewrite so much as putting it where NZ already lives.
Note the proof needs NzOffsetFacts (definition map, loop start/step, SPMD non-negativity), which only a pass can supply — a default-constructed one "proves nothing beyond constant folding". So DeduceTileLoadType is not a viable home; it would reject every dynamic offset.
Independently correct — worth splitting out
The valid_shape → shapes change stands on its own and does not depend on any of the above. It is also a robustness improvement: partition_dims feeds the static type string !pto.partition_tensor_view<1x1x1x16x2x…>, so sourcing it from shapes (required static) rather than valid_shape (may be dynamic) is strictly safer. Existing callers that omit valid_shape get shapes, so there is no behaviour change.
Suggest landing that as its own PR — it need not wait on the offset redesign.
test_matmul_mx.py un-skip
The skip reason shared by all three MX STs is "MX pack/layout requires PTOAS v0.60". test_matmul_mx.py builds MX_A_ZZ itself via _pack_a_scale, so "host-prequant, no quant_mx" does not address that blocker — quant_mx is one way to produce the pack, not the pack itself. Is the CI pin back at v0.60? If it is still v0.57, mx-tests-a5 will fail. The checklist item is unticked, so this looks unverified either way.
Non-blocking
- DRY. The first half of
EmitMxPhysicalOffsetCode(ConstInt branch,EmitCastToIndex,arith.maxsiclamp) duplicatesGetIndexOffsetCodes(pto_ops_shared.cpp:117-139) line for line — and that helper is already in this file'susinglist. - Redundant check. The new
valid_shape <= shapesloop overlapsInferWindowReadValidShape(type_inference.cpp:665-670), which already enforces this for everytile.loadwith a span and a better message. Because the new loop runs first, it shadows that message for MX. It also usesCHECKwhereargs[3]->span_is in hand —error-checking.mdasks for the_SPANform. The only genuinely new predicate isvalid > 0, which applies to all loads and would be better placed in the generic path. - Error class.
check_aligned_if_staticonly catches literals.pl.load(a_s, [i * 8, 0], ...)folded to8bySimplifyreaches the codegenINTERNAL_CHECK_SPAN, so a user error surfaces as "Internal error". The proof-based approach removes this by construction. - Wording.
check_static_aligned(shapes_tuple->..., "load block size")now describesshapes, notvalid_shape; "load block shape" would avoid the ambiguity.
Test coverage
assert "%c16_index" in mlirintest_mx_scale_load_accepts_narrowed_valid_shapedoes not test what its comment claims.%c16_indexappears inalloc_tile, in themake_tensor_viewshape, and in thepartition_viewsizes — the assertion holds even if the partition box shrinks. Assert on thepto.partition_viewline'ssizes = [%c1_index, %c1_index, %c1_index, %c16_index, %c2_index]instead. Same for%c16_index/%c2_indexin the dynamic-offset test: both are shape constants, so only theremui/cf.assert/divuiassertions carry information.- The new
valid_shape <= shapescheck has zero coverage — it fires only when both sides areConstInt, and the test passes a dynamicvalid_rows. The PR describes static partial-M narrowing; avalid_shape=[8, 2]case would exercise both the description and the check, plus avalid_shape=[32, 2]negative. - No negative test that a static misaligned offset is still rejected (
[8, 0]/[0, 1]). That is the boundary thecheck_static_aligned→check_aligned_if_staticswap is most likely to break, and nothing intests/covers it today. - Neither new test reaches
matmul_mx, though the comment justifies narrowedvalid_shapeby whatmatmul_mxconsumes. "TLoad moves 16 rows / tile declaresvalid_row8" is the combination that needs on-device evidence. tests/ut/codegen/test_mx_ops_codegen.pyhas no trailingpytest.mainblock (pre-existing; worth fixing while here).
Suggested split
| Scope | State | |
|---|---|---|
| A | valid_shape → shapes + test fixes |
Landable once the non-blocking items above are addressed |
| B | Dynamic MX offsets via pass 14 + the proof helpers | Needs redesign |
| C | test_matmul_mx.py un-skip |
Pending the PTOAS pin answer |
That way A's benefit isn't held up by B, and B can land where it architecturally belongs.
9352692 to
17b4e32
Compare
…mx ST. Physical partition sizes now follow fractal-aligned load shapes so a narrowed valid_shape only affects tile metadata. Keep ConstInt MX offsets and EmitMxPhysicalView. Drop the module-level skip on test_matmul_mx (PTOAS v0.57 is enough for host-prequant); leave quant_mx suites skipped.
17b4e32 to
def1a59
Compare
yanghaoran29
changed the title
|
@YunjiQin Thanks again — following your Suggested split, this PR is now A + C only. Scope after force-push (
|
Extend pass 14 to block MX_A_ZZ/MX_B_NN into packed rank-5 and rewrite tile.load windows with IsProvableMultipleOf/NonNegative (including FloorDiv forms like ks=k0//32). Drop EmitMxPhysicalView so codegen uses the generic make_tensor_view path; unprovable offsets fail at compile time. Builds on hw-native-sys#2620 (shapes-sized physical TLoad box / narrowed valid_shape metadata).
Extend pass 14 to block MX_A_ZZ/MX_B_NN into packed rank-5 and rewrite tile.load windows with IsProvableMultipleOf/NonNegative (including FloorDiv forms like ks=k0//32). Use per-phase function stamps for idempotence, preserve ND/MX backing aliases and Submit return types, and validate canonical physical load geometry. Drop EmitMxPhysicalView so codegen uses the generic make_tensor_view path; unprovable offsets fail at compile time and narrowed valid shapes remain tile metadata. Builds on hw-native-sys#2620 (shapes-sized physical TLoad box / narrowed valid_shape metadata).
Extend pass 14 to block MX_A_ZZ/MX_B_NN into packed rank-5 and rewrite tile.load windows with IsProvableMultipleOf/NonNegative (including FloorDiv forms like ks=k0//32). Use per-phase function stamps for idempotence, preserve ND/MX backing aliases and Submit return types, and validate canonical physical load geometry. Drop EmitMxPhysicalView so codegen uses the generic make_tensor_view path; unprovable offsets fail at compile time and narrowed valid shapes remain tile metadata. Builds on hw-native-sys#2620 (shapes-sized physical TLoad box / narrowed valid_shape metadata).
Extend pass 14 to block MX_A_ZZ/MX_B_NN into packed rank-5 and rewrite tile.load windows with IsProvableMultipleOf/NonNegative (including FloorDiv forms like ks=k0//32). Follow alignment facts through outlined scalar parameters, require every call site to prove them, preserve ND/MX backing aliases and Submit return types, and validate canonical physical load geometry. Drop EmitMxPhysicalView so codegen uses the generic make_tensor_view path; unprovable offsets fail at compile time and narrowed valid shapes remain tile metadata. Builds on hw-native-sys#2620 (shapes-sized physical TLoad box / narrowed valid_shape metadata).
Extend pass 14 to block MX_A_ZZ/MX_B_NN into packed rank-5 and rewrite tile.load windows with IsProvableMultipleOf/NonNegative (including FloorDiv forms like ks=k0//32). Follow alignment facts through outlined scalar parameters, require every call site to prove them, preserve ND/MX backing aliases and Submit return types, and validate canonical physical load geometry. Drop EmitMxPhysicalView so codegen uses the generic make_tensor_view path; unprovable offsets fail at compile time and narrowed valid shapes remain tile metadata. Builds on hw-native-sys#2620 (shapes-sized physical TLoad box / narrowed valid_shape metadata).
2 participants
Summary
partition_view/ TLoad box sizes now follow fractal-aligned load shapes. A narrower valid_shape only updates tile metadata (e.g. partial M formatmul_mx) and does not shrink the hardware transfer.EmitMxPhysicalView(no dynamic-offset /cf.assertpath in this PR).tests/st/runtime/ops/test_matmul_mx.py(host-prequant; noquant_mx). Leavetest_quant_mx.py/test_quantized_matmul_mx.pyskipped under the PTOAS v0.57 pin.feat/mx-pass14-dynamic-offsetsfor a follow-up PR.Testing
onboard-tests-a5→mx-tests-a5SUCCESS under PTOAS v0.57tests/ut/codegen/test_mx_ops_codegen.pynarrowed-valid_shape casespre-commit/clang-tidy/unit-tests