[SWP] Attempt to move all scheduling logic to a scheduling pass (triton-lang#4618)

The main purpose is to move all scheduling related logic including
scheduleDeps/scheduleDistOne/scheduleRemaining to a new loopScheduling
pass. The new pass will generate (stage, cluster) attributes for each
operation inside a loop.

For operations that are created during createAsyncOps, we manually add
the (stage, cluster) based on the attributes prior to lowering. In most
cases, the lowered operations should have (stage, cluster) of the
original loadOp, or should have (stage, cluster) of the first use of the
loadOp.

This patch also gets rid of prefetchCluster, instead it prefetches one
stage before the actual use
setStageCluster(forOp, wait, stageForFirstUse - 1, clusterForFirstUse +
1);
setStageCluster(forOp, viewLoad, stageForFirstUse - 1,
clusterForFirstUse + 1);
comparing to
      schedule.insert(wait, numStages - 2, prefetchCluster);
      schedule.insert(viewLoad, numStages - 2, prefetchCluster);

At end of createAsyncOps, we make sure all operations inside the loop
have (stage, cluster) attributes.

There is no need to maintain CoarseSchedule in SWP, instead we will just
use (stage, cluster) attributes.

Author: manman-ren

include/triton/Dialect/TritonGPU/Transforms/Passes.td

@@ -179,4 +179,17 @@ def TritonGPUOptimizeAccumulatorInit: Pass<"tritongpu-optimize-accumulator-init"
                            "mlir::triton::TritonDialect"];
 }
 
+def TritonGPULoopScheduling: Pass<"tritongpu-loop-scheduling", "mlir::ModuleOp"> {
+  let summary = "Generate loop scheduling for SWP";
+
+  let description = "This pass sets up stages and clustering for software pipelining.";
+
+  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
+                           "mlir::triton::TritonDialect"];
+  let options = [
+    Option<"numStages", "num-stages",
+           "int32_t", /*default*/"3",
+           "number of pipeline stages">
+  ];
+}
 #endif

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

@@ -8,6 +8,8 @@ namespace mlir {
 namespace triton {
 
 static const char *kNumStagesAttrName = "tt.num_stages";
+static const char *kLoopStageAttrName = "loop.stage";
+static const char *kLoopClusterAttrName = "loop.cluster";
 
 /// Function to mask operations during scheduling.
 Operation *predicateOp(RewriterBase &rewriter, Operation *op, Value pred);
@@ -29,6 +31,11 @@ void addOps(scf::ForOp forOp, int stage,
 /// mutable.
 void replaceUsesAndPropagateType(OpBuilder &builder, Operation *oldUse,
                                  Value val);
+
+// Return the minClusterId and maxClusterId for the given ForOp.
+std::pair<int, int> getMinMaxCluster(scf::ForOp &forOp);
+std::pair<int, int> getStageCluster(Operation *op);
+void setStageCluster(scf::ForOp &forOp, Operation *op, int stage, int cluster);
 } // namespace triton
 } // namespace mlir
 

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

@@ -101,8 +101,16 @@ class CoarseSchedule {
   createFinalSchedule(scf::ForOp forOp);
   void dump();
   bool empty() { return opToStageAndCluster.size() == 0; }
+  void serialize(scf::ForOp &forOp);
+  // Create a CoarseSchedule based on forOp's <stage, cluster>.
+  void deSerialize(scf::ForOp &forOp);
 };
 
+// Add dependencies of anchor ops to the coarse schedule. Schedule them to
+// the same stage and ordering cluster as the anchor op.
+void scheduleDependencies(scf::ForOp forOp, CoarseSchedule &schedule,
+                          int numStages);
+
 } // namespace triton
 } // namespace mlir
 #endif // TRITON_TRITONGPU_TRANSFORM_PIPELINE_SCHEDULE_H_

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

@@ -192,6 +192,10 @@ bool isPureUnaryInlineAsm(Operation *op);
 // read the compute capability from the module attributes
 int getNVIDIAComputeCapability(Operation *module);
 
+std::optional<mlir::triton::gpu::SharedEncodingAttr>
+getSharedEncIfAllUsersAreDotEnc(Value val, bool &incompatible);
+
+bool loadIsMMAv3(Operation *loadOp);
 } // namespace mlir
 
 #endif // TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -3,6 +3,7 @@ add_triton_library(TritonGPUTransforms
   Coalesce.cpp
   F32DotTC.cpp
   CombineTensorSelectAndIf.cpp
+  LoopScheduling.cpp
   ReduceDataDuplication.cpp
   OptimizeAccumulatorInit.cpp
   OptimizeDotOperands.cpp