Adds flag for old transpose behavior on reduction conversion. (#320)

Preferably Triton Shared should not be doing any transposes for
optimizations because we lack the wider analysis to understand if these
transformations are going to be helpful or not. That will be a part of
future work.

For now, we need the old transpose behavior on the Reduce Conversion
where it just transposes rank - 1 reductions to rank - 2. This old
behavior is guarded by a added with a flag that is default to the new
behavior.

Author: red1bluelost

include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp

@@ -1179,9 +1179,15 @@ struct MatmulConverter : public OpConversionPattern<triton::DotOp> {
 };
 
 struct ReduceConverter : public OpConversionPattern<triton::ReduceOp> {
-  using OpConversionPattern<triton::ReduceOp>::OpConversionPattern;
+
+  ReduceConverter(MLIRContext *context, bool transposeToRank0 = true,
+                  PatternBenefit benefit = 1)
+      : OpConversionPattern(context, benefit),
+        transposeToRank0(transposeToRank0) {}
 
 private:
+  bool transposeToRank0;
+
   llvm::SmallVector<Operation *> getRedOps(triton::ReduceOp redOp) const {
     auto reduceBlock = redOp.getBody();
     return llvm::map_to_vector(reduceBlock->without_terminator(),
@@ -1307,19 +1313,32 @@ struct ReduceConverter : public OpConversionPattern<triton::ReduceOp> {
     auto rank = sourceType.getRank();
     auto isVectorReduce = (rank == 1);
 
-    // if it is not a vector reduce, we can transpose the source
-    // so that the reduction axis is the first dimension.
-    if (!isVectorReduce && axis != 0) {
-      SmallVector<int32_t> order;
-      order.reserve(rank);
-      order.push_back(axis);
-      for (int i = 0; i < rank; ++i) {
-        if (i != axis) {
-          order.push_back(i);
+    // For now we are transposing reductions from Triton Shared as an
+    // optimization. This should not be the job of Triton Shared so moving
+    // forward this will be removed. Doing the transpose here lacks a wider
+    // scope of analysis that might indicate that the transpose to a given axis
+    // is not optimal.
+    if (transposeToRank0) {
+      // if it is not a vector reduce, we can transpose the source
+      // so that the reduction axis is the first dimension.
+      if (!isVectorReduce && axis != 0) {
+        SmallVector<int32_t> order;
+        order.reserve(rank);
+        order.push_back(axis);
+        for (int i = 0; i < rank; ++i) {
+          if (i != axis) {
+            order.push_back(i);
+          }
         }
+        source = getTransposedValue(source, op.getLoc(), rewriter, order);
+        axis = 0;
+      }
+    } else {
+      // preserving old behavior until we remove the transpose entirely.
+      if (axis == rank - 1 && !isVectorReduce) {
+        source = getTransposedValue(source, op.getLoc(), rewriter);
+        axis = rank - 2;
       }
-      source = getTransposedValue(source, op.getLoc(), rewriter, order);
-      axis = 0;
     }
 
     bool convertToF32Precision = requiresF32Conversion(resType, rop);

include/triton-shared/Conversion/TritonArithToLinalg/Passes.td

@@ -16,6 +16,10 @@ def TritonArithToLinalg : Pass<"triton-arith-to-linalg", "mlir::ModuleOp"> {
              "Convert tt.assert to cf.assert">,
       Option<"tensorPtrToLinalg", "tensor-ptr-to-linalg", "bool", /*default*/"false",
              "Convert triton ops on tensor of pointers to linalg.generic">,
+      Option<"transposeReduceToRank0", "transpose-reduce-to-rank0", "bool", /*default*/"true",
+             "Transpose reductions to rank 0 which collapses remaining dimensions. Otherwise "
+             "transpose any rank - 1 reduction to rank - 2 as the old behavior. These transposes "
+             "are planned for removal as Triton Shared should not be responsible for this.">,
   ];
 }
 

include/triton-shared/Conversion/TritonArithToLinalg/TritonArithToLinalg.h

@@ -18,13 +18,15 @@ void populateTritonArithToLinalgCanonicalizationPatterns(
 void populateTritonArithToLinalgConversionPatterns(bool pidsToFuncArgs,
                                                    bool addptrToLinalg,
                                                    bool assertToCf,
+                                                   bool transposeReduceToRank0,
                                                    RewritePatternSet &patterns);
 
 // Expand the triton pointer ops operating on pointers to linalg
 void populateTritonTensorPtrConversionPatterns(RewritePatternSet &patterns);
 
 std::unique_ptr<OperationPass<ModuleOp>>
-createTritonArithToLinalgPass(bool tensorPtrToLinalg = false);
+createTritonArithToLinalgPass(bool tensorPtrToLinalg = false,
+                              bool transposeReduceToRank0 = true);
 
 } // namespace triton
 } // namespace mlir

lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp

@@ -50,7 +50,7 @@ void mlir::triton::populateTritonTensorPtrConversionPatterns(
 
 void mlir::triton::populateTritonArithToLinalgConversionPatterns(
     bool pidsToFuncArgs, bool addptrToLinalg, bool assertToCf,
-    RewritePatternSet &patterns) {
+    bool transposeReduceToRank0, RewritePatternSet &patterns) {
 
   if (pidsToFuncArgs) {
     patterns.add<GetProgramIDConverter, GetNumProgramsConverter>(
@@ -96,7 +96,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns(
   // aren't always multiple of 2s, which are sub-optimal for certain hardwares.
   patterns.add<ArgMinConverter>(patterns.getContext());
   patterns.add<ArgMaxConverter>(patterns.getContext());
-  patterns.add<ReduceConverter>(patterns.getContext());
+  patterns.add<ReduceConverter>(patterns.getContext(), transposeReduceToRank0);
 
   // Note: the ordering here matters!
   // These patterns are added last to they will be tried last.

lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp

@@ -186,7 +186,8 @@ class TritonArithToLinalgPass
     }
 
     triton::populateTritonArithToLinalgConversionPatterns(
-        pidsToFuncArgs, addptrToLinalg, assertToCf, patterns);
+        pidsToFuncArgs, addptrToLinalg, assertToCf, transposeReduceToRank0,
+        patterns);
 
     if (pidsToFuncArgs) {
       for (auto func : getOperation().getOps<triton::FuncOp>()) {
@@ -244,8 +245,10 @@ class TritonArithToLinalgPass
 } // namespace
 
 std::unique_ptr<OperationPass<ModuleOp>>
-triton::createTritonArithToLinalgPass(bool tensorPtrToLinalg) {
+triton::createTritonArithToLinalgPass(bool tensorPtrToLinalg,
+                                      bool transposeReduceToRank0) {
   TritonArithToLinalgOptions options;
   options.tensorPtrToLinalg = tensorPtrToLinalg;
+  options.transposeReduceToRank0 = transposeReduceToRank0;
   return std::make_unique<TritonArithToLinalgPass>(options);
 }