add collapse shape pass for linalg{broadcast,transpose,fill,reduce} (#302)

linalg.generic with broadcast
before
```
    %79 = linalg.generic {indexing_maps = [#map7, #map4], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%76 : tensor<1x1x1x1x1x1x1x1x2x2xi32>) outs(%77 : tensor<2x2x2x2x2x2x2x2x2x2xi32>) attrs =  {broadcastDims = array<i64: 0, 1, 2, 3, 4, 5, 6, 7>} {
    ^bb0(%in: i32, %out: i32):
      linalg.yield %in : i32
    } -> tensor<2x2x2x2x2x2x2x2x2x2xi32>
```
after
```
    %collapsed_30 = tensor.collapse_shape %89 [[0, 1, 2, 3, 4, 5, 6, 7], [8, 9]] : tensor<1x1x1x1x1x1x1x1x2x2xi32> into tensor<1x4xi32>
    %92 =  tensor.collapse_shape %77 [[0, 1, 2, 3, 4, 5, 6, 7], [8, 9]]  : tensor<2x2x2x2x2x2x2x2x2x2xi32> into tensor<256x4xi32>
    %93 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel"]} ins(%collapsed_30 : tensor<1x4xi32>) outs(%92 : tensor<256x4xi32>) attrs =  {broadcastDims = array<i64: 0>} {
    ^bb0(%in: i32, %out: i32):
      linalg.yield %in : i32
    } -> tensor<256x4xi32>
    %expanded_31 = tensor.expand_shape %93 [[0, 1, 2, 3, 4, 5, 6, 7], [8, 9]] output_shape [2, 2, 2, 2, 2, 2, 2, 2, 2, 2] : tensor<256x4xi32> into tensor<2x2x2x2x2x2x2x2x2x2xi32>
```

linalg.transpose
before
```
    %transposed = linalg.transpose ins(%expanded_2 : tensor<2x2x2x2x2x2x2x2x2x2xi64>) outs(%66 : tensor<2x2x2x2x2x2x2x2x2x2xi64>) permutation = [0, 1, 2, 3, 4, 5, 6, 7, 9, 8] 
```
after
```
    %collapsed = tensor.collapse_shape %expanded_17 [[0, 1, 2, 3, 4, 5, 6, 7], [8], [9]] : tensor<2x2x2x2x2x2x2x2x2x2xi64> into tensor<256x2x2xi64>
    %77 = tensor.collapse_shape %66 ...: 
    %transposed = linalg.transpose ins(%collapsed : tensor<256x2x2xi64>) outs(%77 : tensor<256x2x2xi64>) permutation = [0, 2, 1] 
    %expanded_22 = tensor.expand_shape  %transposed ...
```

linalg.fill
before
```
    %13 = linalg.fill ins(%c1_i32 : i32) outs(%12 : tensor<1x1x1x1x1x2x1xi32>) -> tensor<1x1x1x1x1x2x1xi32>
```
after
```
    %17 = tensor.collapse_shape %12 ...
    %18 = linalg.fill ins(%c1_i32 : i32) outs(%17 : tensor<2xi32>) -> tensor<2xi32>
    %expanded_6 = tensor.expand_shape %18 [[0, 1, 2, 3, 4, 5, 6]] output_shape [1, 1, 1, 1, 1, 2, 1] : tensor<2xi32> into tensor<1x1x1x1x1x2x1xi32>
```

linalg.reduce
before
```
    %reduced = linalg.reduce ins(%transposed : tensor<2x2x2x2x2x2x2x2x2x2xi64>) outs(%68 : tensor<2x2x2x2x2x2x2x2x2xi64>) dimensions = [8] 
      (%in: i64, %init: i64) {
        %311 = arith.xori %in, %init : i64
        linalg.yield %311 : i64
      }
```
after
```
    %collapsed_20 = tensor.collapse_shape %expanded_19 [[0, 1, 2, 3, 4, 5, 6, 7], [8]] : tensor<2x2x2x2x2x2x2x2x2xi64> into tensor<256x2xi64>
    %reduced = linalg.reduce ins(%transposed : tensor<256x2x2xi64>) outs(%collapsed_20 : tensor<256x2xi64>) dimensions = [1] 
      (%in: i64, %init: i64) {
        %377 = arith.xori %in, %init : i64
        linalg.yield %377 : i64
      }
    %expanded_21 = tensor.expand_shape %reduced [[0, 1, 2, 3, 4, 5, 6, 7], [8]] output_shape [2, 2, 2, 2, 2, 2, 2, 2, 2] : tensor<256x2xi64> into tensor<2x2x2x2x2x2x2x2x2xi64>
```

Author: enjustli

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

@@ -13,6 +13,7 @@
 #include "triton-shared/Analysis/PtrAnalysis.h"
 #include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h"
 #include "triton-shared/Utils/Utils.h"
+#include "triton-shared/Conversion/TritonArithToLinalg/ConversionTools.h"
 
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
@@ -109,10 +110,6 @@ static Value getScalarValue(Value operand, Location loc,
   return nullptr;
 }
 
-static SmallVector<utils::IteratorType> getNParallelLoopsAttrs(unsigned n) {
-  return SmallVector<utils::IteratorType>(n, utils::IteratorType::parallel);
-}
-
 // if order is empty, transpose the last two dimensions
 // otherwise, use the provided order.
 // The order must be a permutation of the source rank.
@@ -656,49 +653,6 @@ struct BroadcastConverter : public OpConversionPattern<triton::BroadcastOp> {
 private:
   using OpConversionPattern<triton::BroadcastOp>::OpConversionPattern;
 
-  SmallVector<int64_t> getBroadcastDims(RankedTensorType src,
-                                        RankedTensorType dst) const {
-    SmallVector<int64_t> broadcastDims;
-    auto srcShape = src.getShape();
-    auto dstShape = dst.getShape();
-
-    for (size_t i = 0; i < srcShape.size(); i++) {
-      if (dstShape[i] != srcShape[i]) {
-        assert(srcShape[i] == 1);
-        broadcastDims.push_back(i);
-      }
-    }
-    assert(!broadcastDims.empty() && "cannot identify broadcast dimension");
-    return broadcastDims;
-  }
-
-  // Broadcasts input tensor based on TosaToLinalg's broadcastToShape
-  AffineMap getBroadcastAffineMap(MLIRContext *context,
-                                  ArrayRef<int64_t> inputShape,
-                                  ArrayRef<int64_t> broadcastToShape) const {
-
-    assert(broadcastToShape.size() >= inputShape.size());
-
-    // Create affine map and shapes for tensor initialization.
-    SmallVector<AffineExpr> outExpr;
-
-    size_t diff = broadcastToShape.size() - inputShape.size();
-    for (size_t i = 0; i < broadcastToShape.size(); i++) {
-      if (i < diff) {
-        continue;
-      }
-      size_t j = i - diff;
-      if (inputShape[j] == 1) {
-        // Broadcast singleton dimension
-        outExpr.push_back(mlir::getAffineConstantExpr(0, context));
-        continue;
-      }
-      // Non-broadcast case
-      outExpr.push_back(mlir::getAffineDimExpr(i, context));
-    }
-    return AffineMap::get(broadcastToShape.size(), 0, outExpr, context);
-  }
-
 public:
   LogicalResult
   matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor,

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

@@ -0,0 +1,59 @@
+#ifndef TRITON_CONVERSION_TRITONARITHTOLINALG_CONVERSIONTOOLS_H
+#define TRITON_CONVERSION_TRITONARITHTOLINALG_CONVERSIONTOOLS_H
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+
+namespace mlir {
+namespace triton {
+
+static inline SmallVector<utils::IteratorType> getNParallelLoopsAttrs(unsigned n) {
+  return SmallVector<utils::IteratorType>(n, utils::IteratorType::parallel);
+}
+
+static inline SmallVector<int64_t> getBroadcastDims(RankedTensorType src,
+                                                    RankedTensorType dst) {
+  SmallVector<int64_t> broadcastDims;
+  auto srcShape = src.getShape();
+  auto dstShape = dst.getShape();
+
+  for (size_t i = 0; i < srcShape.size(); i++) {
+    if (dstShape[i] != srcShape[i]) {
+      assert(srcShape[i] == 1);
+      broadcastDims.push_back(i);
+    }
+  }
+  assert(!broadcastDims.empty() && "cannot identify broadcast dimension");
+  return broadcastDims;
+}
+
+// Broadcasts input tensor based on TosaToLinalg's broadcastToShape
+static inline AffineMap
+getBroadcastAffineMap(MLIRContext *context, ArrayRef<int64_t> inputShape,
+                      ArrayRef<int64_t> broadcastToShape) {
+
+  assert(broadcastToShape.size() >= inputShape.size());
+
+  // Create affine map and shapes for tensor initialization.
+  SmallVector<AffineExpr> outExpr;
+
+  size_t diff = broadcastToShape.size() - inputShape.size();
+  for (size_t i = 0; i < broadcastToShape.size(); i++) {
+    if (i < diff) {
+      continue;
+    }
+    size_t j = i - diff;
+    if (inputShape[j] == 1) {
+      // Broadcast singleton dimension
+      outExpr.push_back(mlir::getAffineConstantExpr(0, context));
+      continue;
+    }
+    // Non-broadcast case
+    outExpr.push_back(mlir::getAffineDimExpr(i, context));
+  }
+  return AffineMap::get(broadcastToShape.size(), 0, outExpr, context);
+}
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_CONVERSION_TRITONARITHTOLINALG_CONVERSIONTOOLS_H

include/triton-shared/Conversion/TritonToLinalgExperimental/CollapseShape.h

@@ -0,0 +1,22 @@
+//===----------------------------------------------------------------------===//
+//
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TRITON_CONVERSION_TRITONTOLINALG_CollapseShape_H
+#define TRITON_CONVERSION_TRITONTOLINALG_CollapseShape_H
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace triton {
+
+std::unique_ptr<OperationPass<ModuleOp>> createCollapseShapePass();
+
+} // namespace triton
+} // namespace mlir
+
+#endif // TRITON_CONVERSION_TRITONTOLINALG_CollapseShape_H

include/triton-shared/Conversion/TritonToLinalgExperimental/Passes.h

@@ -11,6 +11,7 @@
 #include "triton-shared/Conversion/TritonToLinalgExperimental/TritonToLinalgExperimental.h"
 #include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h"
 #include "triton-shared/Conversion/TritonToLinalgExperimental/TritonToPtr.h"
+#include "triton-shared/Conversion/TritonToLinalgExperimental/CollapseShape.h"
 
 namespace mlir {
 namespace triton {

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

@@ -15,7 +15,9 @@ def TritonToLinalgExperimental : Pass<"triton-to-linalg-experimental", "mlir::Mo
   let constructor = "triton::createTritonToLinalgExperimentalPass()";
   let options = [
       Option<"enableMakeGatherScatterTensorPtr", "enable-make-gather-scatter", "bool", /*default*/"true",
-             "Enable make_gather_scatter_tptr support">
+             "Enable make_gather_scatter_tptr support">,
+      Option<"enableCollapseShape", "enable-collapse-shape", "bool", /*default*/"false",
+             "Enable collapse shape pass">,
   ];
 }
 
@@ -29,4 +31,8 @@ def TritonToPtr : Pass<"triton-to-ptr", "mlir::ModuleOp"> {
   let constructor = "triton::createTritonToPtrPass()";
 }
 
+def CollapseShape : Pass</*cli-arg*/"collapse-shape", /*Op*/"mlir::ModuleOp"> {
+  let summary = "Compress tensor dimensions to improve linalg{broadcast,transpose,fill,reduce} efficiency";
+  let constructor = "triton::createCollapseShapePass()";
+}
 #endif

lib/Conversion/TritonToLinalgExperimental/CMakeLists.txt

@@ -8,6 +8,7 @@ add_triton_library(TritonToLinalgExperimental
   TritonToLinalgExperimentalPass.cpp
   ReconcilePtrCastsPass.cpp
   TritonToPtrPass.cpp
+  CollapseShape.cpp
 
   DEPENDS
   TritonToLinalgExperimentalConversionPassIncGen