Implement conversion from FMA dot operand to linear layout

include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td

@@ -766,6 +766,7 @@ for
       // Block encoding is dense stride layout. The elements per thread are contiguous.
       return getSizePerThread();
     };
+    SmallVector<unsigned> getRepOrderForOperand(int opIdx) const;
   }];
 
   let hasCustomAssemblyFormat = 1;

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM/SharedToDotOperandFMA.cpp

@@ -98,13 +98,13 @@ void storeValuesInLinearVector(PatternRewriter &rewriter, Location loc,
   }
 }
 
-void verifyCTALayout(CTALayoutAttr ctaLayout) {
+bool verifyCTALayout(CTALayoutAttr ctaLayout) {
   auto ctaSplit = ctaLayout.getCTASplitNum();
   for (auto split : ctaSplit) {
     if (split != 1)
-      llvm::report_fatal_error("tensors splited in CGA(thread group clusters) "
-                               "are not supported in FMA dot yet.");
+      return false;
   }
+  return true;
 }
 
 /// Get a linear offset of first element loaded by thread.
@@ -216,7 +216,8 @@ Value loadFMAOp(Value srcVal, Value llVal, BlockedEncodingAttr dLayout,
                 Value thread, Location loc,
                 const LLVMTypeConverter *typeConverter,
                 ConversionPatternRewriter &rewriter, const int dotOpNo) {
-  verifyCTALayout(dLayout.getCTALayout());
+  if (!verifyCTALayout(dLayout.getCTALayout()))
+    return Value();
 
   DimIdx dim;
   dim.batch = 0;
@@ -292,6 +293,11 @@ Value loadFMAOp(Value srcVal, Value llVal, BlockedEncodingAttr dLayout,
   auto numBTiles = std::max(1u, B / shapePerCTABTile);
   auto numNonKTiles = std::max(1u, NonK / shapePerCTANonKTile);
 
+  // Found discrepancy in this case,
+  // use linear layout based converter for this case
+  if (numBTiles != 1 || numNonKTiles != 1)
+    return Value();
+
   auto perThreadShape =
       getElemsPerThreadInOp(opTensorShape, shapePerCTATile, sizePerThread);
 

lib/Conversion/TritonGPUToLLVM/DotOpToLLVM/FMA.cpp

@@ -13,24 +13,54 @@ using ::mlir::triton::gpu::expandMatrixShapeWithBatch;
 using ::mlir::triton::gpu::getShapePerCTA;
 using ::mlir::triton::gpu::getSizePerThread;
 
-using ValueTableFMA = std::map<std::tuple<int, int, int>, Value>;
+/// \brief spatial position of repetition and register of a given value
+struct OperandValueKey {
+  unsigned bRepIdx, nonKRepIdx;
+  unsigned bIdx, nonKIdx, kIdx;
+
+  bool operator==(const OperandValueKey &other) const {
+    return (bRepIdx == other.bRepIdx && nonKRepIdx == other.nonKRepIdx &&
+            bIdx == other.bIdx && nonKIdx == other.nonKIdx &&
+            kIdx == other.kIdx);
+  }
+};
+
+template <> struct std::hash<OperandValueKey> {
+  std::size_t operator()(const OperandValueKey &k) const {
+    return std::hash<unsigned>()(k.bRepIdx) ^
+           (std::hash<unsigned>()(k.nonKRepIdx) << 1) ^
+           (std::hash<unsigned>()(k.bIdx) << 2) ^
+           (std::hash<unsigned>()(k.nonKIdx) << 3) ^
+           (std::hash<unsigned>()(k.kIdx) << 4);
+  }
+};
+
+using ValueTableFMA = std::unordered_map<OperandValueKey, Value>;
 
-static ValueTableFMA
-getValueTableFromStructFMA(Value val, ArrayRef<unsigned> perTileShape,
-                           unsigned kDim, unsigned nonKDim,
-                           ConversionPatternRewriter &rewriter, Location loc,
-                           ArrayRef<unsigned> order) {
+static ValueTableFMA getValueTableFromStructFMA(
+    Value val, ArrayRef<unsigned> perRepShape, ArrayRef<unsigned> repetitions,
+    unsigned kDim, unsigned nonKDim, ConversionPatternRewriter &rewriter,
+    Location loc, ArrayRef<unsigned> inRepOrder, ArrayRef<unsigned> repOrder) {
   ValueTableFMA res;
   auto elems = unpackLLElements(loc, val, rewriter);
-  assert(perTileShape.size() == 3);
-  assert(elems.size() == product(perTileShape));
+  assert(perRepShape.size() == 3);
+  auto numElemsRep = product(perRepShape);
+  assert(elems.size() == numElemsRep * product(repetitions));
   assert(kDim == 1 || kDim == 2);
   assert(nonKDim == 1 || nonKDim == 2);
   const unsigned bDim = 0;
 
   for (unsigned idx = 0; idx < elems.size(); ++idx) {
-    auto spatialIdx = mlir::LLVM::delinearize(idx, perTileShape, order);
-    res[{spatialIdx[bDim], spatialIdx[nonKDim], spatialIdx[kDim]}] = elems[idx];
+    auto inRepLinearIdx = idx % numElemsRep;
+    auto repLinearIdx = idx / numElemsRep;
+    auto inRepSpatialIdx =
+        mlir::LLVM::delinearize(inRepLinearIdx, perRepShape, inRepOrder);
+    auto repSpatialIdx =
+        mlir::LLVM::delinearize(repLinearIdx, repetitions, repOrder);
+    OperandValueKey key{repSpatialIdx[0], repSpatialIdx[nonKDim],
+                        inRepSpatialIdx[0], inRepSpatialIdx[nonKDim],
+                        inRepSpatialIdx[kDim]};
+    res[key] = elems[idx];
   }
   return res;
 }
@@ -54,46 +84,61 @@ LogicalResult convertFMADot(triton::DotOp op, triton::DotOp::Adaptor adaptor,
 
   BlockedEncodingAttr dLayout =
       cast<BlockedEncodingAttr>(dTensorTy.getEncoding());
-  auto order = expandMatrixOrderWithBatch(dLayout.getOrder());
+  // TODO process A and B operand separately
+  auto inRepOrder = expandMatrixOrderWithBatch(dLayout.getOrder());
+  auto repOrder = expandMatrixOrderWithBatch(dLayout.getRepOrder());
   auto cc = unpackLLElements(loc, adaptor.getC(), rewriter);
 
   Value llA = adaptor.getA();
   Value llB = adaptor.getB();
 
   auto sizePerThread =
       expandMatrixShapeWithBatch(ArrayRef(getSizePerThread(dLayout)));
+  auto numElemsPerThread = product(sizePerThread);
   auto shapePerCTATile =
       expandMatrixShapeWithBatch(ArrayRef(getShapePerCTATile(dLayout)));
 
   unsigned K = aShapePerCTA[2];
 
-  unsigned perThreadShape[3];
+  unsigned threadTileShape[3];
+  unsigned repetitions[3];
   for (int i = 0; i < 3; ++i) {
-    unsigned numRep = dShapePerCTA[i] / shapePerCTATile[i];
-    numRep = std::max(static_cast<unsigned>(1), numRep);
-    perThreadShape[i] = numRep * sizePerThread[i];
+    repetitions[i] =
+        ceil(dShapePerCTA[i], static_cast<int64_t>(shapePerCTATile[i]));
   }
 
   auto has = getValueTableFromStructFMA(
-      llA, {perThreadShape[0], perThreadShape[1], K},
-      /*kDim*/ 2, /*nonKDim*/ 1, rewriter, loc, order);
+      llA, {sizePerThread[0], sizePerThread[1], K},
+      {repetitions[0], repetitions[1], 1},
+      /*kDim*/ 2, /*nonKDim*/ 1, rewriter, loc, inRepOrder, repOrder);
   auto hbs = getValueTableFromStructFMA(
-      llB, {perThreadShape[0], K, perThreadShape[2]},
-      /*kDim*/ 1, /*nonKDim*/ 2, rewriter, loc, order);
+      llB, {sizePerThread[0], K, sizePerThread[2]},
+      {repetitions[0], 1, repetitions[2]},
+      /*kDim*/ 1, /*nonKDim*/ 2, rewriter, loc, inRepOrder, repOrder);
 
   SmallVector<Value> acc = cc;
 
-  for (unsigned b = 0; b < perThreadShape[0]; ++b)
-    for (unsigned m = 0; m < perThreadShape[1]; ++m)
-      for (unsigned n = 0; n < perThreadShape[2]; ++n) {
-        SmallVector<unsigned> multiDimAccumIdx = {b, m, n};
-        unsigned linearAccumIdx =
-            linearize(multiDimAccumIdx, perThreadShape, order);
-        for (unsigned k = 0; k < K; ++k) {
-          acc[linearAccumIdx] = rewriter.create<LLVM::FMulAddOp>(
-              loc, has[{b, m, k}], hbs[{b, n, k}], acc[linearAccumIdx]);
-        }
-      }
+  for (unsigned bRep = 0; bRep < repetitions[0]; ++bRep)
+    for (unsigned mRep = 0; mRep < repetitions[1]; ++mRep)
+      for (unsigned nRep = 0; nRep < repetitions[2]; ++nRep)
+        for (unsigned b = 0; b < sizePerThread[0]; ++b)
+          for (unsigned m = 0; m < sizePerThread[1]; ++m)
+            for (unsigned n = 0; n < sizePerThread[2]; ++n) {
+              SmallVector<unsigned> multiDimAccumIdx = {b, m, n};
+              unsigned linearInRepIdx =
+                  linearize(multiDimAccumIdx, sizePerThread, inRepOrder);
+              SmallVector<unsigned> multiDimRepIdx = {bRep, mRep, nRep};
+              unsigned linearRepIdx =
+                  linearize(multiDimRepIdx, repetitions, repOrder);
+              unsigned linearAccumIdx =
+                  linearInRepIdx + linearRepIdx * numElemsPerThread;
+              for (unsigned k = 0; k < K; ++k) {
+                auto aOp = has[{bRep, mRep, b, m, k}];
+                auto bOp = hbs[{bRep, nRep, b, n, k}];
+                acc[linearAccumIdx] = rewriter.create<LLVM::FMulAddOp>(
+                    loc, aOp, bOp, acc[linearAccumIdx]);
+              }
+            }
 
   auto res = packLLElements(loc, typeConverter, acc, rewriter, dTensorTy);
   rewriter.replaceOp(op, res);

lib/Conversion/TritonGPUToLLVM/MemoryOpToLLVM.cpp

@@ -126,7 +126,7 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
             LinearEncodingAttr>(dstLayout))
       return true;
     if (auto dot = dyn_cast<DotOperandEncodingAttr>(dstLayout)) {
-      if (isa<MmaEncodingTrait>(dot.getParent()))
+      if (isa<MmaEncodingTrait, BlockedEncodingAttr>(dot.getParent()))
         return true;
     }
     return false;
@@ -164,6 +164,8 @@ struct LocalLoadOpConversion : public ConvertOpToLLVMPattern<LocalLoadOp> {
     Value res = SharedToDotOperandFMA::convertLayout(
         dotLayout.getOpIdx(), op.getSrc(), adaptor.getSrc(), blockedLayout,
         thread, loc, getTypeConverter(), rewriter);
+    if (!res)
+      return failure();
     rewriter.replaceOp(op, res);
     return success();
   }

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1000,6 +1000,14 @@ unsigned NvidiaMmaEncodingAttr::getTotalElemsPerThread(ArrayRef<int64_t> shape,
   return product<unsigned>(getElemsPerThread(shape, eltTy));
 }
 
+// Blocked encoding
+
+SmallVector<unsigned>
+BlockedEncodingAttr::getRepOrderForOperand(int opIdx) const {
+  auto rank = getWarpsPerCTA().size();
+  return getOrderForDotOperand(opIdx, rank, /*kMajor*/ true);
+}
+
 //
 
 SmallVector<unsigned>

lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp

@@ -650,6 +650,45 @@ BlockedEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
   return combineCtaCgaWithShape(ctaLayout, getCTALayout(), shape);
 }
 
+std::optional<LinearLayout>
+fmaDotToLinearLayout(DotOperandEncodingAttr operandLayout,
+                     ArrayRef<int64_t> shape) {
+  int rank = shape.size();
+  auto blocked = cast<BlockedEncodingAttr>(operandLayout.getParent());
+  MLIRContext *ctx = operandLayout.getContext();
+
+  // TODO: introduce registerOrder or use getOrder(operandLayout)
+  // Currently this order is used in legacy converter, because we do not
+  // have access to full dot operand layout, only parent part.
+  auto regOrder = blocked.getOrder();
+  // TODO: use operandLayout.getThreadOrder()
+  auto threadOrder = blocked.getThreadOrder();
+  auto warpOrder = blocked.getWarpOrder();
+  auto repOrder = blocked.getRepOrder();
+
+  StringAttr kReg = S("register");
+  StringAttr kLane = S("lane");
+  StringAttr kWarp = S("warp");
+
+  SmallVector<unsigned> threadSize = blocked.getSizePerThread();
+  auto kDimIdx = operandLayout.getOpIdx() == 0 ? rank - 1 : rank - 2;
+  threadSize[kDimIdx] = shape[kDimIdx];
+  auto threadShape = blocked.getThreadsPerWarp();
+  auto warpShape = blocked.getWarpsPerCTA();
+
+  SmallVector<StringAttr> repDimNames =
+      permuteDimNames(standardOutDimNames(ctx, rank), repOrder);
+
+  LinearLayout ctaLayout = identityStandardND(kReg, threadSize, regOrder)
+                               .transposeOuts(repDimNames) *
+                           identityStandardND(kLane, threadShape, threadOrder)
+                               .transposeOuts(repDimNames) *
+                           identityStandardND(kWarp, warpShape, warpOrder)
+                               .transposeOuts(repDimNames);
+
+  return combineCtaCgaWithShape(ctaLayout, getCTALayout(operandLayout), shape);
+}
+
 LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef<unsigned> tileShape,
                            unsigned kWidth, ArrayRef<unsigned> order,
                            ArrayRef<unsigned> repOrder) {
@@ -750,9 +789,12 @@ LinearLayout nvidiaDotToLinearLayout(ArrayRef<int64_t> shape,
 std::optional<LinearLayout>
 DotOperandEncodingAttr::toLinearLayout(ArrayRef<int64_t> shape) const {
   auto parent = getParent();
-  if (auto mfmaLayout = llvm::dyn_cast<AMDMfmaEncodingAttr>(parent)) {
+  if (auto blockedLayout = mlir::dyn_cast<BlockedEncodingAttr>(parent)) {
+    return fmaDotToLinearLayout(*this, shape);
+  }
+  if (auto mfmaLayout = mlir::dyn_cast<AMDMfmaEncodingAttr>(parent)) {
     return mfmaDotToLinearLayout(*this, shape);
-  } else if (auto wmmaLayout = llvm::dyn_cast<AMDWmmaEncodingAttr>(parent)) {
+  } else if (auto wmmaLayout = mlir::dyn_cast<AMDWmmaEncodingAttr>(parent)) {
     return wmmaDotOperandToLinearLayout(*this, shape);
   } else if (auto mma = mlir::dyn_cast<NvidiaMmaEncodingAttr>(parent)) {
     return nvidiaDotToLinearLayout(shape, *this);

test/Conversion/amd/decompose-unsupported-conversions.mlir

@@ -97,11 +97,11 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
 #blocked = #ttg.blocked<{sizePerThread = [1, 32], threadsPerWarp = [32, 2], warpsPerCTA = [2, 2], order = [1, 0]}>
 #blocked1 = #ttg.blocked<{sizePerThread = [1, 32], threadsPerWarp = [32, 2], warpsPerCTA = [4, 1], order = [1, 0]}>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "hip:gfx940", "ttg.threads-per-warp" = 64 : i32} {
-  tt.func @neg_blocked_to_dot_op_incompatible_warp_gfx940(%arg0: tensor<32x32xf16, #blocked>) {
+  tt.func @neg_blocked_to_dot_op_incompatible_warp_gfx940(%arg0: tensor<128x128xf16, #blocked>) {
     // CHECK-NOT: ttg.convert_layout
     // CHECK: ttg.local_alloc
     // CHECK: ttg.local_load
-    %0 = ttg.convert_layout %arg0 : tensor<32x32xf16, #blocked> -> tensor<32x32xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked1}>>
+    %0 = ttg.convert_layout %arg0 : tensor<128x128xf16, #blocked> -> tensor<128x128xf16, #ttg.dot_op<{opIdx = 0, parent = #blocked1}>>
     tt.return
   }
 }

unittest/Dialect/TritonGPU/LinearLayoutConversionsTest.cpp

@@ -311,6 +311,58 @@ TEST_F(LinearLayoutConversionsTest, Blocked4D) {
                     {S("dim0"), S("dim1"), S("dim2"), S("dim3")}));
 }
 
+TEST_F(LinearLayoutConversionsTest, BlockedDotOperandLhs) {
+  auto parent = blocked(/*size*/ {2, 4}, /*threads*/ {8, 4}, /*warps*/ {2, 4},
+                        /*ctas*/ {1, 1}, /*splits*/ {1, 1}, /*order*/ {1, 0},
+                        /*cta order*/ {1, 0});
+  auto dotOperand = dot(parent, /*idx*/ 0, /*kWidth*/ 0);
+  EXPECT_EQ(
+      toLinearLayout({32, 16}, dotOperand),
+      LinearLayout({{S("register"), {{0, 1}, {0, 2}, {0, 4}, {0, 8}, {1, 0}}},
+                    {S("lane"), {{0, 0}, {0, 0}, {2, 0}, {4, 0}, {8, 0}}},
+                    {S("warp"), {{0, 0}, {0, 0}, {16, 0}}},
+                    {S("block"), {}}},
+                   {S("dim0"), S("dim1")}));
+}
+
+TEST_F(LinearLayoutConversionsTest, BlockedDotOperandLhs1) {
+  auto parent = blocked(/*size*/ {1, 1}, /*threads*/ {1, 32}, /*warps*/ {1, 1},
+                        /*ctas*/ {1, 1}, /*splits*/ {1, 1}, /*order*/ {1, 0},
+                        /*cta order*/ {1, 0});
+  auto dotOperand = dot(parent, /*idx*/ 0, /*kWidth*/ 0);
+  EXPECT_EQ(toLinearLayout({1, 64}, dotOperand),
+            LinearLayout({{S("register"),
+                           {{0, 1}, {0, 2}, {0, 4}, {0, 8}, {0, 16}, {0, 32}}},
+                          {S("lane"), {{0, 0}, {0, 0}, {0, 0}, {0, 0}, {0, 0}}},
+                          {S("warp"), {}},
+                          {S("block"), {}}},
+                         {S("dim0"), S("dim1")}));
+
+  auto dotOperand1 = dot(parent, /*idx*/ 1, /*kWidth*/ 0);
+  EXPECT_EQ(toLinearLayout({64, 64}, dotOperand1),
+            LinearLayout(
+                {{S("register"),
+                  {{1, 0}, {2, 0}, {4, 0}, {8, 0}, {16, 0}, {32, 0}, {0, 32}}},
+                 {S("lane"), {{0, 1}, {0, 2}, {0, 4}, {0, 8}, {0, 16}}},
+                 {S("warp"), {}},
+                 {S("block"), {}}},
+                {S("dim0"), S("dim1")}));
+}
+
+TEST_F(LinearLayoutConversionsTest, BlockedDotOperandRhs) {
+  auto parent = blocked(/*size*/ {2, 4}, /*threads*/ {8, 4}, /*warps*/ {2, 4},
+                        /*ctas*/ {1, 1}, /*splits*/ {1, 1}, /*order*/ {1, 0},
+                        /*cta order*/ {1, 0});
+  auto dotOperand = dot(parent, /*idx*/ 1, /*kWidth*/ 0);
+  EXPECT_EQ(toLinearLayout({16, 64}, dotOperand),
+            LinearLayout({{S("register"),
+                           {{0, 1}, {0, 2}, {1, 0}, {2, 0}, {4, 0}, {8, 0}}},
+                          {S("lane"), {{0, 4}, {0, 8}, {0, 0}, {0, 0}, {0, 0}}},
+                          {S("warp"), {{0, 16}, {0, 32}, {0, 0}}},
+                          {S("block"), {}}},
+                         {S("dim0"), S("dim1")}));
+}
+
 TEST_F(LinearLayoutConversionsTest, MMAv2_16x16) {
   EXPECT_EQ(toLinearLayout({16, 16},
                            mma(2, 0, {16, 8}, {1, 1}, {1, 1}, {1, 1}, {0, 1})),