[Blackwell] Support optional scale TMAs in warp specialization for tl.dot_scaled (triton-lang#6551)

[Blackwell] Support optional scale TMAs in warp specialization for tl.dot_scaled

This enables automatic warp specialization for block scaled workloads.

Author: csullivan

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/LoadMMASpecialization.cpp

@@ -207,6 +207,17 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
                              "loads for `tt.dot` operands");
   }
 
+  SmallVector<Operation *> aScaleChain, bScaleChain;
+  auto scaledMMAOp = dyn_cast<ttng::TCGen5MMAScaledOp>(mmaOp.getOperation());
+  if (scaledMMAOp) {
+    if (failed(
+            findSingleChainToLoad(loop, scaledMMAOp.getAScale(), aScaleChain)))
+      aScaleChain.clear();
+    if (failed(
+            findSingleChainToLoad(loop, scaledMMAOp.getBScale(), bScaleChain)))
+      bScaleChain.clear();
+  }
+
   ttng::TMEMAllocOp oldAccAlloc =
       mmaOp.getAccumulator().getDefiningOp<ttng::TMEMAllocOp>();
   if (!oldAccAlloc)
@@ -218,7 +229,9 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
 
   // Determine if the MMA accumulator can be multibuffered.
   auto isLoadPipelineable = [&](Operation *op) {
-    return llvm::is_contained({aChain.back(), bChain.back()}, op);
+    return llvm::is_contained(llvm::to_vector(llvm::concat<Operation *>(
+                                  aChain, bChain, aScaleChain, bScaleChain)),
+                              op);
   };
   bool accIsMultiBuffered =
       // All operand feeds are pipelineable.
@@ -280,16 +293,27 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
   Partition *mmaPartition = schedule.addPartition(numStages);
 
   // Multi-buffer the loads.
-  auto [loadIndex, loadPhase] = addIndexAndPhase(b, loop, numStages);
+  BlockArgument loadIndex;
+  BlockArgument loadPhase;
+  std::tie(loadIndex, loadPhase) = addIndexAndPhase(b, loop, numStages);
+
+  auto allocate = [&](const SmallVector<Operation *> &chain)
+      -> std::tuple<Operation *, RankedTensorType, SharedEncodingTrait, Value> {
+    if (chain.empty())
+      return {nullptr, RankedTensorType(), SharedEncodingTrait(), Value()};
+
+    Operation *load = chain.back();
+    auto type = cast<RankedTensorType>(load->getResult(0).getType());
+    SharedEncodingTrait enc = getSharedEncoding(chain.back());
+    Value alloc = createAlloc(loop, type, load->getLoc(), enc, numStages);
+
+    return {load, type, enc, alloc};
+  };
 
-  Operation *aLoad = aChain.back();
-  Operation *bLoad = bChain.back();
-  auto aType = cast<RankedTensorType>(aLoad->getResult(0).getType());
-  auto bType = cast<RankedTensorType>(bLoad->getResult(0).getType());
-  SharedEncodingTrait aEnc = getSharedEncoding(aChain.back());
-  SharedEncodingTrait bEnc = getSharedEncoding(bChain.back());
-  Value aAlloc = createAlloc(loop, aType, aLoad->getLoc(), aEnc, numStages);
-  Value bAlloc = createAlloc(loop, bType, bLoad->getLoc(), bEnc, numStages);
+  auto [aLoad, aType, aEnc, aAlloc] = allocate(aChain);
+  auto [bLoad, bType, bEnc, bAlloc] = allocate(bChain);
+  auto [aScaleLoad, aScaleType, aScaleEnc, aScaleAlloc] = allocate(aScaleChain);
+  auto [bScaleLoad, bScaleType, bScaleEnc, bScaleAlloc] = allocate(bScaleChain);
 
   // Share the same set of barriers for both.
   Value emptyBars = createBarrierAlloc(loop, numStages);
@@ -304,9 +328,23 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
   int loadSizeInBytes =
       product(aType.getShape()) * aType.getElementTypeBitWidth() / 8 +
       product(bType.getShape()) * bType.getElementTypeBitWidth() / 8;
+  if (aScaleLoad)
+    loadSizeInBytes += product(aScaleType.getShape()) *
+                       aScaleType.getElementTypeBitWidth() / 8;
+  if (bScaleLoad)
+    loadSizeInBytes += product(bScaleType.getShape()) *
+                       bScaleType.getElementTypeBitWidth() / 8;
 
   // Insert before the group of loads.
