Merge commit '570f24d016702cbfe1179beae3eb03d24e9c6b40'

Author: anmyachev

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -2012,6 +2012,13 @@ SwizzledSharedEncodingAttr AMDMfmaEncodingAttr::composeSharedLayoutForOperand(
     ArrayRef<unsigned> sharedOrder, unsigned vectorSize, unsigned elemBitWidth,
     bool needTrans) const {
   int kDimIndex = operandIdx == 0 ? 1 : 0;
+
+  // Disable swizzling for scales
+  if (operandIdx >= 2) {
+    return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder,
+                                           ctaLayout);
+  }
+
   if (needTrans)
     kDimIndex = 1 - kDimIndex;
 

lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp

@@ -148,6 +148,94 @@ class SinkTMEMLoad : public OpRewritePattern<TMEMTokenLoadOp> {
   }
 };
 
+// Combine back TMEM alloc and store. This is equivalent but gives us a more
+// canonical form to do further optimizations.
+class CombineTMEMStoreAndAlloc : public OpRewritePattern<TMEMTokenStoreOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TMEMTokenStoreOp store,
+                                PatternRewriter &rewriter) const override {
+    if (!matchPattern(store.getPred(), m_One()))
+      return failure();
+    auto alloc = store.getDep().getDefiningOp<TMEMTokenAllocOp>();
+    if (!alloc)
+      return failure();
+    if (alloc->getBlock() != store->getBlock())
+      return failure();
+    alloc.getSrcMutable().assign(store.getSrc());
+    rewriter.replaceOp(store, alloc.getToken());
+    return success();
+  }
+};
+
+// Hoists a tmem alloc outside an if op like this:
+// %0 = scf.if {
+//   %1, %token0 = tmem.alloc %init
+//   ...
+//   %2 = tmem.load %1, %token1
+//   scf.yield %2
+// } else {
+//   scf.yield %init
+// }
+// ->
+// %a, %token0 = tmem.alloc %init
+// %token2 = scf.if {
+//
+//   ...
+//   scf.yield %token1
+// } else {
+//   scf.yield %token0
+// }
+// %2 = tmem.load %a, %token2
+class HoistTMEMAllocOutOfIf : public OpRewritePattern<ttng::TMEMAllocOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ttng::TMEMAllocOp alloc,
+                                PatternRewriter &rewriter) const override {
+    if (!alloc.getToken())
+      return failure();
+    Value init = alloc.getSrc();
+    if (!init)
+      return failure();
+    auto ifOp = dyn_cast<scf::IfOp>(alloc->getParentOp());
+    if (!ifOp)
+      return failure();
+    auto thenOp = ifOp.thenBlock()->getTerminator();
+    auto elseOp = ifOp.elseBlock()->getTerminator();
+    SmallVector<int> yieldArgs;
+    for (auto [thenOperand, elseOperand] :
+         llvm::zip(thenOp->getOpOperands(), elseOp->getOpOperands())) {
+      auto load = thenOperand.get().getDefiningOp<TMEMTokenLoadOp>();
+      if (!load || load.getSrc() != alloc.getResult())
+        continue;
+      if (elseOperand.get() != init)
+        continue;
+      yieldArgs.push_back(thenOperand.getOperandNumber());
+    }
+    if (yieldArgs.empty())
+      return failure();
+    // Since init is used in the else terminator we know that it dominates the
+    // if op.
+    alloc->moveBefore(ifOp);
+    rewriter.setInsertionPointAfter(ifOp);
+    for (int argNo : yieldArgs) {
+      auto load =
+          cast<TMEMTokenLoadOp>(thenOp->getOperand(argNo).getDefiningOp());
+      auto newLoad = cast<TMEMTokenLoadOp>(rewriter.clone(*load));
+      rewriter.modifyOpInPlace(ifOp, [&] {
+        ifOp->getResult(argNo).replaceAllUsesWith(newLoad.getResult());
+        newLoad.getDepMutable().assign(ifOp->getResult(argNo));
+        thenOp->setOperand(argNo, load.getToken());
+        elseOp->setOperand(argNo, alloc.getToken());
+        ifOp->getResult(argNo).setType(newLoad.getToken().getType());
+      });
+    }
+    return success();
+  }
+};
+
 // Remove loop-carried tensor dependencies if they are fed immediately into a
 // TMEM store by pulling the store into the previous iteration.
 class RotateTMEMStoreInLoop : public OpRewritePattern<TMEMTokenStoreOp> {
@@ -412,11 +500,29 @@ struct HoistTMEMAlloc
     mlir::RewritePatternSet patterns(&getContext());
     patterns.add<RotateTMEMStoreInLoop, RotateTMEMLoadInLoop,
                  CombineTMEMLoadAndStore, CombineTMEMStoreAndSelect,
-                 SinkTMEMLoad, RemoveUnusedTMEMLoad>(&getContext());
+                 SinkTMEMLoad, RemoveUnusedTMEMLoad, CombineTMEMStoreAndAlloc,
+                 HoistTMEMAllocOutOfIf>(&getContext());
     scf::ForOp::getCanonicalizationPatterns(patterns, &getContext());
     if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
       llvm_unreachable("Failed to hoist tmem_store");
     }
