[Mlir-commits] [mlir] [mlir][tensor] Enhance SimplifyPackToExpandShape for unit dim cases. (PR #79247)
Han-Chung Wang
llvmlistbot at llvm.org
Tue Jan 23 19:20:55 PST 2024
https://github.com/hanhanW created https://github.com/llvm/llvm-project/pull/79247
None
>From caa67d56d6859ed5b9feed9c2190e80b74c7237a Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 24 Jan 2024 11:18:46 +0800
Subject: [PATCH] [mlir][tensor] Enhance SimplifyPackToExpandShape for unit dim
cases.
---
.../Transforms/PackAndUnpackPatterns.cpp | 54 ++++++++++++++++---
.../Dialect/Tensor/simplify-pack-unpack.mlir | 51 ++++++++++++++++++
2 files changed, 97 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
index 06be017f24b823e..354fe0f0b20a608 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
#include "llvm/Support/Debug.h"
namespace mlir {
@@ -22,6 +23,12 @@ static bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
}
+/// Returns the number of shape sizes that is either dynamic or greater than 1.
+static int64_t getNumGtOneDims(ArrayRef<int64_t> shape) {
+ return llvm::count_if(
+ shape, [](int64_t v) { return ShapedType::isDynamic(v) || v > 1; });
+}
+
/// Packing one-dimensional tensor can be expressed as an expand shape op.
struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
using OpRewritePattern<PackOp>::OpRewritePattern;
@@ -34,11 +41,9 @@ struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
reassociation);
}
- LogicalResult matchAndRewrite(PackOp packOp,
- PatternRewriter &rewriter) const override {
- if (packOp.getPaddingValue())
- return rewriter.notifyMatchFailure(packOp, "expects no padding value");
-
+ /// Returns success() if it is only packing on the innermost dimension.
+ LogicalResult isPackOneInnerMostDim(RewriterBase &rewriter,
+ PackOp packOp) const {
auto outerDimsPerm = packOp.getOuterDimsPerm();
if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
return rewriter.notifyMatchFailure(
@@ -46,14 +51,47 @@ struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
"expects outer_dims_perm is empty or an identity permutation");
}
- RankedTensorType sourceType = packOp.getSourceType();
- RankedTensorType destType = packOp.getDestType();
+ int64_t srcRank = packOp.getSourceRank();
ArrayRef<int64_t> dimsPos = packOp.getInnerDimsPos();
- if (dimsPos.size() != 1 || (dimsPos[0] + 1 != sourceType.getRank())) {
+ if (dimsPos.size() != 1 || (dimsPos[0] + 1 != srcRank)) {
return rewriter.notifyMatchFailure(
packOp, "expects packing at the innermost dimension");
}
+ return success();
+ }
+
+ /// Returns success() if there is only 1 dimension size in source being
+ /// greater than 1 and packing only happens on the dimension. It assumes that
+ /// the pack op does not have padding value.
+ LogicalResult isPack1DSrc(RewriterBase &rewriter, PackOp packOp) const {
+ ArrayRef<int64_t> srcShape = packOp.getSourceType().getShape();
+ if (getNumGtOneDims(srcShape) > 1) {
+ return rewriter.notifyMatchFailure(
+ packOp, "expects source is not 1D tensor with unit dims");
+ }
+ // The pack op does not have padding value. Non-unit inner tile size must be
+ // be used by the non-unit dimension.
+ SmallVector<int64_t> innerTiles = packOp.getStaticTiles();
+ if (getNumGtOneDims(innerTiles) > 1) {
+ return rewriter.notifyMatchFailure(
+ packOp, "expects has at most one non-unit inner tiles");
+ }
+
+ return success();
+ }
+
+ LogicalResult matchAndRewrite(PackOp packOp,
+ PatternRewriter &rewriter) const override {
+ if (packOp.getPaddingValue())
+ return rewriter.notifyMatchFailure(packOp, "expects no padding value");
+
+ if (failed(isPackOneInnerMostDim(rewriter, packOp)) &&
+ failed(isPack1DSrc(rewriter, packOp)))
+ return failure();
+
+ RankedTensorType sourceType = packOp.getSourceType();
+ RankedTensorType destType = packOp.getDestType();
auto reassociation =
getReassociationIndicesForReshape(sourceType, destType);
if (!reassociation)
diff --git a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
index 82bfe6fe8689ab5..00bcacb54d01210 100644
--- a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
@@ -83,6 +83,57 @@ func.func @single_first_inner_dim_packing(%arg0: tensor<256x5xf32>) -> tensor<8x
// -----
+// CHECK-LABEL: func.func @pack_1d_to_1d
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK: return %[[EXPANDED]]
+func.func @pack_1d_to_1d(%arg0 : tensor<1x32xf32>) -> tensor<1x32x1x1xf32> {
+ %empty = tensor.empty() : tensor<1x32x1x1xf32>
+ %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
+ : tensor<1x32xf32> -> tensor<1x32x1x1xf32>
+ return %pack : tensor<1x32x1x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_1x32_to_1x16x1x2
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK: return %[[EXPANDED]]
+func.func @pack_1x32_to_1x16x1x2(%arg0 : tensor<1x32xf32>) -> tensor<1x16x1x2xf32> {
+ %empty = tensor.empty() : tensor<1x16x1x2xf32>
+ %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 2] into %empty
+ : tensor<1x32xf32> -> tensor<1x16x1x2xf32>
+ return %pack : tensor<1x16x1x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_32x1_to_16x1x2x1
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
+// CHECK: return %[[EXPANDED]]
+func.func @pack_32x1_to_16x1x2x1(%arg0 : tensor<32x1xf32>) -> tensor<1x16x2x1xf32> {
+ %empty = tensor.empty() : tensor<1x16x2x1xf32>
+ %pack = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1] into %empty
+ : tensor<32x1xf32> -> tensor<1x16x2x1xf32>
+ return %pack : tensor<1x16x2x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_32x1_to_16x1x1x2
+// CHECK-NOT: tensor.expand_shape
+// CHECK: tensor.pack
+func.func @pack_32x1_to_16x1x1x2(%arg0 : tensor<32x1xf32>) -> tensor<16x1x1x2xf32> {
+ %empty = tensor.empty() : tensor<16x1x1x2xf32>
+ %pack = tensor.pack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
+ : tensor<32x1xf32> -> tensor<16x1x1x2xf32>
+ return %pack : tensor<16x1x1x2xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @unpack_1d_to_collapse
// CHECK-SAME: %[[ARG0:.+]]: tensor<8x32xf32>)
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<8x32xf32> into tensor<256xf32>
More information about the Mlir-commits
mailing list