[TritonGPU] Add a separate pass for optimizing partition num warps (triton-lang#6323)

PartitionLoops no longer changes the number of warps of the partitions
-- they are always left at the default number. A separate pass runs
after the important DCE to analyze the IR and determine the number of
warps needed. As before, if there are no tensor computations in the
partition, the number of warps is set to 1.

However, the pass will now also look at tensor computations and shrink
the number of warps. It does this by doing a very rough estimate of the
register usage of the partition and comparing it against the pool of
available registers given the number of warps. If it can reduce the
number of warps, the algorithm iterates to fixed point. This is
important, e.g., for TMA gather or scatter ops that can still use tensor
ops, but generally small 1D tensors.

This still assumes uniform register allocation per warp. In the future,
we can make the pass much more sophisticated by allowing nonuniform
register allocation with setmaxnreg.

The PR implements a pretty major hack to relayout IR by wiping them out
and rerunning `convert-triton-to-tritongpu` which sadly breaks layering
of TTGIR vs TTIR.

Author: Mogball

include/triton/Conversion/TritonToTritonGPU/Passes.td

@@ -36,7 +36,10 @@ def ConvertTritonToTritonGPU: Pass<"convert-triton-to-tritongpu", "mlir::ModuleO
               "number of ctas in a cga">,
         Option<"target", "target",
               "std::string", /*default*/"\"\"",
-              "the GPU target, e.g., cuda:80, hip:gfx942">
+              "the GPU target, e.g., cuda:80, hip:gfx942">,
+        Option<"enableSourceRemat", "enable-source-remat",
+               "bool", /*default*/"false",
+               "enable trivial source rematerialization">,
    ];
 }
 

include/triton/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.h

@@ -18,7 +18,8 @@ std::unique_ptr<OperationPass<ModuleOp>> createConvertTritonToTritonGPUPass();
 // Create the pass with numWarps set explicitly.
 std::unique_ptr<OperationPass<ModuleOp>>
 createConvertTritonToTritonGPUPass(const std::string &target, int numWarps,
-                                   int threadsPerWarp = 32, int numCTAs = 1);
+                                   int threadsPerWarp = 32, int numCTAs = 1,
+                                   bool enableSourceRemat = false);
 
 } // namespace triton
 } // namespace mlir

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

@@ -155,6 +155,16 @@ def TritonGPUPartitionLoops : Pass<"tritongpu-partition-loops", "mlir::ModuleOp"
   let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect"];
 }
 