+
+    // TODO: currently some code assumes that a mutable tmem alloc doesn't have
+    // an initial value. As a workaround we break up the op in order to keep
+    // this form for the downstream passes. We should remove this once the
+    // downstread passes are fixed.
+    m.walk([&](ttng::TMEMAllocOp alloc) {
+      if (alloc.getType().getMutableMemory() && alloc.getSrc()) {
+        OpBuilder builder(alloc);
+        builder.setInsertionPointAfter(alloc);
+        auto store = builder.create<ttng::TMEMStoreOp>(
+            alloc.getLoc(), builder.getType<AsyncTokenType>(),
+            alloc.getResult(), alloc.getToken(), alloc.getSrc(),
+            builder.create<arith::ConstantIntOp>(alloc.getLoc(), 1, 1));
+        alloc.getToken().replaceAllUsesExcept(store.getToken(), store);
+        alloc.getSrcMutable().clear();
+      }
+    });
   }
 };
 

python/src/passes.h

@@ -36,3 +36,8 @@
 #define ADD_PASS_OPTION_WRAPPER_4(name, builder, ty0, ty1, ty2, ty3)           \
   m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2,          \
                  ty3 val3) { pm.addPass(builder({val0, val1, val2, val3})); })
+
+#define ADD_PASS_OPTION_WRAPPER_5(name, builder, ty0, ty1, ty2, ty3, ty4)      \
+  m.def(name,                                                                  \
+        [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2, ty3 val3,      \
+           ty4 val4) { pm.addPass(builder({val0, val1, val2, val3, val4})); })

python/tutorials/gluon/01-attention-forward.py

@@ -173,6 +173,18 @@ def issue_async_tma_load(smem, bar, desc, offset):
     tma.async_copy_global_to_shared(desc, [offset, 0], bar, smem)
 
 
