[IR][BACKEND] Add the trans attribute to LoadMatrix and StoreMatrix ops (#5467)

Following @ThomasRaoux's suggestion, this way we ensure that the trans
attribute is not discarded and keeps the IR clear

Author: Jokeren

test/Conversion/nvgpu_to_llvm.mlir

@@ -41,6 +41,8 @@ llvm.func @cluster_id() -> i32 {
 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
+  // CHECK: stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [$0], {$1, $2, $3, $4};
+  nvgpu.stmatrix %ptr, %i, %i, %i, %i {trans} : !llvm.ptr<3>, i32, i32, i32, i32
   llvm.return
 }
 
@@ -50,7 +52,11 @@ llvm.func @stmatrix(%i: i32, %ptr: !llvm.ptr<3>) {
 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)>
+  // CHECK: ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {$0, $1, $2, $3}, [$4];
+  %1 = nvgpu.ldmatrix %ptr {trans} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
+  %2 = llvm.extractvalue %1[0] : !llvm.struct<(i32, i32, i32, i32)>
+  %3 = llvm.insertvalue %2, %0[0] : !llvm.struct<(i32, i32, i32, i32)>
+  llvm.return %3 : !llvm.struct<(i32, i32, i32, i32)>
 }
 
 // -----

third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUOps.td

@@ -101,12 +101,19 @@ def NVGPU_ClusterWaitOp : NVGPU_Op<"cluster_wait", []> {
 }
 
 def NVGPU_StoreMatrixOp : NVGPU_Op<"stmatrix", [MemoryEffects<[MemWrite]>]> {
-  let arguments = (ins LLVM_PointerShared:$addr, Variadic<I32>:$vals);
+  let arguments = (
+    ins LLVM_PointerShared:$addr,
+    Variadic<I32>:$vals,
+    UnitAttr:$trans
+  );
   let assemblyFormat = "operands attr-dict `:` type(operands)";
 }
 
 def NVGPU_LoadMatrixOp : NVGPU_Op<"ldmatrix", [MemoryEffects<[MemRead]>]> {
-  let arguments = (ins LLVM_PointerShared:$addr);
+  let arguments = (
+    ins LLVM_PointerShared:$addr,
+    UnitAttr:$trans
+  );
   let results = (outs LLVM_AnyStruct:$result);
   let assemblyFormat = "$addr attr-dict `:` functional-type($addr, $result)";
 }

third_party/nvidia/lib/NVGPUToLLVM/NVGPUToLLVMPass.cpp

@@ -243,10 +243,7 @@ class MatrixOpPattern : public OpRewritePattern<MatrixOpType> {
   LogicalResult matchAndRewrite(MatrixOpType op,
                                 PatternRewriter &rewriter) const override {
     unsigned vecSize = getVectorSize(op);
-    bool trans = op->hasAttr("trans")
-                     ? op->template getAttrOfType<BoolAttr>("trans").getValue()
-                     : false;
-
+    bool trans = op.getTrans();
     // Template method for PTX assembly generation
     std::string ptxAsm =
         (llvm::Twine(ConcreteMatrixOpPattern::kOpCode) +

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ConvertLayoutOpToLLVM/SharedToDotOperandMMAv2OrV3.cpp

@@ -341,8 +341,8 @@ MMA16816SmemLoader::loadX4(int batch, int mat0, int mat1, ArrayRef<Value> ptrs,
       stridedOffset = add(
           stridedOffset, mul(i32_val(batch * warpsPerCTA[0]), smemBatchOffset));
     Value readPtr = gep(ptr_ty(ctx, 3), shemTy, ptr, stridedOffset);
-    auto ldMatrixOp = rewriter.create<nvgpu::LoadMatrixOp>(loc, resTy, readPtr);
-    ldMatrixOp->setAttr("trans", rewriter.getBoolAttr(needTrans));
+    auto ldMatrixOp =
+        rewriter.create<nvgpu::LoadMatrixOp>(loc, resTy, readPtr, 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)};