[Mlir-commits] [mlir] [mlir][tensor] Enhance SimplifyUnPackToCollapseShape for unit dim cases. (PR #79262)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jan 24 00:04:10 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor
Author: Han-Chung Wang (hanhanW)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/79262.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp (+75-10)
- (modified) mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir (+102)
``````````diff
diff --git a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
index 06be017f24b823e..c3147ed1847d0df 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)
@@ -77,8 +115,9 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
operand, reassociation);
}
- LogicalResult matchAndRewrite(UnPackOp unpackOp,
- PatternRewriter &rewriter) const override {
+ /// Returns success() if it is unpacking at the innermost dimension.
+ LogicalResult isUnpackAtInnerMostDim(RewriterBase &rewriter,
+ UnPackOp unpackOp) const {
auto outerDimsPerm = unpackOp.getOuterDimsPerm();
if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
return rewriter.notifyMatchFailure(
@@ -97,6 +136,32 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
unpackOp, "expects unpacking at the innermost dimension");
}
+ return success();
+ }
+
+ /// Returns success() if it unpacks a 1D or 2D source operand to a 1D
+ /// destination.
+ LogicalResult isUnpackTo1DDest(RewriterBase &rewriter,
+ UnPackOp unpackOp) const {
+ ArrayRef<int64_t> destShape = unpackOp.getDestType().getShape();
+ if (getNumGtOneDims(destShape) > 1)
+ return failure();
+
+ SmallVector<int64_t> innerTiles = unpackOp.getStaticTiles();
+ if (getNumGtOneDims(innerTiles) > 1)
+ return failure();
+
+ return success();
+ }
+
+ LogicalResult matchAndRewrite(UnPackOp unpackOp,
+ PatternRewriter &rewriter) const override {
+ if (failed(isUnpackAtInnerMostDim(rewriter, unpackOp)) &&
+ failed(isUnpackTo1DDest(rewriter, unpackOp)))
+ return failure();
+
+ RankedTensorType sourceType = unpackOp.getSourceType();
+ RankedTensorType destType = unpackOp.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..c731dd6ea44a4da 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>
@@ -164,3 +215,54 @@ func.func @single_first_inner_dim_unpacking(%arg0: tensor<8x5x32xf32>) -> tensor
%0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<8x5x32xf32> -> tensor<256x5xf32>
return %0 : tensor<256x5xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_1d_to_1d
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK: return %[[COLLAPSED]]
+func.func @unpack_1d_to_1d(%arg0 : tensor<1x32x1x1xf32>) -> tensor<1x32xf32> {
+ %empty = tensor.empty() : tensor<1x32xf32>
+ %unpack = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
+ : tensor<1x32x1x1xf32> -> tensor<1x32xf32>
+ return %unpack : tensor<1x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_1x2x1x16_to_1x32
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK: return %[[COLLAPSED]]
+func.func @unpack_1x2x1x16_to_1x32(%arg0 : tensor<1x2x1x16xf32>) -> tensor<1x32xf32> {
+ %empty = tensor.empty() : tensor<1x32xf32>
+ %unpack = tensor.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 16] into %empty
+ : tensor<1x2x1x16xf32> -> tensor<1x32xf32>
+ return %unpack : tensor<1x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_16x1x2x1_to_32x1
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
+// CHECK: return %[[COLLAPSED]]
+func.func @unpack_16x1x2x1_to_32x1(%arg0 : tensor<1x16x2x1xf32>) -> tensor<32x1xf32> {
+ %empty = tensor.empty() : tensor<32x1xf32>
+ %unpack = tensor.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1] into %empty
+ : tensor<1x16x2x1xf32> -> tensor<32x1xf32>
+ return %unpack : tensor<32x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_16x1x1x2_to_32x1
+// CHECK-NOT: tensor.collapse_shape
+// CHECK: tensor.unpack
+func.func @unpack_16x1x1x2_to_32x1(%arg0 : tensor<16x1x1x2xf32>) -> tensor<32x1xf32> {
+ %empty = tensor.empty() : tensor<32x1xf32>
+ %unpack = tensor.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
+ : tensor<16x1x1x2xf32> -> tensor<32x1xf32>
+ return %unpack : tensor<32x1xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/79262
More information about the Mlir-commits
mailing list