[Mlir-commits] [mlir] [mlir][tensor] Fold producer linalg transpose with consumer tensor pack (PR #75658)
Prathamesh Tagore
llvmlistbot at llvm.org
Fri Dec 15 13:26:47 PST 2023
https://github.com/meshtag created https://github.com/llvm/llvm-project/pull/75658
Successor to https://github.com/llvm/llvm-project/pull/74206
Partial fix to https://github.com/openxla/iree/issues/15367
>From c7706210742e65deef4afea9a0caaf0715b152b6 Mon Sep 17 00:00:00 2001
From: meshtag <prathameshtagore at gmail.com>
Date: Fri, 15 Dec 2023 07:42:19 +0000
Subject: [PATCH 1/2] Add support for folding consumer pack with producer
transpose
---
.../FoldIntoPackAndUnpackPatterns.cpp | 132 ++++++++++++++----
.../Tensor/fold-into-pack-and-unpack.mlir | 115 +++++++++++++++
2 files changed, 217 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
index e4509b331beeac..d9dc365c7e85cd 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
@@ -21,6 +21,57 @@ static bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
}
+/// Helper function to generate an equivalent permutation map for
+/// `linalg.transpose` and `tensor.pack` which will be used after their folding
+/// into a `tensor.pack`.
+static bool getRemappedPermutationForTransposeAndPack(
+ PackOp packOp, linalg::TransposeOp transposeOp,
+ SmallVector<int64_t> &newOuterDimsPermVec,
+ SmallVector<int64_t> &newInnerDimsPosVec,
+ SmallVector<OpFoldResult> &newMixedInnerTilesVec,
+ bool isTransposeProducer) {
+ bool foldingPossible = true;
+ auto innerDimsPos = packOp.getInnerDimsPos();
+ auto mixedInnerTiles = packOp.getMixedTiles();
+ auto outerDimsPerm = packOp.getOuterDimsPerm();
+ auto transposePerm = transposeOp.getPermutation();
+ int64_t srcRank = packOp.getSourceRank();
+
+ // Note: if isTransposeProducer = true, transposePerm.size() = srcRank, else
+ // transposePerm.size() > srcRank
+
+ // Process transpose operation for non-tiled outer dimensions
+ for (unsigned int i = 0; i < srcRank; ++i) {
+ int64_t remappedPosition =
+ isTransposeProducer ? (!outerDimsPerm.empty() ? outerDimsPerm[i] : i)
+ : transposePerm[i];
+
+ if (remappedPosition >= srcRank) {
+ foldingPossible = false;
+ return foldingPossible;
+ }
+
+ remappedPosition =
+ isTransposeProducer
+ ? transposePerm[remappedPosition]
+ : (!outerDimsPerm.empty() ? outerDimsPerm[remappedPosition]
+ : remappedPosition);
+
+ newOuterDimsPermVec.push_back(remappedPosition);
+ }
+
+ // Process transpose operation for tiled inner dimensions
+ for (unsigned int i = srcRank; i < srcRank + mixedInnerTiles.size(); ++i) {
+ int64_t remappedPosition =
+ isTransposeProducer ? i - srcRank : transposePerm[i] - srcRank;
+
+ newMixedInnerTilesVec.push_back(mixedInnerTiles[remappedPosition]);
+ newInnerDimsPosVec.push_back(innerDimsPos[remappedPosition]);
+ }
+
+ return foldingPossible;
+}
+
/// Fold a `pad` -> `pack` into `pack` if they have the same padding values and
/// the pad op has zero low paddings, or if `pack` has no padding values.
struct FoldPadWithPackOp : public OpRewritePattern<PackOp> {
@@ -96,39 +147,19 @@ struct FoldProducerPackWithConsumerLinalgTransposeOp
if (!packOp)
return failure();
- auto innerDimsPos = packOp.getInnerDimsPos();
- auto mixedInnerTiles = packOp.getMixedTiles();
- auto outerDimsPerm = packOp.getOuterDimsPerm();
- auto transposePerm = transposeOp.getPermutation();
SmallVector<int64_t> newOuterDimsPermVec;
SmallVector<int64_t> newInnerDimsPosVec;
SmallVector<OpFoldResult> newMixedInnerTilesVec;
- int64_t srcRank = packOp.getSourceRank();
-
- // Process transpose operation for non-tiled outer dimensions
- for (unsigned int i = 0; i < srcRank; ++i) {
- int64_t remappedPosition = transposePerm[i];
-
- // If tensor.pack has outer_dims_perm attribute, then consider it during
- // index remapping.
- if (!outerDimsPerm.empty()) {
- if (transposePerm[i] >= srcRank) {
- return rewriter.notifyMatchFailure(
- transposeOp,
- "Cannot fold in tensor.pack if a tile dimension was transposed "
- "with a non-tile dimension in linalg.transpose.");
- }
- remappedPosition = outerDimsPerm[remappedPosition];
- }
-
- newOuterDimsPermVec.push_back(remappedPosition);
- }
- // Process transpose operation for tiled inner dimensions
- for (unsigned int i = srcRank; i < transposePerm.size(); ++i) {
- int64_t remappedPosition = transposePerm[i] - srcRank;
- newMixedInnerTilesVec.push_back(mixedInnerTiles[remappedPosition]);
- newInnerDimsPosVec.push_back(innerDimsPos[remappedPosition]);
+ bool foldingPossible = getRemappedPermutationForTransposeAndPack(
+ packOp, transposeOp, newOuterDimsPermVec, newInnerDimsPosVec,
+ newMixedInnerTilesVec, /*isTransposeProducer*/ false);
+
+ if (!foldingPossible) {
+ return rewriter.notifyMatchFailure(
+ transposeOp,
+ "Cannot fold in tensor.pack if a tile dimension was transposed "
+ "with a non-tile dimension in linalg.transpose.");
}
Value output = packOp.createDestinationTensor(
@@ -142,11 +173,52 @@ struct FoldProducerPackWithConsumerLinalgTransposeOp
return success();
}
};
+
+/// Fold 'transpose' -> 'pack' into 'pack' since 'pack' already has transpose
+/// semantics.
+struct FoldConsumerPackWithProducerLinalgTransposeOp
+ : public OpRewritePattern<PackOp> {
+ using OpRewritePattern<PackOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(PackOp packOp,
+ PatternRewriter &rewriter) const override {
+ auto transposeOp = packOp.getSource().getDefiningOp<linalg::TransposeOp>();
+
+ if (!transposeOp)
+ return failure();
+
+ SmallVector<int64_t> newOuterDimsPermVec;
+ SmallVector<int64_t> newInnerDimsPosVec;
+ SmallVector<OpFoldResult> newMixedInnerTilesVec;
+
+ bool foldingPossible = getRemappedPermutationForTransposeAndPack(
+ packOp, transposeOp, newOuterDimsPermVec, newInnerDimsPosVec,
+ newMixedInnerTilesVec, /*isTransposeProducer*/ true);
+
+ if (!foldingPossible) {
+ return rewriter.notifyMatchFailure(
+ transposeOp,
+ "Cannot fold in tensor.pack if a tile dimension was transposed "
+ "with a non-tile dimension in linalg.transpose.");
+ }
+
+ Value output = packOp.createDestinationTensor(
+ rewriter, packOp.getLoc(), transposeOp.getOperand(0),
+ newMixedInnerTilesVec, newInnerDimsPosVec, newOuterDimsPermVec);
+
+ rewriter.replaceOpWithNewOp<PackOp>(
+ packOp, transposeOp.getOperand(0), output, newInnerDimsPosVec,
+ newMixedInnerTilesVec, packOp.getPaddingValue(), newOuterDimsPermVec);
+
+ return success();
+ }
+};
} // namespace
void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns) {
patterns.insert<FoldUnpackWithExtractSliceOp, FoldPadWithPackOp,
- FoldProducerPackWithConsumerLinalgTransposeOp>(
+ FoldProducerPackWithConsumerLinalgTransposeOp,
+ FoldConsumerPackWithProducerLinalgTransposeOp>(
patterns.getContext());
}
diff --git a/mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir b/mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir
index ca4eb4ff679445..ed101883a40f9a 100644
--- a/mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/fold-into-pack-and-unpack.mlir
@@ -345,3 +345,118 @@ func.func @tensor_pack_linalg_transpose_fold_dynamic_outer_dims_tile_dims_tile_s
// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [2, 1, 3, 0] inner_dims_pos = [3, 1, 2] inner_tiles = [%[[ARG3]], %[[ARG1]], %[[ARG2]]] into %[[INIT]] : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
// CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>
// CHECK: }
+
+// -----
+
+func.func @linalg_transpose_tensor_pack_fold(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x57x56x2x32xf32> {
+ %0 = tensor.empty() : tensor<1x56x57x64xf32>
+ %transposed = linalg.transpose
+ ins(%arg0 : tensor<56x57x1x64xf32>)
+ outs(%0 : tensor<1x56x57x64xf32>)
+ permutation = [2, 0, 1, 3]
+
+ %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
+ %pack = tensor.pack %transposed
+ outer_dims_perm = [0, 2, 1, 3]
+ inner_dims_pos = [3]
+ inner_tiles = [32]
+ into %1 : tensor<1x56x57x64xf32> -> tensor<1x57x56x2x32xf32>
+ return %pack : tensor<1x57x56x2x32xf32>
+}
+// CHECK: func @linalg_transpose_tensor_pack_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 @linalg_transpose_tensor_pack_fold_with_padding(%arg0: tensor<56x57x1x55xf32>, %padding: f32) -> tensor<1x57x56x2x32xf32> {
+ %0 = tensor.empty() : tensor<1x56x57x55xf32>
+ %transpose = linalg.transpose
+ ins(%arg0 : tensor<56x57x1x55xf32>)
+ outs(%0 : tensor<1x56x57x55xf32>)
+ permutation = [2, 0, 1, 3]
+
+ %1 = tensor.empty() : tensor<1x57x56x2x32xf32>
+ %pack = tensor.pack %transpose padding_value(%padding : f32)
+ outer_dims_perm = [0, 2, 1, 3]
+ inner_dims_pos = [3]
+ inner_tiles = [32]
+ into %1 : tensor<1x56x57x55xf32> -> tensor<1x57x56x2x32xf32>
+ return %pack : tensor<1x57x56x2x32xf32>
+}
+// CHECK: func @linalg_transpose_tensor_pack_fold_with_padding(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x55xf32>, %[[PADDING:.+]]: f32)
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x57x56x2x32xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] padding_value(%[[PADDING]] : f32)
+// 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 @linalg_transpose_tensor_pack_fold_no_outer_dims_perm(%arg0: tensor<56x57x1x64xf32>) -> tensor<1x56x57x2x32xf32> {
+ %0 = tensor.empty() : tensor<1x56x57x64xf32>
+ %transposed = linalg.transpose
+ ins(%arg0 : tensor<56x57x1x64xf32>)
+ outs(%0 : tensor<1x56x57x64xf32>)
+ permutation = [2, 0, 1, 3]
+
+ %1 = tensor.empty() : tensor<1x56x57x2x32xf32>
+ %pack = tensor.pack %transposed
+ inner_dims_pos = [3]
+ inner_tiles = [32]
+ into %1 : tensor<1x56x57x64xf32> -> tensor<1x56x57x2x32xf32>
+ return %pack : tensor<1x56x57x2x32xf32>
+}
+// CHECK: func @linalg_transpose_tensor_pack_fold_no_outer_dims_perm(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<56x57x1x64xf32>)
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x56x57x2x32xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]]
+// CHECK-SAME: outer_dims_perm = [2, 0, 1, 3]
+// CHECK-SAME: inner_dims_pos = [3] inner_tiles = [32]
+// CHECK-SAME: into %[[INIT]]
+// CHECK: return %[[PACK]]
+
+// -----
+
+func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_sizes(%arg0: tensor<?x?x?x?xf32>, %transpose_dest: tensor<?x?x?x?xf32>, %pack_dest: tensor<?x?x?x?x?x?x?xf32>, %tile_p : index, %tile_q : index, %tile_r : index) -> tensor<?x?x?x?x?x?x?xf32> {
+ %transposed = linalg.transpose
+ ins(%arg0 : tensor<?x?x?x?xf32>)
+ outs(%transpose_dest : tensor<?x?x?x?xf32>)
+ permutation = [2, 3, 0, 1]
+
+ %pack = tensor.pack %transposed
+ outer_dims_perm = [3, 0, 2, 1]
+ inner_dims_pos = [1, 3, 2]
+ inner_tiles = [%tile_p, %tile_q, %tile_r]
+ into %pack_dest : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
+ return %pack : tensor<?x?x?x?x?x?x?xf32>
+}
+// CHECK: #[[map:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+// CHECK: module {
+// CHECK: func.func @linalg_transpose_tensor_pack_fold_dynamic_outer_dims_tile_dims_tile_sizes(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>, %[[TRANSPOSE_DEST:.+]]: tensor<?x?x?x?xf32>,
+// CHECK-SAME: %[[PACK_DEST:.+]]: tensor<?x?x?x?x?x?x?xf32>,
+// CHECK-SAME: %[[ARG1:.+]]: index, %[[ARG2:.+]]: index,
+// CHECK-SAME: %[[ARG3:.+]]: index)
+// CHECK: %[[c0:.+]] = arith.constant 0 : index
+// CHECK: %[[c1:.+]] = arith.constant 1 : index
+// CHECK: %[[c2:.+]] = arith.constant 2 : index
+// CHECK: %[[c3:.+]] = arith.constant 3 : index
+// CHECK: %[[dim:.+]] = tensor.dim %[[ARG0]], %[[c0]] : tensor<?x?x?x?xf32>
+// CHECK: %[[dim_0:.+]] = tensor.dim %[[ARG0]], %[[c1]] : tensor<?x?x?x?xf32>
+// CHECK: %[[dim_1:.+]] = tensor.dim %[[ARG0]], %[[c2]] : tensor<?x?x?x?xf32>
+// CHECK: %[[dim_2:.+]] = tensor.dim %[[ARG0]], %[[c3]] : tensor<?x?x?x?xf32>
+// CHECK: %[[mapped_dim0:.+]] = affine.apply #[[map:.+]]()[%[[dim_0]], %[[ARG1]]]
+// CHECK: %[[mapped_dim1:.+]] = affine.apply #[[map:.+]]()[%[[dim_2]], %[[ARG2]]]
+// CHECK: %[[mapped_dim2:.+]] = affine.apply #[[map:.+]]()[%[[dim_1]], %[[ARG3]]]
+// CHECK: %[[INIT:.+]] = tensor.empty(%[[mapped_dim0]], %[[mapped_dim2]], %[[dim]], %[[mapped_dim1]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : tensor<?x?x?x?x?x?x?xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [1, 2, 0, 3] inner_dims_pos = [1, 3, 2] inner_tiles = [%[[ARG1]], %[[ARG2]], %[[ARG3]]] into %[[INIT]] : tensor<?x?x?x?xf32> -> tensor<?x?x?x?x?x?x?xf32>
+// CHECK: return %[[PACK]] : tensor<?x?x?x?x?x?x?xf32>
+// CHECK: }
>From 9b40138f3df4adc0a9874251da026f3e61790937 Mon Sep 17 00:00:00 2001
From: meshtag <prathameshtagore at gmail.com>
Date: Fri, 15 Dec 2023 21:19:15 +0000
Subject: [PATCH 2/2] Rectify incorrect failure message
---
.../Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp | 8 ++------
1 file changed, 2 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
index d9dc365c7e85cd..2c45cd3500fa94 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/FoldIntoPackAndUnpackPatterns.cpp
@@ -195,12 +195,8 @@ struct FoldConsumerPackWithProducerLinalgTransposeOp
packOp, transposeOp, newOuterDimsPermVec, newInnerDimsPosVec,
newMixedInnerTilesVec, /*isTransposeProducer*/ true);
- if (!foldingPossible) {
- return rewriter.notifyMatchFailure(
- transposeOp,
- "Cannot fold in tensor.pack if a tile dimension was transposed "
- "with a non-tile dimension in linalg.transpose.");
- }
+ if (!foldingPossible)
+ return failure();
Value output = packOp.createDestinationTensor(
rewriter, packOp.getLoc(), transposeOp.getOperand(0),
More information about the Mlir-commits
mailing list