[Pipeliner] Implement more sophisticated partitioning strategy for attention (#6660)

This PR implements a more sophisticated partitioning and scheduling
strategy, aimed at generating a better schedule for forward attention.

This PR enables composing the pipeliner with warp specialization by
tweaking `scheduleLoops` to consider a `ttg.assigned_stage` attribute.
Normally, the pipeline scheduler determines stages based on a DAG
between latency ops, but this allows some higher level pass to directly
inject an assigned stage into `scheduleKeyOps`. This allows, for
example, warp specialization to place `load K` into stage 0 and `load V`
into stage 2, and `MMA QK` into stage 0 and `MMA PV` into stage 2, in
their respective partitions.

This PR also reworks partition assignment to generate more than 1 vector
partition. Previously, the partitioning strategy was generalized to take
the body of the loop and "outline" async loads and async MMAs into their
own partition, keeping the rest of the loop body in the default
partition.

Now, partition assignment also considers high latency synchronous
operations, starting with `math.exp2` with a large number of elements.
It does a first-order partition assignment, placing local users and
dependencies of loads and MMAs into the same partition, and then
clusters the remaining operations together. The clusters are then
either:

1. Assigned entirely to the source partition
2. Assigned entirely to the sink partition
3. Assigned into a wholly new partition 
4. Rematerialized into sink partitions along critical paths

The idea is simply to simultaneously reduce the critical path between
each latency operation. This strategy successfully automatically derives
the same schedule CUTLASS uses for FMHA (correction partition, etc.).
There is currently no cost model for deciding between rematerialization
or sending intermediates over shared memory/placing them in their own
partition, but we can probably reuse the one @apgoucher implemented for
`remove-layout-conversions` in
triton-lang/triton#6667.

This also slightly tweaks the kernel code in `06-fused-attention.py` to
reduce register pressure. This achieves close to 700 TFLOPS on DHEAD=64
and around 960-1080 TFLOPS on DHEAD=128.

TODO:
- [ ] Write lit tests for new scheduler
- [ ] Write integration tests for MFHA

---------

Co-authored-by: Chris Sullivan <[email protected]>

Author: Mogball

include/triton/Dialect/TritonGPU/Transforms/Partition.h

@@ -61,8 +61,6 @@ class WarpSchedule {
 
   // Create a new partition with a stage.
   Partition *addPartition(unsigned stage);
-  // Give each partition a new index and order. The indices must be unique.
-  void reorderPartitions(ArrayRef<unsigned> order);
   // Update the op to partition mapping.
   void updatePartitions();
 
@@ -74,15 +72,25 @@ class WarpSchedule {
   Partition *getPartition(unsigned idx);
   // Get the partition at the index.
   const Partition *getPartition(unsigned idx) const;
+  // Insert an operation into a partition.
+  void insert(Partition *partition, Operation *op);
   // Return an iterator range over the partitions.
   auto getPartitions() { return llvm::make_pointee_range(partitions); }
   // Return an iterator range over the partitions.
   auto getPartitions() const { return llvm::make_pointee_range(partitions); }
+  // Get the number of partitions.
+  unsigned getNumPartitions() const { return partitions.size(); }
   // Get the root partition.
   Partition *getRootPartition() { return rootPartition.get(); }
   // Get the root partition.
   const Partition *getRootPartition() const { return rootPartition.get(); }
 
+  // Return true if an operation is assigned to a partition.
+  bool isScheduled(Operation *op) const;
+  // Schedule an operation to a partition if it is not already scheduled. Return
+  // true if the operation was scheduled.
+  bool trySchedule(Partition *partition, Operation *op);
+
   // Deserialize a warp schedule from an `scf.for` op using the attributes
   // tagged on operations in its body.
   static FailureOr<WarpSchedule> deserialize(scf::ForOp loop);

include/triton/Dialect/TritonGPU/Transforms/PipeliningUtility.h

@@ -19,6 +19,7 @@ static const char *kWarpSpecializeAttrName = "tt.warp_specialize";
 static const char *kLoopStageAttrName = "loop.stage";
 static const char *kLoopClusterAttrName = "loop.cluster";
 static const char *kScheduledMaxStageAttrName = "tt.scheduled_max_stage";
+static const char *kAssignedStageAttrName = "ttg.assigned_stage";
 
 //===----------------------------------------------------------------------===//
 // Hoisting Utilities

lib/Analysis/Membar.cpp

@@ -1,6 +1,7 @@
 #include "triton/Analysis/Membar.h"
 #include "triton/Analysis/Alias.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
@@ -216,6 +217,13 @@ void MembarAnalysis::update(Operation *op, BlockInfo *blockInfo,
         }
       }
     }
+    // If this op is may be signalling other threads asynchronously, make sure
+    // all shared memory transactions are complete beforehand.
+    if (isa<triton::nvidia_gpu::ArriveBarrierOp>(op)) {
+      Interval<size_t> allIntervals(0, std::numeric_limits<size_t>::max());
+      curBlockInfo.syncWriteIntervals[allIntervals].insert(op);
+      curBlockInfo.syncReadIntervals[allIntervals].insert(op);
+    }
     scratchBufferId = allocation->getBufferId(op);
   }
 

