[Mlir-commits] [mlir] [MLIR] Fix the PTX generation bug for StMatrixOp (PR #148250)

[llvmlistbot at llvm.org]

https://github.com/Pecco-314 created https://github.com/llvm/llvm-project/pull/148250

According to the [PTX documents](https://docs.nvidia.com/cuda/parallel-thread-execution/), the syntax of stmatrix should be:
```
stmatrix.sync.aligned.shape.num{.trans}{.ss}.type [p], r;

.shape  = {.m8n8, .m16n8};
.num    = {.x1, .x2, .x4};
.ss     = {.shared{::cta}};
.type   = {.b16, .b8};
```
However, the current code will generate the PTX like "stmatrix.sync.aligned.x4.m8n8.shared.b16". It seems like a bug.

>From dacee0847b660af16f5c8f45bcb83708d6916b0c Mon Sep 17 00:00:00 2001
From: Gao Yanfeng <gaoyanfeng at linux.alibaba.com>
Date: Fri, 11 Jul 2025 22:51:45 +0800
Subject: [PATCH] [MLIR] Fix the PTX generation bug for StMatrixOp

---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 6895e946b8a45..b27c03ec2c63f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2000,13 +2000,13 @@ def NVVM_StMatrixOp: NVVM_PTXBuilder_Op<"stmatrix">,
   let extraClassDefinition = [{
     std::string $cppClass::getPtx() {
       int d = getSources().size();
-      std::string ptx = "stmatrix.sync.aligned";
+      std::string ptx = "stmatrix.sync.aligned.m8n8";
       ptx += ".x" + std::to_string(d);
       if (getLayout() == NVVM::MMALayout::col)
         ptx += ".trans";
-      if(d == 1) ptx += ".m8n8.shared.b16 [%0], {%1};";
-      if(d == 2) ptx += ".m8n8.shared.b16 [%0], {%1, %2};";
-      if(d == 4) ptx += ".m8n8.shared.b16 [%0], {%1, %2, %3, %4};";
+      if(d == 1) ptx += ".shared.b16 [%0], {%1};";
+      if(d == 2) ptx += ".shared.b16 [%0], {%1, %2};";
+      if(d == 4) ptx += ".shared.b16 [%0], {%1, %2, %3, %4};";
       return ptx;
     }
   }];