diff --git a/python/tvm/backend/cuda/ptx/table.py b/python/tvm/backend/cuda/ptx/table.py index 209651c59a4b..fb15a4ca1671 100644 --- a/python/tvm/backend/cuda/ptx/table.py +++ b/python/tvm/backend/cuda/ptx/table.py @@ -6991,6 +6991,25 @@ def _check_tcgen05_mma_block_scale_block(m): OperandSlot("p", kind="addr", allow_imm_offset=True), ), ), + # movmatrix per PTX ISA 9.7.15.5.14 -- transpose one distributed m8n8 + # matrix whose 16-bit elements are carried by one b32 register per lane. + # + # movmatrix.sync.aligned.m8n8.trans.b16 d, a; + InstructionEntry( + name="movmatrix", + slots=( + ModifierSlot("sync", ("sync",)), + ModifierSlot("aligned", ("aligned",)), + ModifierSlot("shape", ("m8n8",)), + ModifierSlot("trans", ("trans",)), + ModifierSlot("type", ("b16",)), + ), + cert_arch="sm_75", + operands=( + OperandSlot("d", rw="w", dtype="b32"), + OperandSlot("a", dtype="b32"), + ), + ), # stmatrix per PTX ISA 9.7.15.5.16 -- the store mirror of ldmatrix. One # syntax line; the two shapes split into two entries because their type and # target floors differ. diff --git a/python/tvm/script/tirx.pyi b/python/tvm/script/tirx.pyi index 89c3d9fd1136..111a77ff172b 100644 --- a/python/tvm/script/tirx.pyi +++ b/python/tvm/script/tirx.pyi @@ -1043,6 +1043,23 @@ class _Chain_mov: u64: _Chain_mov def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... +class _Chain_movmatrix: + """`movmatrix` — sync∈{sync}; aligned∈{aligned}; shape∈{m8n8}; trans∈{trans}; type∈{b16}""" + + aligned: _Chain_movmatrix + b16: _Chain_movmatrix + m8n8: _Chain_movmatrix + sync: _Chain_movmatrix + trans: _Chain_movmatrix + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... + class _Chain_mul: """`mul` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b) @@ -2242,6 +2259,7 @@ class _PTX: min: _Chain_min mma: _Chain_mma mov: _Chain_mov + movmatrix: _Chain_movmatrix mul: _Chain_mul mul24: _Chain_mul24 multimem_ld_reduce: _Chain_multimem_ld_reduce diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index e4e2fce5ba0b..93993e11b9e4 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/tests/python/tirx/codegen/test_codegen_cuda.py @@ -492,6 +492,7 @@ def main( T.ptx.st.shared.v4.b32(U32.data, U32[0], U32[1], U32[2], U32[3]) T.ptx.st_bulk.weak.shared__cta(U32.data, T.uint64(16)) T.ptx.fns.b32(U32[0], U32[0], U32[1], I32[0]) + T.ptx.movmatrix.sync.aligned.m8n8.trans.b16(U32[1], U32[0]) T.ptx.stmatrix.sync.aligned.m16n8.x1.trans.shared.b8(U32.data, U32[0]) F32[1] = T.cuda.uint_as_float(U32[0]) @@ -519,6 +520,7 @@ def main( "st.shared.v4.b32", "st.bulk.weak.shared::cta", "fns.b32", + "movmatrix.sync.aligned.m8n8.trans.b16", "stmatrix.sync.aligned.m16n8.x1.trans.shared.b8", "ld.global.f32", "tvm_builtin_cuda_ldg_f32", diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py index fe1ffaebe0db..f6902c233163 100644 --- a/tests/python/tirx/codegen/test_ptx_dialect.py +++ b/tests/python/tirx/codegen/test_ptx_dialect.py @@ -2329,7 +2329,7 @@ def test_ptx_all_variants_render_unique(): _, helper, _ = render_variant(entry, *args, addr_offsets=addr_offsets) assert helper not in names, f"address-offset helper name collision: {helper}" names.add(helper) - assert total == 200018 # update when the table grows + assert total == 200027 # update when the table grows def test_ptx_no_instruction_registered_twice():