[mlir][tensor] Fold consumer linalg transpose with producer tensor pack

mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/IR/PatternMatch.h"
@@ -81,10 +82,91 @@ struct FoldUnpackWithExtractSliceOp : public OpRewritePattern<ExtractSliceOp> {
     return success();
   }
 };
+
+/// Fold 'pack' -> 'transpose' into 'pack' since 'pack' already has transpose
+/// semantics.
+struct FoldProducerPackWithConsumerLinalgTransposeOp
+    : public OpRewritePattern<linalg::TransposeOp> {
+  using OpRewritePattern<linalg::TransposeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp,
+                                PatternRewriter &rewriter) const override {
+    auto transposeInputTensor = transposeOp.getOperand(0);
+    auto packOp = transposeInputTensor.getDefiningOp<PackOp>();
+
+    if (!packOp)
+      return failure();
+
+    auto packInnerDimsPos = packOp.getInnerDimsPos();
+    auto packInnerTiles = packOp.getStaticInnerTiles();
+    auto packOuterDimsPerm = packOp.getOuterDimsPerm();
+    auto transposePerm = transposeOp.getPermutation();
+    SmallVector<int64_t> newPackOuterDimsPermVec;
+    SmallVector<int64_t> newPackInnerDimsPosVec;
+    SmallVector<int64_t> newPackInnerTilesVec;
+
+    // Variable for storing translated position after considering original
+    // outer_dims_perm and permutation attributes of tensor.pack and
+    // linalg.transpose.
+    int64_t translatedPosition;
+
+    // Process transpose operation for non-tiled outer dimensions of the tensor.
+    for (unsigned int i = 0; i < transposePerm.size() - packInnerTiles.size();
+         ++i) {
+      // If tensor.pack has outer_dims_perm attribute, then consider it during
+      // index translation.
+      if (packOuterDimsPerm.size()) {
+        // Note: static_cast is added around transposePerm[i] to suppress the
+        // compiler warning of comparison between variables of different types.
+        if (static_cast<unsigned long>(transposePerm[i]) <
+            packOuterDimsPerm.size())
+          translatedPosition = packOuterDimsPerm[transposePerm[i]];
+        else
+          return rewriter.notifyMatchFailure(
+              transposeOp,
+              "Cannot fold in tensor.pack if a tile dimension was transposed "
+              "with a non-tile dimension in linalg.transpose.");
+      } else
+        translatedPosition = transposePerm[i];
+
+      newPackOuterDimsPermVec.push_back(translatedPosition);
+    }
+
+    // Process transpose operation for tiled inner dimensions of the tensor.
+    for (unsigned int i = transposePerm.size() - packInnerTiles.size();
+         i < transposePerm.size(); ++i) {
+      translatedPosition =
+          transposePerm[i] - (transposePerm.size() - packInnerTiles.size());
+
+      newPackInnerTilesVec.push_back(packInnerTiles[translatedPosition]);
+      newPackInnerDimsPosVec.push_back(packInnerDimsPos[translatedPosition]);
+    }
+
+    SmallVector<OpFoldResult> opFoldResultsTiles;
+    opFoldResultsTiles.reserve(newPackInnerTilesVec.size());
+
+    transform(newPackInnerTilesVec, std::back_inserter(opFoldResultsTiles),
+              [&rewriter](int64_t value) {
+                return IntegerAttr::get(IndexType::get(rewriter.getContext()),
+                                        value);
+              });
+
+    Value output = packOp.createDestinationTensor(
+        rewriter, transposeOp.getLoc(), packOp.getSource(), opFoldResultsTiles,
+        newPackInnerDimsPosVec, newPackOuterDimsPermVec);
+
+    rewriter.replaceOpWithNewOp<PackOp>(
+        transposeOp, packOp.getSource(), output, newPackInnerDimsPosVec,
+        opFoldResultsTiles, packOp.getPaddingValue(), newPackOuterDimsPermVec);
+
+    return success();
+  }
+};
 } // namespace
 
 void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns) {
-  patterns.insert<FoldUnpackWithExtractSliceOp, FoldPadWithPackOp>(
+  patterns.insert<FoldUnpackWithExtractSliceOp, FoldPadWithPackOp,
+                  FoldProducerPackWithConsumerLinalgTransposeOp>(
       patterns.getContext());
 }
 

mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir

@@ -114,3 +114,102 @@ func.func @pad_pack_different_padding_value(%src: tensor<16641x16xf32>) -> tenso
 // CHECK-LABEL: func.func @pad_pack_different_padding_value
 // CHECK:         tensor.pad
 // CHECK:         tensor.pack
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {
+  %0 = tensor.empty() : tensor<56x2x1x57x32xf32>
+  %pack = tensor.pack %arg0
+    outer_dims_perm = [0, 3, 2, 1]
+    inner_dims_pos = [3]
+    inner_tiles = [32]
+    into %0 : tensor<56x57x1x64xf32> -> tensor<56x2x1x57x32xf32>
+
+  %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x2x1x57x32xf32>)
+    outs(%1 : tensor<1x57x56x2x32xf32>)
+    permutation = [2, 3, 0, 1, 4]
+  return %transposed : tensor<1x57x56x2x32xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME:      outer_dims_perm = [2, 1, 0, 3]
+// CHECK-SAME:      inner_dims_pos = [3] inner_tiles = [32] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x2x56x57x32xf32> {
+  %0 = tensor.empty() : tensor<56x57x1x2x32xf32>
+  %pack = tensor.pack %arg0
+    inner_dims_pos = [3]
+    inner_tiles = [32]
+    into %0 : tensor<56x57x1x64xf32> -> tensor<56x57x1x2x32xf32>
+
+  %1 = tensor.empty() : tensor<1x2x56x57x32xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x57x1x2x32xf32>)
+    outs(%1 : tensor<1x2x56x57x32xf32>)
+    permutation = [2, 3, 0, 1, 4]
+  return %transposed : tensor<1x2x56x57x32xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold_no_outer_dims_perm(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<1x2x56x57x32xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME:      outer_dims_perm = [2, 3, 0, 1]
+// CHECK-SAME:      inner_dims_pos = [3] inner_tiles = [32] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_tile_dims_transpose(%arg0: tensor<56x64x4x64xf32>) -> tensor<2x2x56x2x32x32x2xf32> {
+  %0 = tensor.empty() : tensor<56x2x2x2x32x2x32xf32>
+  %pack = tensor.pack %arg0
+    outer_dims_perm = [0, 1, 2, 3]
+    inner_dims_pos = [1, 2, 3]
+    inner_tiles = [32, 2, 32]
+    into %0 : tensor<56x64x4x64xf32> -> tensor<56x2x2x2x32x2x32xf32>
+
+  %1 = tensor.empty() : tensor<2x2x56x2x32x32x2xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x2x2x2x32x2x32xf32>)
+    outs(%1 : tensor<2x2x56x2x32x32x2xf32>)
+    permutation = [2, 3, 0, 1, 6, 4, 5]
+  return %transposed : tensor<2x2x56x2x32x32x2xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold_tile_dims_transpose(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x64x4x64xf32>)
+//      CHECK:   %[[INIT:.+]] = tensor.empty() : tensor<2x2x56x2x32x32x2xf32>
+//      CHECK:   %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME:      outer_dims_perm = [2, 3, 0, 1]
+// CHECK-SAME:      inner_dims_pos = [3, 1, 2] inner_tiles = [32, 32, 2] 
+// CHECK-SAME:       into %[[INIT]]
+//      CHECK:   return %[[PACK]]
+
+// -----
+
+func.func @tensor_pack_linalg_transpose_fold_tile_dims_outer_dims_transpose(%arg0: tensor<56x64x4x64xf32>) -> tensor<2x2x2x2x32x32x56xf32> {
+  %0 = tensor.empty() : tensor<56x2x2x2x32x2x32xf32>
+  %pack = tensor.pack %arg0
+    outer_dims_perm = [0, 1, 2, 3]
+    inner_dims_pos = [1, 2, 3]
+    inner_tiles = [32, 2, 32]
+    into %0 : tensor<56x64x4x64xf32> -> tensor<56x2x2x2x32x2x32xf32>
+
+  %1 = tensor.empty() : tensor<2x2x2x2x32x32x56xf32>
+  %transposed = linalg.transpose
+    ins(%pack : tensor<56x2x2x2x32x2x32xf32>)
+    outs(%1 : tensor<2x2x2x2x32x32x56xf32>)
+    permutation = [2, 3, 5, 1, 6, 4, 0]
+  return %transposed : tensor<2x2x2x2x32x32x56xf32>
+}
+//      CHECK: func @tensor_pack_linalg_transpose_fold_tile_dims_outer_dims_transpose(
+// CHECK-SAME:     %[[ARG0:.+]]: tensor<56x64x4x64xf32>)
+//      CHECK:   tensor.pack
+//      CHECK:   linalg.transpose