[WS] Update RewritePartitionDepdencies to insert arefs (#7561)

* Integrates Remove Rewrite multiplicity
([PR7371](triton-lang/triton#7371))
 * Teaches `RewritePartitionDependencies` to insert arefs.
* ~~Lower `nvws::async_complete` to `mbarrier.arrive`/`tcgen05.commit`~~
[`lower-aref` directly lowers to `ttng.arrive/commit`]

~~We currently cannot lower `nvws::async_complete` to `mbarrier.arrive`
because the `LowerAref` pass computes the arrival count using
`num_warps`, which may be modified later by the `OptimizePartitionWarps`
pass. The arrival count then disagrees with the actual number of threads
in the group, leading to hang.~~ `nvws::async_complete` is removed

Data indicates that using multi-threaded `mbarrier.arrive` yields
meaningful performance improvements, compated to singe-threaded
`ttng.arrive_barrier` (see tables below in comments). However, resolving
this issue is orthogonal to the current PR and will be addressed in a
follow-up.



<!---
The core Triton is a small number of people, and we receive many PRs
(thank
you!).  To help us review your code more quickly, **if you are a new
contributor (less than 3 PRs merged) we ask that you complete the
following
tasks and include the filled-out checklist in your PR description.**

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->

# New contributor declaration
- [X] I am not making a trivial change, such as fixing a typo in a
comment.

- [X] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [X] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.

- Select one of the following.
  - [X] I have added tests.
    - `/test` for `lit` tests
    - `/unittest` for C++ tests
    - `/python/test` for end-to-end tests
  - [ ] This PR does not need a test because `FILL THIS IN`.

- Select one of the following.
  - [ ] I have not added any `lit` tests.
- [X] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
including the "tests should be minimal" section. (Usually running Python
code
    and using the instructions it generates is not minimal.)

Author: 3gx

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

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

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

@@ -3,6 +3,7 @@
 #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"
 
@@ -42,6 +43,8 @@ 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,57 +191,6 @@ 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,10 +499,4 @@ void PartitionLoops::runOnOperation() {
     if (failed(partitionLoop(loop)))
       return signalPassFailure();
   }
-
-  OpPassManager pm;
-  pm.addPass(mlir::triton::createNVWSLowerWarpGroup());
-
-  if (failed(runPipeline(pm, getOperation())))
-    return signalPassFailure();
 }