[WS] enable e2e aref (#8262)

* enable e2e aref flow
* remove `load-mma-specialization`
* update `partition-scheduling` to place unannotated ops into partition
of their users
* remove canonicalization passes from `automatic-warp-specialization`
pass
* improve `assign-stage-phase` to track whether `aref.buffer` is from
`put` or `get`, and do not thread stage/phase through control flow where
unnecessary
* `aref-tmem-insertion` doesn't require annotated on `tmem_alloc` w/o
source, and if `tmem_alloc` with src has same partition as its user, we
do not use arefs for this.

~pending issue with `schedule-loops` and this test
https://github.com/triton-lang/triton/blob/main/test/TritonGPU/automatic-warp-specialization.mlir#L118
when there is if-stmt (being discussed). Mark PR draft for now.~

Author: 3gx

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

@@ -176,24 +176,6 @@ def TritonGPUPartitionScheduling : Pass<"tritongpu-partition-scheduling", "mlir:
   }];
 }
 
-def TritonGPULoadMMASpecialization : Pass<"tritongpu-load-mma-specialization", "mlir::ModuleOp"> {
-  let summary = "load MMA specialization";
-
-  let description = [{
-    The `tritongpu-load-mma-specialization` pass looks for matmul loops in the
-    module and attempts to create a partition schedule, separating async loads
-    and async MMAs into separate partitions.
-  }];
-
-  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
-                           "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"];
-
-  let options = [
-    Option<"numStages", "num-stages", "int32_t", /*default*/"3",
-           "number of pipeline stages">
-  ];
-}
-
 def TritonGPUF32DotTC : Pass<"tritongpu-F32DotTC", "mlir::ModuleOp"> {
   let summary = "3xTF32 trick";
 

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -29,7 +29,6 @@ add_triton_library(TritonGPUTransforms
   Utility.cpp
   LayoutPropagationUtility.cpp
   WarpSpecialization/AutomaticWarpSpecialization.cpp
-  WarpSpecialization/LoadMMASpecialization.cpp
   WarpSpecialization/Partition.cpp
   WarpSpecialization/OptimizePartitionWarps.cpp
   WarpSpecialization/PartitionBuilder.cpp

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

@@ -36,11 +36,7 @@ void AutomaticWarpSpecialization::runOnOperation() {
   OpPassManager pm;
   pm.addPass(createTritonGPUPartitionScheduling());
   pm.addPass(createNVWSInsertAref());
-#if 0
   pm.addPass(createNVWSInsertTmemAref());
-#else
-  pm.addPass(createTritonGPULoadMMASpecialization({numStages}));
-#endif
   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.
@@ -50,19 +46,6 @@ void AutomaticWarpSpecialization::runOnOperation() {
   pm.addPass(createNVWSLowerAref({numStages}));
   pm.addPass(createTritonGPUPartitionLoops());
   pm.addPass(createNVWSLowerWarpGroup());
-  if (failed(runPipeline(pm, getOperation())))
-    return signalPassFailure();
-
-  // Cleanup code generated by warp specialization.
-  RewritePatternSet patterns(&getContext());
-  populateForOpDeadArgumentElimination(patterns);
-  scf::ForOp::getCanonicalizationPatterns(patterns, &getContext());
-  scf::IfOp::getCanonicalizationPatterns(patterns, &getContext());
-  WarpSpecializeOp::getCanonicalizationPatterns(patterns, &getContext());
-  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
-    return signalPassFailure();
-
-  pm.clear();
   pm.addPass(createTritonGPUOptimizePartitionWarps());
   pm.addPass(createTritonGPUScheduleLoops());
   if (failed(runPipeline(pm, getOperation())))