[IR][BACKEND] Introduce nvgpu.ldmatrix IR (#5442)

Jokeren · web-flow · commit 43f1ad488d88 · 2024-12-19T18:00:19.000Z
The purpose is to replace the legacy way of using `ldmatrix` through
PTXBuilder.
diff --git a/test/Conversion/nvgpu_to_llvm.mlir b/test/Conversion/nvgpu_to_llvm.mlir
@@ -37,15 +37,24 @@ llvm.func @cluster_id() -> i32 {
 
 // -----
 
-// CHECK-LABEL: @st_matrix
-llvm.func @st_matrix(%i: i32, %ptr: !llvm.ptr<3>) {
+// CHECK-LABEL: @stmatrix
+llvm.func @stmatrix(%i: i32, %ptr: !llvm.ptr<3>) {
   // CHECK: stmatrix.sync.aligned.m8n8.x4.shared.b16 [$0], {$1, $2, $3, $4};
   nvgpu.stmatrix %ptr, %i, %i, %i, %i : !llvm.ptr<3>, i32, i32, i32, i32
   llvm.return
 }
 
 // -----
 
+// CHECK-LABEL: @ldmatrix
+llvm.func @ldmatrix(%ptr: !llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)> {
+  // CHECK: ldmatrix.sync.aligned.m8n8.x4.shared.b16 {$0, $1, $2, $3}, [$4];
+  %0 = nvgpu.ldmatrix %ptr : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
+  llvm.return %0 : !llvm.struct<(i32, i32, i32, i32)>
+}
+
+// -----
+
 !struct_128xf32 = !llvm.struct<(
   f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32,
   f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32, f32,
diff --git a/test/Conversion/tritongpu_to_llvm.mlir b/test/Conversion/tritongpu_to_llvm.mlir
@@ -845,10 +845,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
   tt.func @convert_dot(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) {
     %AA = ttg.local_alloc %A : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
     %BB = ttg.local_alloc %B : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
-    // CHECK: llvm.inline_asm
-    // CHECK: ldmatrix.sync.aligned.m8n8.x4
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4
+    // CHECK: nvgpu.ldmatrix
+    // CHECK: nvgpu.ldmatrix
     %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
@@ -876,8 +874,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
   tt.func @convert_dot_fp8(%A: tensor<16x16xf8E5M2, #blocked0>, %B: tensor<16x16xf8E5M2, #blocked0>) {
     %AA = ttg.local_alloc %A : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
     %BB = ttg.local_alloc %B : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4
+    // CHECK: nvgpu.ldmatrix
     %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf8E5M2, #shared0, #smem> -> tensor<16x16xf8E5M2, #dot_operand_a>
     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf8E5M2, #shared0, #smem> -> tensor<16x16xf8E5M2, #dot_operand_b>
     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
@@ -1177,7 +1174,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   tt.func @matmul_kernel_dot_operand_layout(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},
   %a:!ttg.memdesc<128x32xf16, #shared, #smem>, %b:!ttg.memdesc<32x256xf16, #shared, #smem>) {
     %cst = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #mma>
-    // CHECK: ldmatrix.sync.aligned.m8n8.x4.shared.b16
+    // CHECK: nvgpu.ldmatrix
     %a_mat = ttg.local_load %a : !ttg.memdesc<128x32xf16, #shared, #smem> -> tensor<128x32xf16, #dot_operand_a>
     %b_mat = ttg.local_load %b : !ttg.memdesc<32x256xf16, #shared, #smem> -> tensor<32x256xf16, #dot_operand_b>
 
@@ -1227,11 +1224,9 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
   tt.func @matmul_tf32dot(%ptr:!tt.ptr<f32> {tt.divisibility = 16 : i32},
   %a:!ttg.memdesc<32x16xf32, #shared, #smem>, %b:!ttg.memdesc<16x32xf32, #shared, #smem>) {
     %cst = arith.constant dense<0.000000e+00> : tensor<32x32xf32, #mma>
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4.shared.b16
+    // CHECK: nvgpu.ldmatrix
     // CHECK-SAME: (i32, i32, i32, i32)
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4.shared.b16
+    // CHECK: nvgpu.ldmatrix
     // CHECK-SAME: (i32, i32, i32, i32)
     %a_mat = ttg.local_load %a : !ttg.memdesc<32x16xf32, #shared, #smem> -> tensor<32x16xf32, #dot_operand_a>
     %b_mat = ttg.local_load %b : !ttg.memdesc<16x32xf32, #shared, #smem> -> tensor<16x32xf32, #dot_operand_b>
@@ -1720,10 +1715,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.thr
     %f16_shared = ttg.local_alloc %f16_inp : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
     %i16_shared = ttg.local_alloc %i16_inp : (tensor<16x16xi16, #blocked0>) -> !ttg.memdesc<16x16xi16, #shared0, #smem>
 
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4
-    // CHECK: llvm.inline_asm
-    // CHECK-SAME: ldmatrix.sync.aligned.m8n8.x4
+    // CHECK: nvgpu.ldmatrix
+    // CHECK: nvgpu.ldmatrix
 
     %f16_dot = ttg.local_load %f16_shared : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
     %i16_dot = ttg.local_load %i16_shared : !ttg.memdesc<16x16xi16, #shared0, #smem> -> tensor<16x16xi16, #dot_operand_b>
diff --git a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td
@@ -43,12 +43,10 @@ def NVGPU_WGMMACommitGroupOp : NVGPU_Op<"wgmma_commit_group", []> {
   let assemblyFormat = "attr-dict";
 }
 
-def NVGPU_WGMMAWaitGroupOp : NVGPU_Op<"wgmma_wait_group",
-[DeclareOpInterfaceMethods<InferTypeOpInterface>,
- AllTypesMatch<["input", "output"]>]> {
+def NVGPU_WGMMAWaitGroupOp : NVGPU_Op<"wgmma_wait_group", [DeclareOpInterfaceMethods<InferTypeOpInterface>,
+                                                           AllTypesMatch<["input", "output"]>]> {
   let arguments = (ins LLVM_AnyStruct:$input, I32Attr:$pendings);
   let results = (outs LLVM_AnyStruct:$output);
-  let assemblyFormat = "attr-dict";
   let assemblyFormat = "$input attr-dict `:` type($input)";
 }
 
@@ -103,10 +101,16 @@ def NVGPU_ClusterWaitOp : NVGPU_Op<"cluster_wait", []> {
 }
 
 def NVGPU_StoreMatrixOp : NVGPU_Op<"stmatrix", [MemoryEffects<[MemWrite]>]> {
-  let arguments = (ins LLVM_PointerShared:$addr, Variadic<I32>:$datas);
+  let arguments = (ins LLVM_PointerShared:$addr, Variadic<I32>:$vals);
   let assemblyFormat = "operands attr-dict `:` type(operands)";
 }
 
+def NVGPU_LoadMatrixOp : NVGPU_Op<"ldmatrix", [MemoryEffects<[MemRead]>]> {
+  let arguments = (ins LLVM_PointerShared:$addr);
+  let results = (outs LLVM_AnyStruct:$result);
+  let assemblyFormat = "$addr attr-dict `:` functional-type($addr, $result)";
+}
+
 def NVGPU_ClusterCTAIdOp : NVGPU_Op<"cluster_id", [Pure]> {
   let results = (outs I32:$result);
   let assemblyFormat = "attr-dict";
diff --git a/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp b/third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp
@@ -23,21 +23,20 @@ using ttn::OperandsAndConstraints;
 
 namespace {
 
-const std::string Wgmma_Fence_Op = "wgmma.fence.sync.aligned;";
-const std::string Wgmma_Commit_Group_Op = "wgmma.commit_group.sync.aligned;";
-const std::string Cluster_Wait_Op = "barrier.cluster.wait.aligned;";
-const std::string Fence_Mbarrier_Init_Op =
-    "fence.mbarrier_init.release.cluster;";
-const std::string Cluster_Cta_Id_Op = "{\n"
-                                      ".reg .u32 a<5>;              \n"
-                                      "mov.u32 a0, %cluster_ctaid.x;\n"  // x
-                                      "mov.u32 a1, %cluster_ctaid.y;\n"  // y
-                                      "mov.u32 a2, %cluster_ctaid.z;\n"  // z
-                                      "mov.u32 a3, %cluster_nctaid.x;\n" // nx
-                                      "mov.u32 a4, %cluster_nctaid.y;\n" // ny
-                                      "mad.lo.u32 a1, a2, a4, a1;     \n"
-                                      "mad.lo.u32 $0, a1, a3, a0;     \n"
-                                      "}";
+const std::string kWgmmaFenceOp = "wgmma.fence.sync.aligned;";
+const std::string kWgmmaCommitGroupOp = "wgmma.commit_group.sync.aligned;";
+const std::string kClusterWaitOp = "barrier.cluster.wait.aligned;";
+const std::string kFenceMbarrierInitOp = "fence.mbarrier_init.release.cluster;";
+const std::string kClusterCtaIdOp = "{\n"
+                                    ".reg .u32 a<5>;              \n"
+                                    "mov.u32 a0, %cluster_ctaid.x;\n"  // x
+                                    "mov.u32 a1, %cluster_ctaid.y;\n"  // y
+                                    "mov.u32 a2, %cluster_ctaid.z;\n"  // z
+                                    "mov.u32 a3, %cluster_nctaid.x;\n" // nx
+                                    "mov.u32 a4, %cluster_nctaid.y;\n" // ny
+                                    "mad.lo.u32 a1, a2, a4, a1;     \n"
+                                    "mad.lo.u32 $0, a1, a3, a0;     \n"
+                                    "}";
 
 bool isNumber(const std::string &s) {
   return !s.empty() && std::find_if(s.begin(), s.end(), [](unsigned char c) {
@@ -235,46 +234,141 @@ class ClusterArriveOpPattern : public OpRewritePattern<ttn::ClusterArriveOp> {
   }
 };
 
-class StoreMatrixOpPattern : public OpRewritePattern<ttn::StoreMatrixOp> {
+// Base class for Matrix Operation Patterns
+template <typename MatrixOpType, typename ConcreteMatrixOpPattern>
+class MatrixOpPattern : public OpRewritePattern<MatrixOpType> {
 public:
-  using OpRewritePattern<ttn::StoreMatrixOp>::OpRewritePattern;
+  using OpRewritePattern<MatrixOpType>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(ttn::StoreMatrixOp op,
+  LogicalResult matchAndRewrite(MatrixOpType op,
                                 PatternRewriter &rewriter) const override {
-    return rewriteAsPtxAsm(op, rewriter, getPtxAsm(op),
-                           getOperandsAndConstraints(op));
-  }
-
-  OperandsAndConstraints
-  getOperandsAndConstraints(ttn::StoreMatrixOp op) const {
-    OperandsAndConstraints operandsAndTypes;
-    auto addr = op.getAddr();
-    auto datas = op.getDatas();
-    operandsAndTypes.push_back({addr, "r"});
-    for (unsigned i = 0; i < datas.size(); i++) {
-      operandsAndTypes.push_back({datas[i], "r"});
-    }
-    return operandsAndTypes;
+    unsigned vecSize = getVectorSize(op);
+    bool trans = op->hasAttr("trans")
+                     ? op->template getAttrOfType<BoolAttr>("trans").getValue()
+                     : false;
+
+    // Template method for PTX assembly generation
+    std::string ptxAsm =
+        (llvm::Twine(ConcreteMatrixOpPattern::kOpCode) +
+         getPtxModifiers(vecSize, trans) + " " + getOperands(op, vecSize) + ";")
+            .str();
+
+    OperandsAndConstraints operandAndConstraints =
+        getOperandsAndConstraints(op, vecSize);
+    Constraints outputConstraints = getOutputConstraints(op, vecSize);
+
+    return rewriteAsPtxAsm(op, rewriter, ptxAsm, operandAndConstraints,
+                           outputConstraints);
   }
 
-  std::string getPtxAsm(ttn::StoreMatrixOp op) const {
-    auto datas = op.getDatas();
-    std::string ptxAsm;
-    switch (datas.size()) {
+protected:
+  // Shared helper methods
+  std::string getPtxModifiers(unsigned vecSize, bool trans) const {
+    auto ptxAsmBase = llvm::Twine(".sync.aligned.m8n8");
+    const std::string suffix = trans ? ".trans.shared.b16" : ".shared.b16";
+    switch (vecSize) {
     case 1:
-      ptxAsm = "stmatrix.sync.aligned.m8n8.x1.shared.b16 [$0], {$1};";
-      break;
+      return (ptxAsmBase + ".x1" + suffix).str();
     case 2:
-      ptxAsm = "stmatrix.sync.aligned.m8n8.x2.shared.b16 [$0], {$1, $2};";
-      break;
+      return (ptxAsmBase + ".x2" + suffix).str();
     case 4:
-      ptxAsm =
-          "stmatrix.sync.aligned.m8n8.x4.shared.b16 [$0], {$1, $2, $3, $4};";
-      break;
+      return (ptxAsmBase + ".x4" + suffix).str();
     default:
-      assert(false && "Invalid size");
+      assert(false && "Invalid vector size");
     }
-    return ptxAsm;
+  }
+
+  std::string getPtxRegOperands(unsigned startIdx, unsigned count) const {
+    llvm::SmallString<20> regOperands;
+    llvm::raw_svector_ostream stream(regOperands);
+    stream << "{";
+    for (unsigned i = 0; i < count; i++) {
+      stream << "$" + llvm::utostr(startIdx + i);
+      if (i != count - 1)
+        stream << ", ";
+    }
+    stream << "}";
+    return std::string(regOperands.str());
+  }
+
+  std::string getPtxAddrOperand(unsigned idx) const {
+    return (llvm::Twine("[$") + llvm::utostr(idx) + "]").str();
+  }
+
+  virtual std::string getOperands(MatrixOpType op, unsigned vecSize) const = 0;
+  virtual OperandsAndConstraints
+  getOperandsAndConstraints(MatrixOpType op, unsigned vecSize) const = 0;
+  virtual Constraints getOutputConstraints(MatrixOpType op,
+                                           unsigned vecSize) const = 0;
+  virtual unsigned getVectorSize(MatrixOpType op) const = 0;
+};
+
+// StoreMatrixOp Pattern
+class StoreMatrixOpPattern
+    : public MatrixOpPattern<ttn::StoreMatrixOp, StoreMatrixOpPattern> {
+public:
+  using MatrixOpPattern<ttn::StoreMatrixOp,
+                        StoreMatrixOpPattern>::MatrixOpPattern;
+  static constexpr const char *kOpCode = "stmatrix";
+
+protected:
+  unsigned getVectorSize(ttn::StoreMatrixOp op) const override {
+    return op.getVals().size();
+  }
+
+  std::string getOperands(ttn::StoreMatrixOp op,
+                          unsigned vecSize) const override {
+    return (llvm::Twine(getPtxAddrOperand(0)) + ", " +
+            getPtxRegOperands(1, vecSize))
+        .str();
+  }
+
+  OperandsAndConstraints
+  getOperandsAndConstraints(ttn::StoreMatrixOp op,
+                            unsigned vecSize) const override {
+    OperandsAndConstraints constraints = {{op.getAddr(), "r"}};
+    for (unsigned i = 0; i < vecSize; i++) {
+      constraints.push_back({op.getVals()[i], "r"});
+    }
+    return constraints;
+  }
+
+  Constraints getOutputConstraints(ttn::StoreMatrixOp op,
+                                   unsigned vecSize) const override {
+    return {}; // No output constraints for StoreMatrixOp
+  }
+};
+
+// LoadMatrixOp Pattern
+class LoadMatrixOpPattern
+    : public MatrixOpPattern<ttn::LoadMatrixOp, LoadMatrixOpPattern> {
+public:
+  using MatrixOpPattern<ttn::LoadMatrixOp,
+                        LoadMatrixOpPattern>::MatrixOpPattern;
+  static constexpr const char *kOpCode = "ldmatrix";
+
+protected:
+  unsigned getVectorSize(ttn::LoadMatrixOp op) const override {
+    auto resultType = cast<LLVM::LLVMStructType>(op.getType());
+    return resultType.getBody().size();
+  }
+
+  std::string getOperands(ttn::LoadMatrixOp op,
+                          unsigned vecSize) const override {
+    return (llvm::Twine(getPtxRegOperands(0, vecSize)) + ", " +
+            getPtxAddrOperand(vecSize))
+        .str();
+  }
+
+  OperandsAndConstraints
+  getOperandsAndConstraints(ttn::LoadMatrixOp op,
+                            unsigned vecSize) const override {
+    return {{op.getAddr(), "r"}};
+  }
+
+  Constraints getOutputConstraints(ttn::LoadMatrixOp op,
+                                   unsigned vecSize) const override {
+    return Constraints(vecSize, "=r");
   }
 };
 
@@ -507,17 +601,16 @@ class ConvertNVGPUToLLVM : public ConvertNVGPUToLLVMBase<ConvertNVGPUToLLVM> {
 #define POPULATE_NVGPU_OP(SRC_OP, ASM)                                         \
   patterns.add<NVGPUOpGenericPattern<SRC_OP>>(context, ASM, Constraints(),     \
                                               Constraints());
-    POPULATE_NVGPU_OP(ttn::WGMMAFenceOp, Wgmma_Fence_Op)
-    POPULATE_NVGPU_OP(ttn::WGMMACommitGroupOp, Wgmma_Commit_Group_Op)
-    POPULATE_NVGPU_OP(ttn::ClusterWaitOp, Cluster_Wait_Op)
+    POPULATE_NVGPU_OP(ttn::WGMMAFenceOp, kWgmmaFenceOp)
+    POPULATE_NVGPU_OP(ttn::WGMMACommitGroupOp, kWgmmaCommitGroupOp)
+    POPULATE_NVGPU_OP(ttn::ClusterWaitOp, kClusterWaitOp)
 #undef POPULATE_NVGPU_OP
     patterns.add<NVGPUOpGenericPattern<ttn::ClusterCTAIdOp>>(
-        context, Cluster_Cta_Id_Op, Constraints({"=r"}), Constraints());
+        context, kClusterCtaIdOp, Constraints({"=r"}), Constraints());
 
-    patterns
-        .add<FenceAsyncSharedOpPattern, StoreMatrixOpPattern,
-             ClusterArriveOpPattern, WGMMAOpPattern, WGMMAWaitGroupOpPattern>(
-            context);
+    patterns.add<FenceAsyncSharedOpPattern, LoadMatrixOpPattern,
+                 StoreMatrixOpPattern, ClusterArriveOpPattern, WGMMAOpPattern,
+                 WGMMAWaitGroupOpPattern>(context);
 
     if (applyPatternsAndFoldGreedily(mod, std::move(patterns)).failed())
       signalPassFailure();
diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToDotOperandMMAv2OrV3.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToDotOperandMMAv2OrV3.cpp
@@ -1,6 +1,7 @@
 #include "TritonNVIDIAGPUToLLVM/PTXAsmFormat.h"
 #include "Utility.h"
 #include "mlir/Support/LLVM.h"
+#include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
 
 using namespace mlir;
 
@@ -339,23 +340,10 @@ MMA16816SmemLoader::loadX4(int batch, int mat0, int mat1, ArrayRef<Value> ptrs,
     if (batch != 0)
       stridedOffset = add(
           stridedOffset, mul(i32_val(batch * warpsPerCTA[0]), smemBatchOffset));
-
     Value readPtr = gep(ptr_ty(ctx, 3), shemTy, ptr, stridedOffset);
-
-    PTXBuilder builder;
-    // ldmatrix.m8n8.x4 returns 4x2xfp16(that is 4xb32) elements for a
-    // thread.
-    auto resArgs = builder.newListOperand(4, "=r");
-    auto addrArg = builder.newAddrOperand(readPtr, "r");
-
-    auto ldmatrix = builder.create("ldmatrix.sync.aligned.m8n8.x4")
-                        ->o("trans", needTrans /*predicate*/)
-                        .o("shared.b16");
-    ldmatrix(resArgs, addrArg);
-
-    // The result type is 4xi32, each i32 is composed of 2xf16
-    // elements (adjacent two columns in a row) or a single f32 element.
-    Value resV4 = builder.launch(rewriter, loc, resTy);
+    auto ldMatrixOp = rewriter.create<nvgpu::LoadMatrixOp>(loc, resTy, readPtr);
+    ldMatrixOp->setAttr("trans", rewriter.getBoolAttr(needTrans));
+    auto resV4 = ldMatrixOp.getResult();
     return {extract_val(elemTy, resV4, 0), extract_val(elemTy, resV4, 1),
             extract_val(elemTy, resV4, 2), extract_val(elemTy, resV4, 3)};
   } else {
