[BACKEND] Fix wrong K dimension for dot_scaled op (#6269)

When we have a lhs of 4bits the K dimension calculated was wrong due to
ambiguous meaning of some variable.

Author: ThomasRaoux

test/Conversion/tritongpu_to_llvm_blackwell.mlir

@@ -229,6 +229,44 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 
 // -----
 
+#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 64, transposed = false, elementBitWidth = 8}>
+#shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 64, transposed = false, elementBitWidth = 8}>
+#shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
+#tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, unpacked = true>
+#tmem_scales = #ttng.tensor_memory_scales_encoding<>
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+  // CHECK-LABEL: @tc_gen5_mma_block_scale_fp4_a
+  // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144769664 : i32) : i32
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
+  // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
+  // CHECK: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
+  // CHECK: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32
+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]]
+  tt.func @tc_gen5_mma_block_scale_fp4_a(%a: !ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>,
+                       %b: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
+                       %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
+                       %scale_a: !ttg.memdesc<128x2xi8, #tmem_scales, #ttng.tensor_memory>,
+                       %scale_b: !ttg.memdesc<128x2xi8, #tmem_scales, #ttng.tensor_memory>,
+                       %useAcc: i1,
+                       %pred: i1,
+                       %barrier: !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>) {
+    ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e4m3, %barrier :
+    (!ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>,
+    !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
+    !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
+    !ttg.memdesc<128x2xi8, #tmem_scales, #ttng.tensor_memory>,
+    !ttg.memdesc<128x2xi8, #tmem_scales, #ttng.tensor_memory>,
+    i1,
+    i1,
+    !ttg.memdesc<1xi64, #shared2, #ttg.shared_memory, mutable>) -> ()
+    tt.return
+  }
+}
+
+// -----
+
 #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 64, transposed = false, elementBitWidth = 16, CTAsPerCGA = [2, 1], CTASplitNum = [2, 1], CTAOrder = [1, 0]}>
 #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 64, transposed = false, elementBitWidth = 16, CTAsPerCGA = [1, 2], CTASplitNum = [1, 2], CTAOrder = [1, 0]}>
 #shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0], CTAsPerCGA = [2], CTASplitNum = [1], CTAOrder = [0]}>

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp

@@ -511,11 +511,14 @@ struct TCGen5MMAScaledOpConversion
 
     unsigned int M = dTensorTy.getDimSize(0);
     unsigned int N = dTensorTy.getDimSize(1);
-    int numBitsPerElementA = opKindIsMXFP4 ? getFormatBitSize(op.getAType())
-                                           : aTensorTy.getElementTypeBitWidth();
-    int numBitsPerElementB = opKindIsMXFP4 ? getFormatBitSize(op.getBType())
-                                           : bTensorTy.getElementTypeBitWidth();
-    unsigned int K = (aTensorTy.getDimSize(1) * 8) / numBitsPerElementA;
+    int numBitsUnpackedPerElementA = opKindIsMXFP4
+                                         ? getFormatBitSize(op.getAType())
+                                         : aTensorTy.getElementTypeBitWidth();
+    int numBitsUnpackedPerElementB = opKindIsMXFP4
+                                         ? getFormatBitSize(op.getBType())
+                                         : bTensorTy.getElementTypeBitWidth();
+    unsigned int K =
+        (aTensorTy.getDimSize(1) * 8) / getFormatBitSize(op.getAType());
 
     // Get MMA size based on acc layout.
     auto tensorMemAttr = cast<triton::nvidia_gpu::TensorMemoryEncodingAttr>(
@@ -546,12 +549,12 @@ struct TCGen5MMAScaledOpConversion
     } else {
       aLoader = std::make_unique<DotOpMmaV3SmemLoader>(
           op.getA(), baseA, shapeA, shapeA, zero, 1, transA, aOperandShape,
-          numBitsPerElementA, rewriter, loc);
+          numBitsUnpackedPerElementA, rewriter, loc);
     }
     DotOpMmaV3SmemLoader bLoader =
         DotOpMmaV3SmemLoader(op.getB(), baseB, shapeB, shapeB, zero, 1, transB,
                              {(unsigned)mmaSizeN, (unsigned)mmaSizeK},
-                             numBitsPerElementB, rewriter, loc);
+                             numBitsUnpackedPerElementB, rewriter, loc);
 
     // Only run mma on one thread. We currently use elect as ptxas is not able
     // to detect that tid.x == 0 is true only for 1 thread.