[Mlir-commits] [mlir] f59eef6 - [mlir][tensor] Enhance SimplifyPackToExpandShape for unit dim cases. (#79247)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jan 24 18:47:29 PST 2024
Author: Han-Chung Wang
Date: 2024-01-24T18:47:25-08:00
New Revision: f59eef6515433577d757cf64d2d2f402d95a689e
URL: https://github.com/llvm/llvm-project/commit/f59eef6515433577d757cf64d2d2f402d95a689e
DIFF: https://github.com/llvm/llvm-project/commit/f59eef6515433577d757cf64d2d2f402d95a689e.diff
LOG: [mlir][tensor] Enhance SimplifyPackToExpandShape for unit dim cases. (#79247)
Progress on https://github.com/openxla/iree/issues/16181
Added:
Modified:
mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
index 06be017f24b823e..b404543ddef867e 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
@@ -22,6 +22,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 +40,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 isPackOnInnerMostDim(RewriterBase &rewriter,
+ PackOp packOp) const {
auto outerDimsPerm = packOp.getOuterDimsPerm();
if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
return rewriter.notifyMatchFailure(
@@ -46,14 +50,50 @@ 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 {
+ assert(!packOp.getPaddingValue() &&
+ "expect the op does not have padding value.");
+ ArrayRef<int64_t> srcShape = packOp.getSourceType().getShape();
+ if (getNumGtOneDims(srcShape) > 1) {
+ return rewriter.notifyMatchFailure(
+ packOp, "expects source to have at most one non-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 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(isPackOnInnerMostDim(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..859eb5ee2a7061d 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_1x32_to_1x32x1x1
+// 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_1x32x1x1(%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