[WS] Reimplement PartitionLoops pass supporting general control flow (triton-lang#7415)

`PartitionLoops` is now implemented using the idea of mutual recursion
between for ops, if ops etc partitioning, as pioneered by the Meta WS
branch. It should be ready for nested loops, although that has not been
tested. If we agree that this work is on the right track, I'll work on
manually creating a lit test for nested loops.

The code is mostly ported from [the code split implementation in the
NVWS
branch](https://github.com/triton-lang/triton/blob/aref_auto_ws/third_party/nvidia/lib/Dialect/NVWS/Transforms/ArefCodeSplit.cpp).
As discussed between NV and @jeffniu-openai, we are now emitting
`nvws.warp_group` op as the result of code splitting, and to make the
new implementation a drop-in replacement for the existing one, the
`nvws.warp_group` -> `ttg.warp_specialize` conversion pass by
@mbrookhart is immediately run. Thus, all lit tests for
`tritongpu-partition-loops` pass as is and no e2e workload should break
as the result of this refactoring. Depending on a future direction, we
may move the code spitting part earlier in the pipeline, or retire
`nvws.warp_group` and just emit `ttg.warp_specialize` directly.

For now, I'm adding a dependency on the NVWS dialect to TritonGPU, to
use `nvws.warp_group` op and `nvws-lower-warp-group` from
`PartitionLoops.cpp`. I'll move them to TritonGPU if the new
implementation looks good.

cc @3gx @htyu @manman-ren

---------

Co-authored-by: Jeff Niu <[email protected]>

Author: masahi

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

@@ -2,6 +2,7 @@
 #define TRITON_DIALECT_TRITONGPU_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
+#include "nvidia/include/Dialect/NVWS/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 namespace mlir {

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

@@ -106,7 +106,8 @@ def TritonGPUAutomaticWarpSpecialization : Pass<"tritongpu-automatic-warp-specia
     "mlir::triton::gpu::TritonGPUDialect",
     "mlir::scf::SCFDialect",
     "mlir::arith::ArithDialect",
-    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
+    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
+    "triton::nvws::NVWSDialect"
   ];
 
   let options = [
@@ -143,7 +144,10 @@ def TritonGPUPartitionLoops : Pass<"tritongpu-partition-loops", "mlir::ModuleOp"
     between any of the partitions.
   }];
 
-  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect"];
+  let dependentDialects = [
+    "mlir::triton::gpu::TritonGPUDialect",
+    "triton::nvws::NVWSDialect"
+  ];
 }
 
 def TritonGPUOptimizePartitionWarps : Pass<"tritongpu-optimize-partition-warps", "mlir::ModuleOp"> {

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -46,6 +46,8 @@ add_triton_library(TritonGPUTransforms
   TritonTransforms
   TritonGPUIR
   TritonNvidiaGPUIR
+  NVWSIR
+  NVWSTransforms
   TritonToTritonGPU
   TritonInstrumentIR
   MLIRTransformUtils