Revert "[WS] Update RewritePartitionDepdencies to insert arefs" (#7611)

Reverts triton-lang/triton#7561 as it causes functional failures in our
internal attention tests

Author: ThomasRaoux

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

@@ -130,8 +130,7 @@ def TritonGPURewritePartitionDependencies : Pass<"tritongpu-rewrite-partition-de
     "mlir::triton::gpu::TritonGPUDialect",
     "mlir::scf::SCFDialect",
     "mlir::arith::ArithDialect",
-    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
-    "mlir::triton::nvws::NVWSDialect"
+    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
   ];
 }
 

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

@@ -3,7 +3,6 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
-#include "third_party/nvidia/include/Dialect/NVWS/Transforms/Passes.h"
 #include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 
@@ -43,8 +42,6 @@ void AutomaticWarpSpecialization::runOnOperation() {
   pm.addPass(createSCCPPass());
   pm.addPass(createCSEPass());
   pm.addPass(createTritonGPUPartitionLoops());
-  pm.addPass(createNVWSLowerAref());
-  pm.addPass(createNVWSLowerWarpGroup());
   if (failed(runPipeline(pm, getOperation())))
     return signalPassFailure();
 

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

@@ -191,6 +191,57 @@ LogicalResult WarpSchedule::verify(scf::ForOp loop) const {
   if (failed)
     return failure();
 
+  // Within a loop iteration, the partitions must form a DAG. For example, the
+  // following is invalid:
+  //
+  //   scf.for %i = %lb to %ub step %step
+  //     %0 = op_a()     {ttg.partition = 0}
+  //     %1 = op_b(%0)   {ttg.partition = 1}
+  //     op_c(%1)        {ttg.partition = 0}
+  //
+  PartitionGraph graph(loop, *this);
+  for (auto it = llvm::scc_begin(graph); !it.isAtEnd(); ++it) {
+    if (!it.hasCycle())
+      continue;
+    InFlightDiagnostic diag =
+        mlir::emitWarning(loop.getLoc(), "warp schedule contains a cycle");
+    for (auto [node, use] : *it) {
+      assert(use && "already checked that the root partition has no ancestors");
+      diag.attachNote(use->getOwner()->getLoc())
+          << "operation in partition #" << node->partition->getIndex()
+          << " uses value defined in partition #"
+          << opToPartition.at(use->get().getDefiningOp())->getIndex();
+    }
+    return failure();
+  }
+
+  // Each partition's stage must be strictly less than all of its consumers plus
+  // the distance.
+  for (Partition &partition : getPartitions()) {
+    bool failed = false;
+    auto callback = [&](OpResult output, OpOperand &use, unsigned distance) {
+      Operation *user = loop.getBody()->findAncestorOpInBlock(*use.getOwner());
+      const Partition *consumer = opToPartition.at(user);
+      if (partition.getStage() < consumer->getStage() + distance)
+        return;
+      InFlightDiagnostic diag =
+          mlir::emitWarning(loop.getLoc(), "partition #")
+          << partition.getIndex() << " has stage " << partition.getStage()
+          << " but is consumed by partition #" << consumer->getIndex()
+          << " with stage " << consumer->getStage() << " at distance "
+          << distance;
+      diag.attachNote(use.getOwner()->getLoc())
+          << "use of value defined in partition #" << partition.getIndex()
+          << " at " << distance << " iterations in the future";
+      diag.attachNote(output.getLoc())
+          << "value defined here in partition #" << partition.getIndex();
+      failed = true;
+    };
+    iterateUses(loop, &partition, callback);
+    if (failed)
+      return failure();
+  }
+
   return success();
 }
 

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

@@ -499,4 +499,10 @@ void PartitionLoops::runOnOperation() {
     if (failed(partitionLoop(loop)))
       return signalPassFailure();
   }
+
+  OpPassManager pm;
+  pm.addPass(mlir::triton::createNVWSLowerWarpGroup());
+
+  if (failed(runPipeline(pm, getOperation())))
+    return signalPassFailure();
 }