Summary
Add transpose support to the public MX L1-to-L0 staging operations pto.mte_l1_l0a_mx and pto.mte_l1_l0b_mx.
Motivation / use case
TileLang FP8 blockscaled GEMM needs this for K-on-row layouts. In examples/ascend/test_gemm_padding_irregular_shapes.py, 4 of the 8 FP8 cases use k_on_row=True; their Final TIR requires transposed data and scale-factor staging from L1 to L0.
Regular pto.mte_l1_l0a/b already expose transpose, but the MX wrappers accept only MX coordinate/stride operands and have no transpose field. TileLang currently rejects the MX path when transpose != 0 to avoid emitting semantically incorrect PTODSL.
PR #1235 (pto.raw_fill_l1) resolves the separate OOB-padding L1 fill blocker. MX transpose remains the blocker for the four K-row FP8 cases.
Proposed API / behavior
- Extend the public PTODSL APIs
pto.mte_l1_l0a_mx(...) and pto.mte_l1_l0b_mx(...) with transpose=False.
- Carry the control through public VPTO wrapper ODS, parser/printer, verifier, and pointer normalization.
- Extend wrapper expansion and raw MX load representation so data and E8M0 scale staging use the same transpose semantics.
- Emit the ISA-supported transpose configuration/intrinsic in both VPTO LLVM and CANN900 lowering paths.
- Document the control, supported dtype/layout constraints, and its hardware encoding in the micro-ISA documentation.
- Preserve current behavior for
transpose=False. Reject transpose=True with a precise diagnostic if an ISA/backend combination cannot represent it.
Acceptance tests
- PTODSL frontend tests for both MX wrappers, shape-derived and full-control forms, with
transpose=True.
- VPTO lit coverage verifies the wrapper-to-raw expansion and emitted LLVM/CANN900 intrinsic/configuration for L0A and L0B.
- Negative tests cover unsupported type/layout combinations.
- TileLang source golden verifies
ascend_load_cbuf_to_ca/cb 16-argument MX calls carry the Final TIR transpose flag.
- A5 numerical validation covers all four
k_on_row=True FP8 cases in examples/ascend/test_gemm_padding_irregular_shapes.py.
Alternatives considered
Do not silently clear the transpose flag or reinterpret the scale pointer in TileLang: either option changes data/scale layout semantics and would make K-row FP8 results incorrect. Restricting migration to K-column cases leaves half of the FP8 irregular-shape matrix unsupported.
Additional context
Relevant local paths:
ptodsl/ptodsl/_ops.py (mte_l1_l0a_mx, mte_l1_l0b_mx)
include/PTO/IR/VPTOOps.td (PTO_MteL1L0aMxOp, PTO_MteL1L0bMxOp)
lib/PTO/Transforms/VPTOExpandWrapperOps.cpp
lib/PTO/Transforms/VPTOLLVMEmitter.cpp
lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp
Related: #1235
Summary
Add transpose support to the public MX L1-to-L0 staging operations
pto.mte_l1_l0a_mxandpto.mte_l1_l0b_mx.Motivation / use case
TileLang FP8 blockscaled GEMM needs this for
K-on-row layouts. Inexamples/ascend/test_gemm_padding_irregular_shapes.py, 4 of the 8 FP8 cases usek_on_row=True; their Final TIR requires transposed data and scale-factor staging from L1 to L0.Regular
pto.mte_l1_l0a/balready exposetranspose, but the MX wrappers accept only MX coordinate/stride operands and have no transpose field. TileLang currently rejects the MX path whentranspose != 0to avoid emitting semantically incorrect PTODSL.PR #1235 (
pto.raw_fill_l1) resolves the separate OOB-padding L1 fill blocker. MX transpose remains the blocker for the four K-row FP8 cases.Proposed API / behavior
pto.mte_l1_l0a_mx(...)andpto.mte_l1_l0b_mx(...)withtranspose=False.transpose=False. Rejecttranspose=Truewith a precise diagnostic if an ISA/backend combination cannot represent it.Acceptance tests
transpose=True.ascend_load_cbuf_to_ca/cb16-argument MX calls carry the Final TIR transpose flag.k_on_row=TrueFP8 cases inexamples/ascend/test_gemm_padding_irregular_shapes.py.Alternatives considered
Do not silently clear the transpose flag or reinterpret the scale pointer in TileLang: either option changes data/scale layout semantics and would make K-row FP8 results incorrect. Restricting migration to K-column cases leaves half of the FP8 irregular-shape matrix unsupported.
Additional context
Relevant local paths:
ptodsl/ptodsl/_ops.py(mte_l1_l0a_mx,mte_l1_l0b_mx)include/PTO/IR/VPTOOps.td(PTO_MteL1L0aMxOp,PTO_MteL1L0bMxOp)lib/PTO/Transforms/VPTOExpandWrapperOps.cpplib/PTO/Transforms/VPTOLLVMEmitter.cpplib/PTO/Transforms/VPTOCANN900LLVMEmitter.cppRelated: #1235