+def TritonGPUOptimizePartitionWarps : Pass<"tritongpu-optimize-partition-warps", "mlir::ModuleOp"> {
+  let summary = "optimize the number of warps assigned to partitions";
+
+  let description = [{
+    The `tritongpu-optimize-partition-warps` pass will analyze the partitions
+    of `ttg.warp_specialize` ops and attempts to reduce the number of warps
+    assigned to them and optimize the register usage of the partitions.
+  }];
+}
+
 def TritonGPULoadMMASpecialization : Pass<"tritongpu-load-mma-specialization", "mlir::ModuleOp"> {
   let summary = "load MMA specialization";
 

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

@@ -14,7 +14,7 @@ namespace mlir {
 class TritonGPUTypeConverter : public TypeConverter {
 public:
   TritonGPUTypeConverter(MLIRContext *context, int numWarps, int threadsPerWarp,
-                         int numCTAs);
+                         int numCTAs, bool enableSourceRemat);
   int getNumWarps() const { return numWarps; }
   int getThreadsPerWarp() const { return threadsPerWarp; }
   int getNumCTAs() const { return numCTAs; }

lib/Conversion/TritonGPUToLLVM/ReduceOpToLLVM.cpp

@@ -136,7 +136,6 @@ struct ReduceOpConversion
       uniqueOffsets.insert({offsets[i], i});
     }
 
-    unsigned srcElems = getTotalElemsPerThread(operandType);
     auto *combineOp = &op.getCombineOp();
     auto srcIndices = emitIndices(op.getLoc(), rewriter, targetInfo,
                                   helper.getSrcLayout(), operandType, true);

lib/Conversion/TritonToTritonGPU/CMakeLists.txt

@@ -12,5 +12,4 @@ add_triton_library(TritonToTritonGPU
     TritonIR
     ProtonIR
     TritonGPUIR
-    TritonGPUTransforms
 )

lib/Conversion/TritonToTritonGPU/TritonGPUConversion.cpp

@@ -9,6 +9,7 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 using namespace mlir;
 using namespace mlir::triton::gpu;
@@ -18,7 +19,8 @@ using namespace mlir::triton::gpu;
 //
 TritonGPUTypeConverter::TritonGPUTypeConverter(MLIRContext *context,
                                                int numWarps, int threadsPerWarp,
-                                               int numCTAs)
+                                               int numCTAs,
+                                               bool enableSourceRemat)
     : context(context), numWarps(numWarps), threadsPerWarp(threadsPerWarp),
       numCTAs(numCTAs) {
   addConversion([](Type type) { return type; });
@@ -55,28 +57,28 @@ TritonGPUTypeConverter::TritonGPUTypeConverter(MLIRContext *context,
   //
   // This will be called when (newArgType != origArgType)
   // This will create newArg, and map(origArg, newArg)
-  addArgumentMaterialization([&](OpBuilder &builder,
-                                 RankedTensorType tensorType, ValueRange inputs,
-                                 Location loc) -> Value {
+  addArgumentMaterialization([](OpBuilder &builder, RankedTensorType tensorType,
+                                ValueRange inputs, Location loc) -> Value {
     llvm_unreachable("Argument rematerialization should not happen in Triton "
                      "-> TritonGPU conversion");
     return {};
   });
 
   // If the origValue still has live user(s), use this to
   // convert origValue to newValue
-  addSourceMaterialization([&](OpBuilder &builder, RankedTensorType tensorType,
+  addSourceMaterialization([=](OpBuilder &builder, RankedTensorType tensorType,
                                ValueRange inputs, Location loc) -> Value {
-    llvm_unreachable("Source rematerialization should not happen in Triton -> "
-                     "TritonGPU Conversion");
-    return {};
+    assert(enableSourceRemat && "Source rematerialization should not happen in "
+                                "Triton -> TritonGPU Conversion");
+    return builder.create<UnrealizedConversionCastOp>(loc, tensorType, inputs)
+        .getResult(0);
   });
 
   // This will be called when (desiredType != newOperandType)
   // where, desiredType = typeConverter->convertType(origType)
   // NOTE: only for remapped values.
-  addTargetMaterialization([&](OpBuilder &builder, RankedTensorType tensorType,
-                               ValueRange inputs, Location loc) {
+  addTargetMaterialization([](OpBuilder &builder, RankedTensorType tensorType,
+                              ValueRange inputs, Location loc) {
     auto cast =
         builder.create<triton::gpu::ConvertLayoutOp>(loc, tensorType, inputs);
     return cast.getResult();
@@ -98,7 +100,8 @@ TritonGPUConversionTarget::TritonGPUConversionTarget(
 
   addDynamicallyLegalDialect<arith::ArithDialect, math::MathDialect,
                              triton::TritonDialect, cf::ControlFlowDialect,
-                             scf::SCFDialect, ub::UBDialect>(
+                             scf::SCFDialect, ub::UBDialect,
+                             triton::nvidia_gpu::TritonNvidiaGPUDialect>(
       [&](Operation *op) {
         bool hasLegalRegions = true;
         for (auto &region : op->getRegions()) {

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -10,6 +10,7 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h"
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 #include "llvm/ADT/APSInt.h"
 #include <numeric>
 
@@ -431,50 +432,37 @@ static RankedTensorType getNewIndicesType(RankedTensorType type,
                                newEncoding);
 }
 
-struct TritonDescriptorGatherPattern
-    : public OpConversionPattern<triton::DescriptorGatherOp> {
-  using OpConversionPattern::OpConversionPattern;
-
-  LogicalResult
-  matchAndRewrite(triton::DescriptorGatherOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    auto numThreads = lookupThreadsPerWarp(rewriter);
-    auto numWarps = lookupNumWarps(op);
-    RankedTensorType newType = getNewIndicesType(
-        cast<RankedTensorType>(adaptor.getXOffsets().getType()), numThreads,
-        numWarps);
-    if (!newType)
-      return failure();
-
-    Value newInd = rewriter.create<ConvertLayoutOp>(op.getLoc(), newType,
-                                                    adaptor.getXOffsets());
-    rewriter.replaceOpWithNewOp<triton::DescriptorGatherOp>(
-        op, getTypeConverter()->convertType(op.getType()), adaptor.getDesc(),
-        newInd, adaptor.getYOffset());
-    return success();
-  }
-};
+// Function for converting any gather or scatter op that requires a specific
+// index layout. This also handles converting result types if there are any.
+static LogicalResult convertGatherScatterOp(Operation *op, OpOperand &indices,
+                                            ConversionPatternRewriter &b) {
+  auto type = cast<RankedTensorType>(indices.get().getType());
+  RankedTensorType newType =
+      getNewIndicesType(type, lookupThreadsPerWarp(b), lookupNumWarps(op));
+  if (!newType)
+    return failure();
+  Value index = b.create<ConvertLayoutOp>(op->getLoc(), newType, indices.get());
+  indices.set(index);
+  return success();
+}
 
-struct TritonDescriptorScatterPattern
-    : public OpConversionPattern<triton::DescriptorScatterOp> {
-  using OpConversionPattern::OpConversionPattern;
+template <typename OpT>
+struct GatherScatterOpPattern : public OpConversionPattern<OpT> {
+  using OpConversionPattern<OpT>::OpConversionPattern;
 
   LogicalResult
-  matchAndRewrite(triton::DescriptorScatterOp op, OpAdaptor adaptor,
+  matchAndRewrite(OpT op, typename OpT::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    auto numThreads = lookupThreadsPerWarp(rewriter);
-    auto numWarps = lookupNumWarps(op);
-    RankedTensorType newType = getNewIndicesType(
-        cast<RankedTensorType>(adaptor.getXOffsets().getType()), numThreads,
-        numWarps);
-    if (!newType)
-      return failure();
-
-    Value newInd = rewriter.create<ConvertLayoutOp>(op.getLoc(), newType,
-                                                    adaptor.getXOffsets());
-    rewriter.replaceOpWithNewOp<triton::DescriptorScatterOp>(
-        op, adaptor.getDesc(), newInd, adaptor.getYOffset(), adaptor.getSrc());
-    return success();
+    LogicalResult result = success();
+    rewriter.modifyOpInPlace(op, [&] {
+      for (auto [operand, value] :
+           llvm::zip(op->getOpOperands(), adaptor.getOperands()))
+        operand.set(value);
+      for (OpResult result : op->getOpResults())
+        result.setType(this->typeConverter->convertType(result.getType()));
+      result = convertGatherScatterOp(op, op.getXOffsetsMutable(), rewriter);
+    });
+    return result;
   }
 };
 
@@ -619,10 +607,13 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       GenericOpPattern<triton::ReduceReturnOp>, TritonScanPattern,
       GenericOpPattern<triton::ScanReturnOp>,
       GenericOpPattern<triton::MakeRangeOp>, TritonExpandDimsPattern,
-      TritonTransPattern, TritonDotPattern, TritonDescriptorGatherPattern,
-      TritonDescriptorScatterPattern, GenericOpPattern<triton::LoadOp>,
-      GenericOpPattern<triton::StoreOp>, GenericOpPattern<triton::HistogramOp>,
-      GenericOpPattern<triton::GatherOp>,
+      TritonTransPattern, TritonDotPattern,
+      GatherScatterOpPattern<DescriptorGatherOp>,
+      GatherScatterOpPattern<DescriptorScatterOp>,
+      GatherScatterOpPattern<triton::nvidia_gpu::AsyncTMAGatherOp>,
+      GatherScatterOpPattern<triton::nvidia_gpu::AsyncTMAScatterOp>,
+      GenericOpPattern<triton::LoadOp>, GenericOpPattern<triton::StoreOp>,
+      GenericOpPattern<triton::HistogramOp>, GenericOpPattern<triton::GatherOp>,
       GenericOpPattern<triton::ExternElementwiseOp>,
       GenericOpPattern<triton::PrintOp>, GenericOpPattern<triton::AssertOp>,
       GenericOpPattern<triton::AtomicCASOp>,
@@ -840,11 +831,13 @@ class ConvertTritonToTritonGPU
   ConvertTritonToTritonGPU() = default;
   // constructor with some parameters set explicitly.
   ConvertTritonToTritonGPU(const std::string &target, int numWarps,
-                           int threadsPerWarp, int numCTAs) {
+                           int threadsPerWarp, int numCTAs,
+                           bool enableSourceRemat) {
     this->numWarps = numWarps;
     this->threadsPerWarp = threadsPerWarp;
     this->numCTAs = numCTAs;
     this->target = target;
+    this->enableSourceRemat = enableSourceRemat;
   }
 
   void runOnOperation() override {
@@ -859,7 +852,7 @@ class ConvertTritonToTritonGPU
     ModuleOp mod = getOperation();
     // type converter
     TritonGPUTypeConverter typeConverter(context, numWarps, threadsPerWarp,
-                                         numCTAs);
+                                         numCTAs, enableSourceRemat);
     TritonGPUConversionTarget target(*context, typeConverter);
     // rewrite patterns
     RewritePatternSet patterns(context);
@@ -898,9 +891,10 @@ std::unique_ptr<OperationPass<ModuleOp>>
 mlir::triton::createConvertTritonToTritonGPUPass(const std::string &target,
                                                  int numWarps,
                                                  int threadsPerWarp,
-                                                 int numCTAs) {
-  return std::make_unique<::ConvertTritonToTritonGPU>(target, numWarps,
-                                                      threadsPerWarp, numCTAs);
+                                                 int numCTAs,
+                                                 bool enableSourceRemat) {
+  return std::make_unique<::ConvertTritonToTritonGPU>(
+      target, numWarps, threadsPerWarp, numCTAs, enableSourceRemat);
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -34,6 +34,7 @@ add_triton_library(TritonGPUTransforms
   Utility.cpp
   WarpSpecialization/AutomaticWarpSpecialization.cpp
   WarpSpecialization/LoadMMASpecialization.cpp
+  WarpSpecialization/OptimizePartitionWarps.cpp
   WarpSpecialization/PartitionLoops.cpp
   WarpSpecialization/RewritePartitionDependencies.cpp
 
@@ -47,5 +48,6 @@ add_triton_library(TritonGPUTransforms
   TritonIR
   TritonGPUIR
   TritonNvidiaGPUIR
+  TritonToTritonGPU
   MLIRTransformUtils
 )

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

@@ -51,4 +51,9 @@ void AutomaticWarpSpecialization::runOnOperation() {
   WarpSpecializeOp::getCanonicalizationPatterns(patterns, &getContext());
   if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
     return signalPassFailure();
+
+  pm.clear();
+  pm.addPass(createTritonGPUOptimizePartitionWarps());
+  if (failed(runPipeline(pm, getOperation())))
+    return signalPassFailure();
 }