lib/Conversion/TritonToTritonGPU/RelayoutTritonGPU.cpp

@@ -44,6 +44,11 @@ struct TMEMLoadOpPattern : public OpConversionPattern<ttng::TMEMLoadOp> {
     RankedTensorType type = getTMEMTensorLayout(
         typeConverter, op.getType(), op.getSrc().getType(), lookupNumWarps(op));
     rewriter.modifyOpInPlace(op, [&] { op.getResult().setType(type); });
+    Type resultType = getTypeConverter()->convertType(op.getType());
+    rewriter.setInsertionPointAfter(op);
+    auto cvt = rewriter.create<ConvertLayoutOp>(op.getLoc(), resultType,
+                                                op.getResult());
+    rewriter.replaceAllUsesExcept(op.getResult(), cvt, cvt);
     return success();
   }
 };

lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp

@@ -189,9 +189,13 @@ void tt::CoarseSchedule::serialize(scf::ForOp &forOp) {
   for (auto [op, stage, cluster] : getOpsInOrder(forOp)) {
     setStageCluster(op, stage, *cluster);
   }
+
+  Builder b(forOp.getContext());
+  int maxStages = numStages - 1;
+  if (auto maxStageAttr = tryGetMaxStage(forOp))
+    maxStages = std::max(maxStages, *maxStageAttr);
   forOp->setAttr(mlir::triton::kScheduledMaxStageAttrName,
-                 IntegerAttr::get(IntegerType::get(forOp.getContext(), 32),
-                                  numStages - 1));
+                 b.getI32IntegerAttr(maxStages));
 }
 
 // Create a CoarseSchedule based on forOp's <stage, cluster>.

lib/Dialect/TritonGPU/Transforms/Pipeliner/ScheduleLoops.cpp

@@ -59,6 +59,8 @@ bool hasLatenciesAssigned(scf::ForOp forOp,
   for (auto &op : forOp.getBody()->without_terminator()) {
     if (opLatency.count(&op))
       return true;
+    if (op.getAttr(kAssignedStageAttrName))
+      return true;
   }
   return false;
 }
