[BACKEND] Take indices rather than reps as inputs to {smem,tmem}Load (#8623)

The previous code was doing a bad job at trying to guess the CTA-level
size of the tile being lowered. Here, we completely give up and instead
ask the caller to provide the starting coordinates of the subtensor they
want to lower rather than the reps.

In the passing, we also switch the `tmemLoad` logic to use LinearLayouts

Fixes triton-lang/triton#8606

Author: lezcano

test/Conversion/tritongpu_to_llvm_hopper.mlir

@@ -74,6 +74,24 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32} {
 
 // -----
 
+#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [16, 2], instrShape = [16, 256, 16]}>
+#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}>
+#smem = #ttg.shared_memory
+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 32 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} {
+  // CHECK-LABEL: @warp_group_dot_bf16_32_warps
+  tt.func @warp_group_dot_bf16_32_warps(
+      %a: !ttg.memdesc<256x128xbf16, #shared, #smem>,
+      %b: !ttg.memdesc<128x512xbf16, #shared, #smem>,
+      %acc: tensor<256x512xf32, #mma>) {
+    %res = ttng.warp_group_dot %a, %b, %acc {inputPrecision = 0 : i32, isAsync = true} :
+      !ttg.memdesc<256x128xbf16, #shared, #smem> * !ttg.memdesc<128x512xbf16, #shared, #smem> -> tensor<256x512xf32, #mma>
+    // CHECK: nvgpu.wgmma {{.*}} k = 16 : i32, layoutA = 1 : i32, layoutB = 1 : i32, m = 64 : i32, n = 256 : i32}
+    tt.return
+  }
+}
+
+// -----
+
 #mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], CTAsPerCGA = [1, 1], CTASplitNum = [1, 1], CTAOrder = [1, 0], instrShape = [16, 64, 16]}>
 #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}>
 #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}>

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h

