Skip to content

[Feature] Support transpose in MX L1-to-L0 staging #1236

Description

@jimmychou0

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions