[WS]: partition-scheduler annotates all ops + fixes (#8215)

* all ops, except parent ops (if/reduce), are annotated with partition
sets after `partition-scheduler`
* patch up downstream passes to handle this properly
* partition set of a parent op is inferred in `partition-loop` pass;
also infers partitions set for each result
* extend `partition-loop` to partition if-stmt across different
partition when loops are split
* patch `lower-mma-specialization` to annotate all ops, and patch unit
test
* patch up remainig passes to play well with requirement that all ops
must be annotated in ws-loops
* there are few remaining issues before `aref-tmem-insertion` can be
enabled and `load-mma-specialzation` disabled, which is subject of
subsequent PR

Author: 3gx

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

@@ -17,6 +17,7 @@ class ForOp;
 } // namespace mlir
 
 static constexpr char kPartitionAttrName[] = "ttg.partition";
+static constexpr char kPartitionOutputsAttrName[] = "ttg.partition.outputs";
 static constexpr char kPartitionStagesAttrName[] = "ttg.partition.stages";
 static constexpr char kWarpSpecializeTagAttrName[] = "ttg.warp_specialize.tag";
 

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

@@ -2,34 +2,47 @@
 #define TRITON_TRITONGPU_TRANSFORMS_PARTITIONBUILDER_H
 
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "llvm/ADT/SetVector.h"
 
 namespace mlir::triton::gpu {
 
 class Partition;
 
 using StageCluster = std::optional<std::pair<int, int>>;
 
+// Get the stage and cluster for an operation, if it has one assigned.
+void setStageCluster(OpBuilder &b, Operation *op, StageCluster stageCluster);
+StageCluster getStageCluster(Operation *op);
+
 struct PartitionBuilder : public ImplicitLocOpBuilder {
   using ImplicitLocOpBuilder::ImplicitLocOpBuilder;
 
   Value intCst(int value, unsigned width = 32);
   Value boolCst(bool value);
 
-  void assignStage(Operation *op, StageCluster stageCluster);
   void assignPartition(Operation *op, Partition &partition);
 
   template <typename OpT, typename... Args>
   auto createInto(Partition &partition, StageCluster stageCluster,
                   Args &&...args) {
     auto op = create<OpT>(std::forward<Args>(args)...);
     assignPartition(op, partition);
-    assignStage(op, stageCluster);
+    setStageCluster(*this, op, stageCluster);
     return op;
   }
 };
 
-// Get the stage and cluster for an operation, if it has one assigned.
-StageCluster getStageCluster(Operation *op);
+template <typename OpT, typename... Args>
+OpT createInto(OpBuilder &b, Location loc,
+               std::optional<SetVector<int>> partitionSet,
+               StageCluster stageCluster, Args &&...args) {
+  auto op = b.create<OpT>(loc, std::forward<Args>(args)...);
+  if (partitionSet) {
+    setPartition(op, *partitionSet);
+    setStageCluster(b, op, stageCluster);
+  }
+  return op;
+}
 
 } // namespace mlir::triton::gpu
 

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