-  b.setInsertionPoint(aLoad->isBeforeInBlock(bLoad) ? aLoad : bLoad);
+  SmallVector<Operation *> allLoads{aLoad, bLoad};
+  if (aScaleLoad)
+    allLoads.push_back(aScaleLoad);
+  if (bScaleLoad)
+    allLoads.push_back(bScaleLoad);
+  std::sort(allLoads.begin(), allLoads.end(),
+            [](Operation *a, Operation *b) { return a->isBeforeInBlock(b); });
+  b.setInsertionPoint(allLoads.front());
+
   // Wait for the buffer to be empty and the corresponding barrier to be
   // exhausted.
   Value curEmptyBar = createSingleBufferView(b, emptyBars, loadIndex);
@@ -318,19 +356,21 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
                                            loadSizeInBytes, intCst(true, 1));
 
   // Replace the loads with async copies.
-  b.setInsertionPoint(aLoad);
-  Value aView = createSingleBufferView(b, aAlloc, loadIndex);
-  lowerTMACopy(b, *loadPartition, aLoad, curLoadBar, aView);
-  replaceUsesAndPropagateType(b, *aLoad->user_begin(), aView);
-  aLoad->user_begin()->erase();
-  aLoad->erase();
-
-  b.setInsertionPoint(bLoad);
-  Value bView = createSingleBufferView(b, bAlloc, loadIndex);
-  lowerTMACopy(b, *loadPartition, bLoad, curLoadBar, bView);
-  replaceUsesAndPropagateType(b, *bLoad->user_begin(), bView);
-  bLoad->user_begin()->erase();
-  bLoad->erase();
+  auto lowerLoadAndPropagate = [&](Operation *load, Value alloc,
+                                   Value barrier) {
+    b.setInsertionPoint(load);
+    Value view = createSingleBufferView(b, alloc, loadIndex);
+    lowerTMACopy(b, *loadPartition, load, barrier, view);
+    replaceUsesAndPropagateType(b, *load->user_begin(), view);
+    load->user_begin()->erase();
+    load->erase();
+  };
+  lowerLoadAndPropagate(aLoad, aAlloc, curLoadBar);
+  lowerLoadAndPropagate(bLoad, bAlloc, curLoadBar);
+  if (aScaleLoad)
+    lowerLoadAndPropagate(aScaleLoad, aScaleAlloc, curLoadBar);
+  if (bScaleLoad)
+    lowerLoadAndPropagate(bScaleLoad, bScaleAlloc, curLoadBar);
 
   // Place the remaining users in the MMA partition. Re-acquire the use chain
   // because some ops were invalidated by `replaceUsesAndPropagateType`.