@@ -41,6 +41,8 @@ struct MemDescOperand {
 class DotOpMmaMemLoader {
 public:
   virtual ~DotOpMmaMemLoader() = default;
+  // Given the starting coordinates of the logical tensor (i.e. reps *
+  // ctaTileSize), return the associated memory descriptor for SMEM / TMEM.
   virtual MemDescOperand memLoad(int a, int b,
                                  ConversionPatternRewriter &rewriter,
                                  Location loc) const = 0;
@@ -50,10 +52,8 @@ class DotOpMmaSmemLoader : public DotOpMmaMemLoader {
 public:
   DotOpMmaSmemLoader() = default;
 
-  DotOpMmaSmemLoader(MMASMEMDescriptor desc, Value baseb128, LinearLayout llInv,
-                     ArrayRef<unsigned> instrShape)
-      : desc(desc), baseb128(baseb128), ll(std::move(llInv)),
-        instrShape(instrShape) {}
+  DotOpMmaSmemLoader(MMASMEMDescriptor desc, Value baseb128, LinearLayout llInv)
+      : desc(desc), baseb128(baseb128), ll(std::move(llInv)) {}
 
   static DotOpMmaSmemLoader
   build(Location loc, RewriterBase &rewriter, gpu::MemDescType memTy,
@@ -136,15 +136,6 @@ class DotOpMmaSmemLoader : public DotOpMmaMemLoader {
                             {{kWarp, warpId}})[0]
               .second;
       baseSrcb128 = b.add(baseSrcb128, warpStrideb128);
-      // Increase the instruction shape to describe the size at a block level
-      // as the input just describes it at a warp level
-      int logwgAlongMN = 0;
-      for (int i = 0; i < warpGroupToOffsetb128.getInDimSizeLog2(kWarp); i++) {
-        if (warpGroupToOffsetb128.getBasis(kWarp, i, kOffset) != 0) {
-          logwgAlongMN++;
-        }
-      }
-      instrShape[MNdim] *= (1 << logwgAlongMN);
     }
 
     for (auto [dim, instrSize] : llvm::zip(ll.getInDimNames(), instrShape)) {
@@ -155,22 +146,18 @@ class DotOpMmaSmemLoader : public DotOpMmaMemLoader {
     auto desc = getDescriptor(ll, instrShape, bitwidth, MNdim, mmaVersion);
 
     Value baseb128 = b.zext(i64_ty, b.and_(baseSrcb128, b.i32_val(0x3FFF)));
-    return {desc, baseb128, ll, instrShape};
+    return {desc, baseb128, ll};
   }
 
   Value smemLoad(int a, int b, ConversionPatternRewriter &rewriter,
                  Location loc) const {
     auto *ctx = loc.getContext();
     auto tb = TritonLLVMOpBuilder(loc, rewriter);
     auto dims = to_vector(ll.getInDimNames());
-    assert((a + 1) * instrShape[0] <= ll.getInDimSize(dims[0]));
-    assert((b + 1) * instrShape[1] <= ll.getInDimSize(dims[1]));
     assert(to_vector(ll.getOutDimNames()) ==
            llvm::to_vector(
                ArrayRef<StringAttr>{str_attr("offset"), str_attr("block")}));
-    int32_t totalOffElems = ll.apply({{dims[0], a * instrShape[0]},
-                                      {dims[1], b * instrShape[1]}})[0]
-                                .second;
+    int32_t totalOffElems = ll.apply({{dims[0], a}, {dims[1], b}})[0].second;
     int32_t smemByteOffsetb8 = totalOffElems * desc.bitwidth / 8;
     auto currDesc = desc.descriptor;
     // Take the next 0/1/2/3 bits after the 128b tile
@@ -194,7 +181,6 @@ class DotOpMmaSmemLoader : public DotOpMmaMemLoader {
   MMASMEMDescriptor desc;
   Value baseb128;
   LinearLayout ll;
-  SmallVector<unsigned> instrShape;
 
   static MMASMEMDescriptor getDescriptor(const LinearLayout &ll,
                                          ArrayRef<unsigned> instrShape,
@@ -337,9 +323,9 @@ class DotOpMmaSmemLoader : public DotOpMmaMemLoader {
 class DotOpMmaV5TmemLoader : public DotOpMmaMemLoader {
 public:
   DotOpMmaV5TmemLoader() {}
-  DotOpMmaV5TmemLoader(Value tensor, Value base,
-                       SmallVector<unsigned int> instrShape, bool interleaved,
-                       bool trans);
+  static DotOpMmaV5TmemLoader build(Location loc, RewriterBase &rewriter,
+                                    gpu::MemDescType memTy, Value tmemBase);
+
   MemDescOperand tmemLoad(int a, int b, ConversionPatternRewriter &rewriter,
                           Location loc) const;
 
@@ -349,14 +335,12 @@ class DotOpMmaV5TmemLoader : public DotOpMmaMemLoader {
   }
 
 private:
-  Value base;
-  bool trans;
-  bool interleaved;
-  bool unpacked;
-  SmallVector<unsigned int> instrShape;
-  int numElementsPer32b;
-  int numRepM;
-  int numSlicePerBlockN;
+  DotOpMmaV5TmemLoader(LinearLayout ll, Value address, int bitwidth)
+      : ll(std::move(ll)), address(address), bitwidth(bitwidth) {}
+
+  LinearLayout ll;
+  Value address;
+  int bitwidth;
 };
 
 static Value getOffsetedBase(Value v, gpu::MemDescType memDescTy,

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

@@ -18,44 +18,25 @@ using ::mlir::triton::gpu::SharedLinearEncodingAttr;
 // DotOpMmaV5TmemLoader
 //===----------------------------------------------------------------------===//
 
-mlir::triton::NVIDIA::DotOpMmaV5TmemLoader::DotOpMmaV5TmemLoader(
-    Value tensor, Value base, SmallVector<unsigned int> instrShape,
-    bool interleaved, bool trans)
-    : base(base), instrShape(instrShape), interleaved(interleaved),
-      trans(trans) {
-  auto ty = cast<MemDescType>(tensor.getType());
-  auto tmemEncoding = cast<ttng::TensorMemoryEncodingAttr>(ty.getEncoding());
-  int elTyWidth = ty.getElementTypeBitWidth();
-  unpacked = tmemEncoding.getColStride() != 1;
-  // When using TMEM to store operands mma operands the TMEM block size may be
-  // smaller than mma k block. Therefore we need to adjust the offset
-  // calculation.
-  numSlicePerBlockN = tmemEncoding.getBlockN() / instrShape[1];
-  numElementsPer32b = 32 / (elTyWidth * tmemEncoding.getColStride());
-  auto shapePerCTA = triton::gpu::getShapePerCTA(ty);
-  numRepM = ceil<unsigned>(shapePerCTA[0], instrShape[0]);
+DotOpMmaV5TmemLoader mlir::triton::NVIDIA::DotOpMmaV5TmemLoader::build(
+    Location loc, RewriterBase &rewriter, gpu::MemDescType memTy,
+    Value tmemBase) {
+  auto ctx = loc.getContext();
+  auto ll = toLinearLayout(memTy);
+  auto layout = cast<ttng::TensorMemoryEncodingAttr>(memTy.getEncoding());
+  auto bitwidth = memTy.getElementTypeBitWidth();
+  auto tb = TritonLLVMOpBuilder(loc, rewriter);
+  Value address = tb.ptrtoint(i32_ty, tmemBase);
+  return DotOpMmaV5TmemLoader(ll.pseudoinvert(), address, bitwidth);
 }
 
 MemDescOperand mlir::triton::NVIDIA::DotOpMmaV5TmemLoader::tmemLoad(
     int a, int b, ConversionPatternRewriter &rewriter, Location loc) const {
-  int numRows = 64;
-  if (interleaved || instrShape[0] >= 128)
-    numRows = 128;
-  int numColPerBlock =
-      ((instrShape[0] * numSlicePerBlockN * instrShape[1]) / numRows) /
-      numElementsPer32b;
-  int blockId = a + (b / numSlicePerBlockN) * numRepM;
-  int offset;
-  if (!interleaved) {
-    offset = numColPerBlock * blockId;
-  } else {
-    int blockIdIsOdd = blockId & 1;
-    int blockIdPrevEven = blockId - blockIdIsOdd;
-    offset = numColPerBlock * blockIdPrevEven + ((16 * blockIdIsOdd) << 16);
-  }
-  offset += (b % numSlicePerBlockN) * (instrShape[1] / numElementsPer32b);
-  auto tb = TritonLLVMOpBuilder(loc, rewriter);
-  Value address = tb.ptrtoint(i32_ty, base);
+  auto dims = to_vector(ll.getInDimNames());
+  auto rowCol = ll.apply({{dims[0], a}, {dims[1], b}});
+  int row = rowCol[0].second;
+  int col = rowCol[1].second * bitwidth / 32;
+  int offset = col | (row << 16);
   return {address, offset};
 }
 
@@ -445,8 +426,8 @@ void convertDotImpl(const LLVMTypeConverter &typeConverter,
   std::unique_ptr<DotOpMmaMemLoader> aLoader;
   bool transA = false;
   if (aInTmem) {
-    aLoader = std::make_unique<DotOpMmaV5TmemLoader>(a, baseA, aOperandShape,
-                                                     interleaved, transA);
+    aLoader = std::make_unique<DotOpMmaV5TmemLoader>(
+        DotOpMmaV5TmemLoader::build(loc, rewriter, aTensorTy, baseA));
   } else {
     auto isFp4a = op.numBitsPerElementA == 4;
     aLoader = std::make_unique<DotOpMmaSmemLoader>(DotOpMmaSmemLoader::build(
@@ -479,8 +460,9 @@ void convertDotImpl(const LLVMTypeConverter &typeConverter,
       Value useInitAcc = useDFlag;
       MemDescOperand accAddress = op.getAccAddress(rewriter, loc, m, n, desc);
       for (int k = 0; k < numRepK; k++) {
-        MemDescOperand a = aLoader->memLoad(m, k, rewriter, loc);
-        Value b = bLoader.smemLoad(k, n, rewriter, loc);
+        MemDescOperand a =
+            aLoader->memLoad(m * mmaSizeM, k * mmaSizeK, rewriter, loc);
+        Value b = bLoader.smemLoad(k * mmaSizeK, n * mmaSizeN, rewriter, loc);
         op.createMMAInst(rewriter, loc, accAddress, a, b, elect, useInitAcc,
                          desc, m, n, k);
         useInitAcc = tb.i1_val(1);
@@ -503,6 +485,7 @@ void convertDot(const LLVMTypeConverter &typeConverter,
   MemDescType aTensorTy = op.getA().getType();
   MemDescType bTensorTy = op.getB().getType();
   MemDescType dTensorTy = op.getD().getType();
+  auto dLayout = cast<ttng::TensorMemoryEncodingAttr>(dTensorTy.getEncoding());
   bool twoCTAs = op.getTwoCtas();
 
   DotConversion dot;
@@ -518,12 +501,12 @@ void convertDot(const LLVMTypeConverter &typeConverter,
   dot.numBitsPerElementA = aTensorTy.getElementTypeBitWidth();
   dot.numBitsPerElementB = bTensorTy.getElementTypeBitWidth();
 
+  DotOpMmaV5TmemLoader dLoader =
+      DotOpMmaV5TmemLoader::build(loc, rewriter, dTensorTy, adaptor.getD());
   dot.getAccAddress = [&](ConversionPatternRewriter &rewriter, Location loc,
                           int m, int n, const DotConversion::InstDesc &desc) {
-    DotOpMmaV5TmemLoader dLoader = DotOpMmaV5TmemLoader(
-        op.getD(), adaptor.getD(), {desc.mmaSizeM, desc.mmaSizeN},
-        desc.interleaved, /*trans=*/false);
-    return dLoader.tmemLoad(m, n, rewriter, loc);
+    return dLoader.tmemLoad(m * dLayout.getBlockM(), n * dLayout.getBlockN(),
+                            rewriter, loc);
   };
 
   dot.createMMAInst = [&](ConversionPatternRewriter &rewriter, Location loc,

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

@@ -204,19 +204,21 @@ LogicalResult convertDot(const LLVMTypeConverter *typeConverter,
   auto baseB = getOffsetedBase(loadedB, cast<MemDescType>(bTensorTy),
                                typeConverter, rewriter, loc);
   auto dShapePerCTA = getShapePerCTA(dTensorTy);
-  auto instrShape = mmaEncoding.getInstrShape();
-  auto accSize = 2 * (instrShape[1] / 4);
-  unsigned M = 4 * instrShape[0];
-  unsigned N = instrShape[1];
-  unsigned K = instrShape[2];
-  bool zeroAcc = isZeroConst(c);
   auto instrMNK = mmaEncoding.getInstrShape();
+  auto accSize = 2 * (instrMNK[1] / 4);
+  unsigned M = 4 * instrMNK[0];
+  unsigned N = instrMNK[1];
+  unsigned K = instrMNK[2];
+  bool zeroAcc = isZeroConst(c);
   auto warpSize = mmaEncoding.getWarpsPerCTA();
   auto shapePerCTATile = SmallVector<unsigned>{instrMNK[0] * warpSize[0],
                                                instrMNK[1] * warpSize[1]};
-  int numRepM = ceil<unsigned>(dShapePerCTA[0], shapePerCTATile[0]);
-  int numRepN = ceil<unsigned>(dShapePerCTA[1], shapePerCTATile[1]);
-  int numRepK = ceil<unsigned>(aTensorTy.getShape()[1], instrShape[2]);
+  unsigned mmaSizeM = shapePerCTATile[0];
+  unsigned mmaSizeN = shapePerCTATile[1];
+  unsigned mmaSizeK = instrMNK[2];
+  int numRepM = ceil<unsigned>(dShapePerCTA[0], mmaSizeM);
+  int numRepN = ceil<unsigned>(dShapePerCTA[1], mmaSizeN);
+  int numRepK = ceil<unsigned>(aTensorTy.getShape()[1], mmaSizeK);
   DotOpMmaSmemLoader aLoader;
   SmallVector<Value> structA;
   auto warpGroups = {warpSize[0] / 4, warpSize[1]};
@@ -270,14 +272,14 @@ LogicalResult convertDot(const LLVMTypeConverter *typeConverter,
       for (int k = 0; k < numRepK; ++k) {
         Value a;
         if (aInShared) {
-          a = aLoader.smemLoad(m, k, rewriter, loc);
+          a = aLoader.smemLoad(m * mmaSizeM, k * mmaSizeK, rewriter, loc);
         } else {
           auto aDotOpEnc =
               cast<DotOperandEncodingAttr>(aTensorTy.getEncoding());
           assert(aDotOpEnc.getKWidth() ==
                  32 / aTensorTy.getElementTypeBitWidth());
 
-          unsigned regASize = (instrShape[0] * instrShape[2]) / 32;
+          unsigned regASize = (instrMNK[0] * instrMNK[2]) / 32;
           llvm::SmallVector<Value> regA =
               loadReg(rewriter, loc, structA, (m * numRepK + k) * regASize,
                       regASize, startSequence);
@@ -286,7 +288,7 @@ LogicalResult convertDot(const LLVMTypeConverter *typeConverter,
               SmallVector<Type>(regA.size(), regA[0].getType()));
           a = packLLElements(loc, typeConverter, regA, rewriter, regATy);
         }
-        auto b = bLoader.smemLoad(k, n, rewriter, loc);
+        auto b = bLoader.smemLoad(k * mmaSizeK, n * mmaSizeN, rewriter, loc);
         numLowPrecisionAcc += K;
         // If using native accumulation would cause use to do more low precion
         // accumulation than allowed do a separate allocation.

third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TensorMemoryToLLVM.cpp

@@ -604,9 +604,7 @@ static void copySharedToTmem(ConversionPatternRewriter &rewriter, Location loc,
   }
 
   for (int col = 0; col < cvt.getInDimSize(kCol); col += instrShape[1]) {
-    // smemLoad takes the colRep. It'd be nice to change this but we would need
-    // to change the wgmma and mmav5 lowering
-    auto desc = loader.smemLoad(0, col / instrShape[1], rewriter, loc);
+    auto desc = loader.smemLoad(0, col, rewriter, loc);
     auto tmemAddr =
         b.or_(b.ptrtoint(i32_ty, baseDst), b.i32_val(col * bitwidth / 32),
               /*disjoint=*/true);