@@ -70,12 +72,15 @@ CoarseSchedule scheduleKeyOps(scf::ForOp forOp,
   auto terminator = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
   // Determine all operations that have a non-zero latency
   SmallVector<Operation *> latOps;
+  SmallVector<Operation *> stagedOps;
   for (auto &op : forOp.getBody()->without_terminator()) {
     if (opLatency.count(&op))
       latOps.push_back(&op);
+    if (op.getAttr(kAssignedStageAttrName))
+      stagedOps.push_back(&op);
   }
   // If no latency ops, nothing to schedule
-  if (latOps.empty())
+  if (latOps.empty() && stagedOps.empty())
     return CoarseSchedule(0);
 
   DominanceInfo domInfo(forOp);
@@ -123,6 +128,11 @@ CoarseSchedule scheduleKeyOps(scf::ForOp forOp,
       opToStage[op] = maxDistance - dist;
   }
 
+  for (Operation *op : stagedOps) {
+    auto stageAttr = op->getAttrOfType<IntegerAttr>(kAssignedStageAttrName);
+    opToStage[op] = stageAttr.getInt();
+  }
+
   auto stages = llvm::make_second_range(opToStage);
   int maxStage = *llvm::max_element(stages);
   CoarseSchedule schedule(maxStage + 1);

lib/Dialect/TritonGPU/Transforms/Pipeliner/SoftwarePipeliner.cpp

@@ -73,6 +73,7 @@ static void removeAttributes(ModuleOp moduleOp) {
     op->removeAttr(mlir::triton::kLoopStageAttrName);
     op->removeAttr(mlir::triton::kLoopClusterAttrName);
     op->removeAttr(mlir::triton::kScheduledMaxStageAttrName);
+    op->removeAttr(mlir::triton::kAssignedStageAttrName);
   });
 }
 

lib/Dialect/TritonGPU/Transforms/Utility.cpp

@@ -1402,30 +1402,35 @@ void replaceUsesAndPropagateType(OpBuilder &builder, Operation *oldUse,
   // Save the operand to replace / delete later (avoid iterator invalidation).
   // TODO: can we use an early_inc iterator?
   for (OpOperand &use : oldUse->getUses()) {
+    // Propagate through `ttg.warp_specialize`.
+    if (auto wsOp = dyn_cast<ttg::WarpSpecializeOp>(use.getOwner())) {
+      for (Region *region : wsOp.getPartitionRegions())
+        region->getArgument(use.getOperandNumber()).setType(val.getType());
+    }
+
     // Non-subview/trans ops will be replaced by `val`.
-    if (!isa<triton::gpu::MemDescTransOp, triton::gpu::MemDescSubviewOp>(
-            use.getOwner())) {
+    if (!isa<ttg::MemDescTransOp, ttg::MemDescSubviewOp>(use.getOwner())) {
       operandsToReplace.push_back(&use);
       continue;
     }
+
     Operation *user = use.getOwner();
     // `subview(old_op)` is replaced by a new `subview(val)`.
     OpBuilder::InsertionGuard g(builder);
     builder.setInsertionPoint(user);
     Value newVal;
-    if (auto subview = dyn_cast<triton::gpu::MemDescSubviewOp>(user)) {
-      triton::gpu::MemDescType oldType = subview.getType();
-      bool isMutable =
-          cast<triton::gpu::MemDescType>(val.getType()).getMutableMemory();
-      Type newDstType = triton::gpu::MemDescType::get(
+    if (auto subview = dyn_cast<ttg::MemDescSubviewOp>(user)) {
+      ttg::MemDescType oldType = subview.getType();
+      bool isMutable = cast<ttg::MemDescType>(val.getType()).getMutableMemory();
+      Type newDstType = ttg::MemDescType::get(
           oldType.getShape(), oldType.getElementType(), oldType.getEncoding(),
           oldType.getMemorySpace(), isMutable);
-      newVal = builder.create<triton::gpu::MemDescSubviewOp>(
+      newVal = builder.create<ttg::MemDescSubviewOp>(
           subview.getLoc(), newDstType, val, subview.getOffsets());
       newVal.getDefiningOp()->setAttrs(user->getAttrs());
-    } else if (auto trans = dyn_cast<triton::gpu::MemDescTransOp>(user)) {
-      newVal = builder.create<triton::gpu::MemDescTransOp>(trans.getLoc(), val,
-                                                           trans.getOrder());
+    } else if (auto trans = dyn_cast<ttg::MemDescTransOp>(user)) {
+      newVal = builder.create<ttg::MemDescTransOp>(trans.getLoc(), val,
+                                                   trans.getOrder());
       newVal.getDefiningOp()->setAttrs(user->getAttrs());
     }
     assert(newVal);