feat(lower-tirx): support PTX movmatrix

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.

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

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",

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():