@@ -36,7 +36,11 @@ void AutomaticWarpSpecialization::runOnOperation() {
   OpPassManager pm;
   pm.addPass(createTritonGPUPartitionScheduling());
   pm.addPass(createNVWSInsertAref());
+#if 0
+  pm.addPass(createNVWSInsertTmemAref());
+#else
   pm.addPass(createTritonGPULoadMMASpecialization({numStages}));
+#endif
   pm.addPass(createTritonGPURewritePartitionDependencies());
   // `int-range-optimizations` and SCCP are good at cleaning up loop arithmetic.
   // FIXME: Re-enable integer range analysis once it is fixed.

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

@@ -573,7 +573,7 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
   for (int i = 0; i < nodes.size(); ++i) {
     Node &cur = nodes[i];
     Node &next = nodes[(i + 1) % nodes.size()];
-    if (!samePartition(inBody(cur.op), inBody(next.op))) {
+    if (!samePartition(cur.op, next.op)) {
       cur.barNext = createBarrierAlloc(loop, numMmaStages);
       next.barPrev = cur.barNext;
     }
@@ -616,6 +616,7 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
       continue;
     b.setInsertionPoint(node.op);
     Value view = createSingleBufferView(b, allocOp, node.index);
+    b.assignPartition(view.getDefiningOp(), *partitions.getPartition(node.op));
     if (auto storeOp = dyn_cast<ttng::TMEMStoreOp>(node.op)) {
       storeOp.getDstMutable().assign(view);
       storeOp.getDepMutable().clear();
@@ -671,7 +672,6 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
     Operation *defOp = operand.getDefiningOp();
     if (!defOp || loop.isDefinedOutsideOfLoop(operand))
       continue;
-    defOp = inBody(defOp);
 
     if (partitions.isInRootPartition(defOp)) {
       // If the MMA operand is coming from outside the loop, move the alloc out.
@@ -717,7 +717,7 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
   }
 
   for (Node &node : nodes) {
-    Partition *partition = partitions.getPartition(inBody(node.op));
+    Partition *partition = partitions.getPartition(node.op);
     PartitionBuilder b(node.op->getLoc(), loop);
 
     SmallVector<Operation *> defs;
@@ -746,11 +746,13 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
             domInfo.properlyDominates(mmaOp, userPred.getDefiningOp())) {
           b.restoreInsertionPoint(*incrementPt);
           Value bar = createSingleBufferView(b, node.barPrev, curIndex);
+          b.assignPartition(bar.getDefiningOp(), *partition);
           b.createInto<ttng::WaitBarrierOp>(*partition, nodeStageCluster, bar,
                                             curPhase, userPred);
         } else {
           b.setInsertionPoint(domOp);
           Value bar = createSingleBufferView(b, node.barPrev, node.index);
+          b.assignPartition(bar.getDefiningOp(), *partition);
           b.createInto<ttng::WaitBarrierOp>(*partition, nodeStageCluster, bar,
                                             node.phase, userPred);
         }
@@ -759,6 +761,7 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
         if (isa<scf::IfOp>(domOp->getParentOp()) && accIsMultiBuffered)
           b.setInsertionPointToStart(domOp->getBlock());
         Value bar = createSingleBufferView(b, node.barPrev, node.index);
+        b.assignPartition(bar.getDefiningOp(), *partition);
         b.createInto<ttng::WaitBarrierOp>(*partition, nodeStageCluster, bar,
                                           node.phase);
       }
@@ -767,13 +770,15 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
       if (mmaOp == node.op) {
         b.setInsertionPoint(mmaOp);
         Value bar = createSingleBufferView(b, node.barNext, node.index);
+        b.assignPartition(bar.getDefiningOp(), *partitions.getPartition(mmaOp));
         mmaOp.addCompletionBarrier(bar, userPred);
         mmaOp.setIsAsync(true);
       } else {
         b.setInsertionPointAfter(lastOp);
         if (isa<scf::IfOp>(lastOp->getParentOp()) && accIsMultiBuffered)
           b.setInsertionPoint(lastOp->getBlock()->getTerminator());
         Value bar = createSingleBufferView(b, node.barNext, node.index);
+        b.assignPartition(bar.getDefiningOp(), *partition);
         b.createInto<ttng::ArriveBarrierOp>(*partition, nodeStageCluster, bar,
                                             1);
       }
@@ -799,20 +804,26 @@ static LogicalResult pipelineMMA(scf::ForOp &loop, PipelinedMMA &mma,
     StageCluster srcStageCluster = getStageCluster(domOp);
     b.setInsertionPoint(domOp);
     Value emptyView = createSingleBufferView(b, emptyBar, index);
+    b.assignPartition(emptyView.getDefiningOp(), *partition);
     b.createInto<ttng::WaitBarrierOp>(*partition, srcStageCluster, emptyView,
                                       phase);
 
     b.setInsertionPointAfter(lastOp);
     Value readyView = createSingleBufferView(b, readyBar, index);
+    b.assignPartition(readyView.getDefiningOp(), *partition);
     b.createInto<ttng::ArriveBarrierOp>(*partition, srcStageCluster, readyView,
                                         1);
 
     b.setInsertionPoint(mmaOp);
     Value readyView2 = createSingleBufferView(b, readyBar, index);
+    b.assignPartition(readyView2.getDefiningOp(),
+                      *partitions.getPartition(mmaOp));
     b.createInto<ttng::WaitBarrierOp>(*partitions.getPartition(mmaOp),
                                       getStageCluster(mmaOp), readyView2,
                                       phase);
     Value emptyView2 = createSingleBufferView(b, emptyBar, index);
+    b.assignPartition(emptyView2.getDefiningOp(),
+                      *partitions.getPartition(mmaOp));
     mmaOp.addCompletionBarrier(emptyView2, b.boolCst(true));
     mmaOp.setIsAsync(true);
   }

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

@@ -14,13 +14,6 @@ Value PartitionBuilder::boolCst(bool value) {
   return intCst(value, /*width=*/1);
 }
 
-void PartitionBuilder::assignStage(Operation *op, StageCluster stageCluster) {
-  if (stageCluster) {
-    op->setAttr(kLoopStageAttrName, getI32IntegerAttr(stageCluster->first));
-    op->setAttr(kLoopClusterAttrName, getI32IntegerAttr(stageCluster->second));
-  }
-}
-
 void PartitionBuilder::assignPartition(Operation *op, Partition &partition) {
   setPartition(op, &partition);
 }
@@ -32,3 +25,12 @@ StageCluster triton::gpu::getStageCluster(Operation *op) {
     return std::nullopt;
   return std::make_pair(stageAttr.getInt(), clusterAttr.getInt());
 }
+
+void triton::gpu::setStageCluster(OpBuilder &b, Operation *op,
+                                  StageCluster stageCluster) {
+  if (stageCluster) {
+    op->setAttr(kLoopStageAttrName, b.getI32IntegerAttr(stageCluster->first));
+    op->setAttr(kLoopClusterAttrName,
+                b.getI32IntegerAttr(stageCluster->second));
+  }
+}