[WS] fold rewrite-partition-dependencies into insert-aref (#8619)

folder `rewrite-partition-dependencies` pass into `insert-aref`. fold
corresponding lit test into `insert-aref` lit tests

- `insert-aref` now handles all aref insertio that follows
producer-consumer model
- it also follow `rewrite-partition-dependencies` logic where
`get.enter` is inserted just before consumer in the same block as
producer is located, instead right after producer as it was before.
- insert-tmem-aref handles tmem aref, that follow ownership-transfer
model.

Author: 3gx

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

@@ -126,6 +126,7 @@ void setPartition(Operation *op, const SetVector<int> &partitionIds);
 void setPartitionOutputs(Operation *op,
                          ArrayRef<SetVector<int>> partitionOutputsIds);
 SmallVector<SetVector<int>, 4> getPartitionOutputs(Operation *op);
+SetVector<int> getPartitionIds(OpOperand *use);
 
 } // namespace mlir::triton::gpu
 

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

@@ -121,25 +121,6 @@ def TritonGPUAutomaticWarpSpecialization : Pass<"tritongpu-automatic-warp-specia
   ];
 }
 
-def TritonGPURewritePartitionDependencies : Pass<"tritongpu-rewrite-partition-dependencies", "mlir::ModuleOp"> {
-  let summary = "test pass for rewriting partition dependencies";
-
-  let description = [{
-    The `tritongpu-rewrite-partition-dependencies` pass analyzes the partitions
-    assigned to a loop and their SSA dependencies. It rewrites the dependencies
-    to be passed through shared memory, applying multi-buffering according to
-    the assigned stages of the partitions.
-  }];
-
-  let dependentDialects = [
-    "mlir::triton::gpu::TritonGPUDialect",
-    "mlir::scf::SCFDialect",
-    "mlir::arith::ArithDialect",
-    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
-    "mlir::triton::nvws::NVWSDialect"
-  ];
-}
-
 def TritonGPUPartitionLoops : Pass<"tritongpu-partition-loops", "mlir::ModuleOp"> {
   let summary = "split scheduled loops into `ttg.warp_specialize`";
 

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -34,7 +34,6 @@ add_triton_library(TritonGPUTransforms
   WarpSpecialization/PartitionBuilder.cpp
   WarpSpecialization/PartitionLoops.cpp
   WarpSpecialization/PartitionScheduling.cpp
-  WarpSpecialization/RewritePartitionDependencies.cpp
 
   DEPENDS
   TritonGPUTransformsIncGen

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

@@ -37,7 +37,6 @@ void AutomaticWarpSpecialization::runOnOperation() {
   pm.addPass(createTritonGPUPartitionScheduling());
   pm.addPass(createNVWSInsertAref());
   pm.addPass(createNVWSInsertTmemAref());
-  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.
   // pm.addPass(arith::createIntRangeOptimizationsPass());

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

@@ -249,6 +249,18 @@ SmallVector<SetVector<int>, 4> getPartitionOutputs(Operation *op) {
   return partitionOutputsIds;
 }
 
+SetVector<int> getPartitionIds(OpOperand *use) {
+  auto owner = use->getOwner();
+  if (isa<scf::YieldOp>(owner)) {
+    return getPartitionOutputs(owner->getParentOp())[use->getOperandNumber()];
+  } else if (scf::ForOp forOp = dyn_cast<scf::ForOp>(owner)) {
+    int idx = use->getOperandNumber() - forOp.getNumControlOperands();
+    return idx >= 0 ? getPartitionOutputs(owner)[idx] : *getPartitionIds(forOp);
+  } else {
+    return *getPartitionIds(owner);
+  }
+}
+
 bool hasPartition(Operation *op) { return getPartitionIds(op) != std::nullopt; }
 
 } // namespace mlir::triton::gpu