+@gluon.jit
+def _interleave_n(a, b, size: gl.constexpr, f: gl.constexpr, i: gl.constexpr = 0):
+    if a.shape[1] == size:
+        return f(a, b, i)
+    else:
+        a0, a1 = a.reshape([a.shape[0], 2, a.shape[1] // 2]).permute(0, 2, 1).split()
+        b0, b1 = b.reshape([b.shape[0], 2, b.shape[1] // 2]).permute(0, 2, 1).split()
+        c0 = _interleave_n(a0, b0, size, f, i)
+        c1 = _interleave_n(a1, b1, size, f, i + a.shape[1] // 2)
+        return gl.convert_layout(gl.join(c0, c1).permute(0, 2, 1).reshape(a.shape), a.type.layout)
+
+
 # ===-----------------------------------------------------------------------===#
 # Gluon Attention
 # ===-----------------------------------------------------------------------===#
@@ -556,7 +568,7 @@ def _attn_fwd_mma(config, chnls, descs, M, STAGE: gl.constexpr):
         s0_tmem, s0_bar, s0_producer = s0_producer.acquire()
         p0_tmem = _borrow_s_as_p(config, s0_tmem)
         tcgen05_mma(p0_tmem, v_smem, o0_tmem, use_acc=False, mbarriers=[o0_bar])
-        o_init = False
+        o1_init = False
 
         for _ in range(num_mmas - 1):
             k_smem, k_bar, kv_consumer = kv_consumer.acquire()
@@ -565,43 +577,69 @@ def _attn_fwd_mma(config, chnls, descs, M, STAGE: gl.constexpr):
             o1_tmem, o1_bar, o_producer = o_producer.acquire()
             s1_tmem, s1_bar, s1_producer = s1_producer.acquire()
             p1_tmem = _borrow_s_as_p(config, s1_tmem)
-            tcgen05_mma(p1_tmem, v_smem, o1_tmem, use_acc=o_init, mbarriers=[o1_bar, v_bar])
-            o_init = True
+            tcgen05_mma(p1_tmem, v_smem, o1_tmem, use_acc=o1_init, mbarriers=[o1_bar, v_bar])
+            o1_init = True
 
             tcgen05_mma(q1_smem, k_smem.permute((1, 0)), s1_tmem, use_acc=False, mbarriers=[s1_bar, k_bar])
 
             v_smem, v_bar, kv_consumer = kv_consumer.acquire()
             o0_tmem, o0_bar, o_producer = o_producer.acquire()
             s0_tmem, s0_bar, s0_producer = s0_producer.acquire()
             p0_tmem = _borrow_s_as_p(config, s0_tmem)
-            tcgen05_mma(p0_tmem, v_smem, o0_tmem, use_acc=o_init, mbarriers=[o0_bar])
-            o_init = True
+            tcgen05_mma(p0_tmem, v_smem, o0_tmem, mbarriers=[o0_bar])
 
         tcgen05_commit(q0_bar)
         tcgen05_commit(q1_bar)
 
         o1_tmem, o1_bar, o_producer = o_producer.acquire()
         s1_tmem, s1_bar, s1_producer = s1_producer.acquire()
         p1_tmem = _borrow_s_as_p(config, s1_tmem)
-        tcgen05_mma(p1_tmem, v_smem, o1_tmem, use_acc=o_init, mbarriers=[o1_bar, v_bar, s0_bar, s1_bar])
+        tcgen05_mma(p1_tmem, v_smem, o1_tmem, use_acc=o1_init, mbarriers=[o1_bar, v_bar, s0_bar, s1_bar])
+
+
+@gluon.jit
+def _mask_inner(qk, mask, i: gl.constexpr):
+    mask_i_bit = mask & (1 << i) == 0
+    return gl.where(mask_i_bit, qk, -float("inf"))
+
+
+@gluon.jit
+def _mask_frag(qk, col_limit_right, s: gl.constexpr):
+    col_limit_right_s = col_limit_right - s
+    col_limit_right_cur = max(col_limit_right_s, 0)
+    mask = -1 << col_limit_right_cur
+    return _interleave_n(qk, mask, 1, _mask_inner)
+
+
+@gluon.jit
+def _mask_bits(qk, col_limit_right):
+    # FIXME: This is a more concise implementation (which compiles faster) but
+    # it results in slightly slower code due to the lack of interleaving.
+    offs_n = gl.arange(0, qk.shape[1], layout=gl.SliceLayout(0, qk.type.layout))[None, :]
+    s = offs_n & ~0xf
+    i = offs_n & 0xf
+
+    col_lim_right_s = col_limit_right - s
+    col_lim_right_cur = max(col_lim_right_s, 0)
+    mask = -1 << col_lim_right_cur
+    mask_i_bit = (mask & (1 << i)) == 0
+    return gl.where(mask_i_bit, qk, -float("inf"))
 
 
 @gluon.jit
 def _softmax_inner_loop(tile_id: gl.constexpr, config, prog,  #
                         s_consumer, corr_producer, exp_turnstile, corr_bar,  #
-                        offs_m, offs_n, m_i, l_i0, l_i1, STAGE: gl.constexpr):
+                        offs_m, m_i, l_i0, l_i1, STAGE: gl.constexpr):
     lo, hi = prog.get_loop_bounds(STAGE)
 
     for start_n in range(lo, hi, config.BLOCK_N):
         s_tmem, s_bar, s_consumer = s_consumer.acquire()
         qk = s_tmem.load(config.qk_layout)
 
         if STAGE == 2:
-            # Prevent LLVM from hoisting the partial sums, which triggers spilling.
-            offs_n = gl.inline_asm_elementwise("mov.b32 $0, $0;", "=r,r", [offs_n], dtype=gl.int32, is_pure=True,
-                                               pack=1)
-            mask = offs_m[:, None] < (start_n + offs_n[None, :])
-            qk = gl.where(mask, -1.0e8, qk)
+            col_limit_right = (offs_m - start_n + 1)[:, None].broadcast_to(qk.shape)
+            qk = _interleave_n(qk, col_limit_right, 16, _mask_frag)
+
         m_ij = gl.maximum(m_i, gl.max(qk, 1) * config.qk_scale)
         alpha = gl.exp2(m_i - m_ij)
 
@@ -682,11 +720,8 @@ def _softmax_inner_loop(tile_id: gl.constexpr, config, prog,  #
 @gluon.jit
 def _softmax_tile(tile_id: gl.constexpr, config, M, desc_o, STAGE: gl.constexpr,  #
                   s_chnl, corr_chnl, exp_turnstile):
-    qk_slice_dim0: gl.constexpr = gl.SliceLayout(0, config.qk_layout)
     qk_slice_dim1: gl.constexpr = gl.SliceLayout(1, config.qk_layout)
 
-    offs_n = gl.arange(0, config.BLOCK_N, qk_slice_dim0)
-
     s_consumer = s_chnl.create_consumer()
     corr_producer = corr_chnl.create_producer()
     _, corr_bar, corr_producer = corr_producer.acquire()
@@ -709,11 +744,11 @@ def _softmax_tile(tile_id: gl.constexpr, config, M, desc_o, STAGE: gl.constexpr,
         if STAGE & 1:
             m_i, l_i0, l_i1, corr_bar, s_consumer, corr_producer, exp_turnstile = _softmax_inner_loop(  #
                 tile_id, config, prog, s_consumer, corr_producer, exp_turnstile, corr_bar,  #
-                offs_m, offs_n, m_i, l_i0, l_i1, STAGE=4 - STAGE)
+                offs_m, m_i, l_i0, l_i1, STAGE=4 - STAGE)
         if STAGE & 2:
             m_i, l_i0, l_i1, corr_bar, s_consumer, corr_producer, exp_turnstile = _softmax_inner_loop(  #
                 tile_id, config, prog, s_consumer, corr_producer, exp_turnstile, corr_bar,  #
-                offs_m, offs_n, m_i, l_i0, l_i1, STAGE=2)
+                offs_m, m_i, l_i0, l_i1, STAGE=2)
 
         if config.use_fadd2_reduce:
             l_i = l_i0 + l_i1