@@ -339,9 +379,18 @@ LogicalResult triton::gpu::specializeLoadMMADependencies(scf::ForOp &loop,
   aChain.push_back(mmaOp);
   (void)findSingleChainToLoad(loop, dot.getA(), aChain);
   (void)findSingleChainToLoad(loop, dot.getB(), bChain);
+  if (aScaleLoad) {
+    aScaleChain.clear();
+    (void)findSingleChainToLoad(loop, scaledMMAOp.getAScale(), aScaleChain);
+  }
+  if (bScaleLoad) {
+    bScaleChain.clear();
+    (void)findSingleChainToLoad(loop, scaledMMAOp.getBScale(), bScaleChain);
+  }
 
   // Place users in the MMA partition.
-  auto allUsers = llvm::to_vector(llvm::concat<Operation *>(aChain, bChain));
+  auto allUsers = llvm::to_vector(
+      llvm::concat<Operation *>(aChain, bChain, aScaleChain, bScaleChain));
   for (Operation *user : allUsers)
     mmaPartition->insert(user);
 

test/TritonGPU/load-mma-specialization.mlir

@@ -762,4 +762,59 @@ tt.func @matmul_tma_acc_with_conditional_def_and_use_no_multibuf_flag(
   tt.return
 }
 
+
+
+tt.func @matmul_scaled_rhs_scales_tma(
+  // CHECK-LABEL: @matmul_scaled_rhs_scales_tma
+  %k_tiles: i32,
+  %off_m: i32,
+  %off_n: i32,
+  %a_desc: !tt.tensordesc<tensor<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>>>,
+  %b_desc: !tt.tensordesc<tensor<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>>>,
+  %b_scale_desc: !tt.tensordesc<tensor<128x8xi8, #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [4, 3, 2, 1, 0]}>>>
+) {
+  %true = arith.constant true
+  %c0_i32 = arith.constant 0 : i32
+  %c1_i32 = arith.constant 1 : i32
+  %BLOCK_K = arith.constant 64 : i32
+  %zero = arith.constant dense<0.0> : tensor<128x128xf32, #acc_layout>
+
+  %a_scales_const = arith.constant dense<127> : tensor<128x8xi8, #oper_layout>
+  %a_scales_tmem = ttng.tmem_alloc %a_scales_const : (tensor<128x8xi8, #oper_layout>) -> !ttg.memdesc<128x8xi8, #ttng.tensor_memory_scales_encoding<>, #ttng.tensor_memory>
+
+  %result = scf.for %k = %c0_i32 to %k_tiles step %c1_i32 iter_args(%acc = %zero) -> tensor<128x128xf32, #acc_layout> : i32 {
+    %off_k = arith.muli %k, %BLOCK_K : i32
+
+    // CHECK: %{{[0-9]+}} = ttg.memdesc_subview %{{[0-9]+}}[%arg7, %c0_i32, %c0_i32]
+    // CHECK-NEXT: %{{[0-9]+}} = ttng.tensor_desc_to_tma_ptr %arg3 {ttg.partition = 0 : i32}
+    // CHECK-NEXT: ttng.async_tma_copy_global_to_local %{{[0-9]+}}[%arg1, %{{[0-9]+}}] %{{[0-9]+}}, %{{[0-9]+}}, %true {ttg.partition = 0 : i32}
+    %a_reg = tt.descriptor_load %a_desc[%off_m, %off_k] : !tt.tensordesc<tensor<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>>> -> tensor<128x64xf8E4M3FN, #oper_layout>
+
+    // CHECK-NEXT: %{{[0-9]+}} = ttg.memdesc_subview %{{[0-9]+}}[%arg7, %c0_i32, %c0_i32]
+    // CHECK-NEXT: %{{[0-9]+}} = ttng.tensor_desc_to_tma_ptr %arg4 {ttg.partition = 0 : i32}
+    // CHECK-NEXT: ttng.async_tma_copy_global_to_local %{{[0-9]+}}[%arg2, %{{[0-9]+}}] %{{[0-9]+}}, %{{[0-9]+}}, %true {ttg.partition = 0 : i32}
+    %b_reg = tt.descriptor_load %b_desc[%off_n, %off_k] : !tt.tensordesc<tensor<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>>> -> tensor<128x64xf8E4M3FN, #oper_layout>
+
+    // CHECK-NEXT: %{{[0-9]+}} = ttg.memdesc_subview %{{[0-9]+}}[%arg7, %c0_i32, %c0_i32]
+    // CHECK-NEXT: %{{[0-9]+}} = ttng.tensor_desc_to_tma_ptr %arg5 {ttg.partition = 0 : i32}
+    // CHECK-NEXT: ttng.async_tma_copy_global_to_local %{{[0-9]+}}[%arg1, %c0_i32] %{{[0-9]+}}, %{{[0-9]+}}, %true {ttg.partition = 0 : i32}
+    %b_scales_reg = tt.descriptor_load %b_scale_desc[%off_m, %c0_i32] : !tt.tensordesc<tensor<128x8xi8, #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [4, 3, 2, 1, 0]}>>> -> tensor<128x8xi8, #oper_layout>
+
+    %a_sh = ttg.local_alloc %a_reg : (tensor<128x64xf8E4M3FN, #oper_layout>) -> !ttg.memdesc<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>, #smem>
+    %b_sh_raw = ttg.local_alloc %b_reg : (tensor<128x64xf8E4M3FN, #oper_layout>) -> !ttg.memdesc<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>, #smem>
+    %b_sh = ttg.memdesc_trans %b_sh_raw {order = array<i32: 1, 0>} : !ttg.memdesc<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>, #smem> -> !ttg.memdesc<64x128xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 8}>, #smem>
+
+    %b_scales_tmem = ttng.tmem_alloc %b_scales_reg : (tensor<128x8xi8, #oper_layout>) -> !ttg.memdesc<128x8xi8, #ttng.tensor_memory_scales_encoding<>, #ttng.tensor_memory>
+
+    %c_tmem = ttng.tmem_alloc %acc : (tensor<128x128xf32, #acc_layout>) -> !ttg.memdesc<128x128xf32, #acc_tmem, #ttng.tensor_memory, mutable>
+
+    ttng.tc_gen5_mma_scaled %a_sh, %b_sh, %c_tmem, %a_scales_tmem, %b_scales_tmem, %true, %true lhs = e4m3 rhs = e4m3 : !ttg.memdesc<128x64xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 8}>, #smem>, !ttg.memdesc<64x128xf8E4M3FN, #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 8}>, #smem>, !ttg.memdesc<128x128xf32, #acc_tmem, #ttng.tensor_memory, mutable>, !ttg.memdesc<128x8xi8, #ttng.tensor_memory_scales_encoding<>, #ttng.tensor_memory>, !ttg.memdesc<128x8xi8, #ttng.tensor_memory_scales_encoding<>, #ttng.tensor_memory>
+
+    %c = ttng.tmem_load %c_tmem : !ttg.memdesc<128x128xf32, #acc_tmem, #ttng.tensor_memory, mutable> -> tensor<128x128xf32, #acc_layout>
+    scf.yield %c : tensor<128x128xf32, #acc_layout>
+  } {tt.warp_specialize}
+
+  tt.return
+}
+
 }