[Mlir-commits] [mlir] 0c1c0d5 - [MLIR] Add patterns to bubble-up pack and push-down unpack through collapse/expand shape ops (#85297)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Mar 27 18:32:31 PDT 2024
Author: Jerry Wu
Date: 2024-03-27T21:32:27-04:00
New Revision: 0c1c0d53931636331b59a03ed08f70936835399c
URL: https://github.com/llvm/llvm-project/commit/0c1c0d53931636331b59a03ed08f70936835399c
DIFF: https://github.com/llvm/llvm-project/commit/0c1c0d53931636331b59a03ed08f70936835399c.diff
LOG: [MLIR] Add patterns to bubble-up pack and push-down unpack through collapse/expand shape ops (#85297)
Add DataLayoutPropagation patterns to bubble-up pack and push-down
unpack through collapse/expand shape ops.
---------
Co-authored-by: Quinn Dawkins <quinn.dawkins at gmail.com>
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
mlir/test/Dialect/Linalg/data-layout-propagation.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index 5ceb85e7d9903b..7fd88dec71d491 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include <optional>
@@ -552,6 +553,305 @@ class BubbleUpPackThroughPadOp final : public OpRewritePattern<tensor::PackOp> {
ControlPropagationFn controlFn;
};
+/// Project dimsPos to the inner-most non-unit dim pos with reassocIndices.
+///
+/// For example, given dimsPos [0, 2], reassocIndices [[0, 1], [2, 3]], and
+/// targetShape [16, 16, 32, 1], it returns [1, 2]. Because for pos 0, the
+/// inner-most projected dim in pos [0, 1] is 1. And for pos 2, the inner-most
+/// non-unit projected dims in pos [2, 3] is 2.
+///
+/// If all candidates in a reassociation are unit dims, it chooses the
+/// inner-most dim pos.
+static SmallVector<int64_t>
+projectToInnerMostNonUnitDimsPos(ArrayRef<int64_t> dimsPos,
+ ArrayRef<ReassociationIndices> reassocIndices,
+ ArrayRef<int64_t> targetShape) {
+ SmallVector<int64_t> projectedDimsPos;
+ for (auto pos : dimsPos) {
+ // In the case all dims are unit, this will return the inner-most one.
+ int64_t projectedPos = reassocIndices[pos].back();
+ for (auto i : llvm::reverse(reassocIndices[pos])) {
+ int64_t dim = targetShape[i];
+ if (dim > 1 || ShapedType::isDynamic(dim)) {
+ projectedPos = i;
+ break;
+ }
+ }
+ projectedDimsPos.push_back(projectedPos);
+ }
+ return projectedDimsPos;
+}
+
+/// Check if all dims in dimsPos are divisible by the corresponding tile sizes.
+static bool isDimsDivisibleByTileSizes(ArrayRef<int64_t> dimsPos,
+ ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> tileSizes) {
+ for (auto [pos, tileSize] : llvm::zip_equal(dimsPos, tileSizes)) {
+ int64_t dim = shape[pos];
+ if (ShapedType::isDynamic(dim) || (dim % tileSize) != 0)
+ return false;
+ }
+ return true;
+}
+
+/// Permutate the reassociation indices and reindex them in the sequence order.
+/// Returns the next dim pos in the sequence.
+///
+/// For example, given reassocIndices [[0, 1], [2]] and permutation [1, 0], it
+/// applies the permutation to get [[2], [0, 1]] and reindexes the indices into
+/// [[0], [1, 2]].
+static int64_t applyPermutationAndReindexReassoc(
+ SmallVector<ReassociationIndices> &reassocIndices,
+ ArrayRef<int64_t> permutation) {
+ applyPermutationToVector<ReassociationIndices>(reassocIndices, permutation);
+ int64_t nextPos = 0;
+ for (ReassociationIndices &indices : reassocIndices) {
+ for (auto &index : indices) {
+ index = nextPos;
+ nextPos += 1;
+ }
+ }
+ return nextPos;
+}
+
+/// Bubble up pack op through collapse shape op when the packed dims can be
+/// projected to the dims before collapsing. This is possible when the inner
+/// tile sizes can divide the projected dims.
+///
+/// For example:
+///
+/// %collapsed = tensor.collapse_shape %in [[0, 1], 2]
+/// : tensor<?x16x4xf32> into tensor<?x4xf32>
+/// %pack = tensor.pack %collapsed outer_dims_perm = [0, 1]
+/// inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %empty
+/// : tensor<?x4xf32> -> tensor<?x4x8x1xf32>
+///
+/// can be transformed into:
+///
+/// %pack = tensor.pack %in outer_dims_perm = [1, 2]
+/// inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %empty
+/// : tensor<?x16x4xf32> -> tensor<?x2x4x8x1xf32>
+/// %collapsed = tensor.collapse_shape %pack [[0, 1], 2, 3, 4]
+/// : tensor<?x2x4x8x1xf32> into tensor<?x4x8x1>
+static LogicalResult
+bubbleUpPackOpThroughCollapseShape(tensor::CollapseShapeOp collapseOp,
+ tensor::PackOp packOp,
+ PatternRewriter &rewriter) {
+ SmallVector<int64_t> innerTileSizes = packOp.getStaticTiles();
+ ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
+ ArrayRef<int64_t> outerDimsPerm = packOp.getOuterDimsPerm();
+
+ ArrayRef<int64_t> srcShape = collapseOp.getSrcType().getShape();
+ SmallVector<ReassociationIndices> reassocIndices =
+ collapseOp.getReassociationIndices();
+ // Project inner tile pos to the dim pos before collapsing. For example, if
+ // dims [x, y] is collapsed into [z], packing on dim z can be projected back
+ // to pack on dim y.
+ //
+ // Project to inner-most non-unit dims to increase the chance that they can be
+ // divided by the inner tile sizes. This is correct because for [..., x, 1],
+ // packing on dim 1 is equivalent to packing on dim x.
+ SmallVector<int64_t> projectedInnerDimsPos =
+ projectToInnerMostNonUnitDimsPos(innerDimsPos, reassocIndices, srcShape);
+
+ if (!isDimsDivisibleByTileSizes(projectedInnerDimsPos, srcShape,
+ innerTileSizes)) {
+ return failure();
+ }
+ // Expand the outer dims permutation with the associated source dims for the
+ // new permutation after bubbling. This is because moving a collapsed dim is
+ // equivalent to moving the associated source dims together.
+ SmallVector<int64_t> newOuterDimsPerm;
+ for (auto outerPos : outerDimsPerm) {
+ newOuterDimsPerm.insert(newOuterDimsPerm.end(),
+ reassocIndices[outerPos].begin(),
+ reassocIndices[outerPos].end());
+ }
+
+ auto emptyOp = tensor::PackOp::createDestinationTensor(
+ rewriter, packOp.getLoc(), collapseOp.getSrc(), packOp.getMixedTiles(),
+ projectedInnerDimsPos, newOuterDimsPerm);
+ auto newPackOp = rewriter.create<tensor::PackOp>(
+ packOp.getLoc(), collapseOp.getSrc(), emptyOp, projectedInnerDimsPos,
+ packOp.getMixedTiles(), packOp.getPaddingValue(), newOuterDimsPerm);
+
+ SmallVector<ReassociationIndices> newReassocIndices = reassocIndices;
+ // First apply the permutation on the reassociations of the outer dims.
+ // For example given the permutation [1, 0], the reassociations [[0, 1], [2]]
+ // -> [[0], [1, 2]]
+ int64_t nextPos =
+ applyPermutationAndReindexReassoc(newReassocIndices, outerDimsPerm);
+ // Then add direct mapping for the inner tile dims.
+ for (size_t i = 0; i < innerDimsPos.size(); ++i) {
+ newReassocIndices.push_back({nextPos});
+ nextPos += 1;
+ }
+
+ auto newCollapseOp = rewriter.create<tensor::CollapseShapeOp>(
+ collapseOp.getLoc(), packOp.getType(), newPackOp, newReassocIndices);
+ rewriter.replaceOp(packOp, newCollapseOp);
+
+ return success();
+}
+
+class BubbleUpPackOpThroughReshapeOp final
+ : public OpRewritePattern<tensor::PackOp> {
+public:
+ BubbleUpPackOpThroughReshapeOp(MLIRContext *context, ControlPropagationFn fun)
+ : OpRewritePattern<tensor::PackOp>(context), controlFn(std::move(fun)) {}
+
+ LogicalResult matchAndRewrite(tensor::PackOp packOp,
+ PatternRewriter &rewriter) const override {
+ Operation *srcOp = packOp.getSource().getDefiningOp();
+ // Currently only support when the pack op is the only user.
+ if (!srcOp || !(srcOp->getNumResults() == 1) ||
+ !srcOp->getResult(0).hasOneUse()) {
+ return failure();
+ }
+ // Currently only support static inner tile sizes.
+ if (llvm::any_of(packOp.getStaticTiles(), [](int64_t size) {
+ return ShapedType::isDynamic(size);
+ })) {
+ return failure();
+ }
+
+ // User controlled propagation function.
+ if (!controlFn(srcOp))
+ return failure();
+
+ return TypeSwitch<Operation *, LogicalResult>(srcOp)
+ .Case([&](tensor::CollapseShapeOp op) {
+ return bubbleUpPackOpThroughCollapseShape(op, packOp, rewriter);
+ })
+ .Default([](Operation *) { return failure(); });
+ }
+
+private:
+ ControlPropagationFn controlFn;
+};
+
+/// Push down unpack op through expand shape op when the packed dims can be
+/// projected to the dims after expanding. This is possible when the inner tile
+/// sizes can divide the projected dims.
+///
+/// For example:
+///
+/// %unpack = tensor.unpack %in outer_dims_perm = [0, 1]
+/// inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %empty
+/// : tensor<?x32x8x8xf32> -> tensor<?x256xf32>
+/// %expanded = tensor.expand_shape %unpack [[0, 1], [2]]
+/// : tensor<?x256xf32> into tensor<?x256x256xf32>
+///
+/// can be transformed into:
+///
+/// %expanded = tensor.expand_shape %ain [[0, 1], [2], [3], [4]]
+/// : tensor<?x32x8x8xf32> into tensor<?x32x32x8x8xf32>
+/// %unpack = tensor.unpack %expanded outer_dims_perm = [0, 1, 2]
+/// inner_dims_pos = [1, 2] inner_tiles = [8, 8] into %empty
+/// : tensor<?x32x32x8x8xf32> -> tensor<?x256x256xf32>
+static LogicalResult
+pushDownUnPackOpThroughExpandShape(tensor::UnPackOp unPackOp,
+ tensor::ExpandShapeOp expandOp,
+ PatternRewriter &rewriter) {
+ SmallVector<int64_t> innerTileSizes = unPackOp.getStaticTiles();
+ ArrayRef<int64_t> innerDimsPos = unPackOp.getInnerDimsPos();
+ ArrayRef<int64_t> outerDimsPerm = unPackOp.getOuterDimsPerm();
+
+ ArrayRef<int64_t> dstShape = expandOp.getType().getShape();
+ SmallVector<ReassociationIndices> reassocIndices =
+ expandOp.getReassociationIndices();
+ // Project inner tile pos to the dim pos after expanding. For example, if dims
+ // [z] is expanded into [x, y], unpacking on dim z can be projected to unpack
+ // on dim y.
+ //
+ // Project to inner-most non-unit dims to increase the chance that they can be
+ // divided by the inner tile sizes. This is correct because for [..., x, 1],
+ // unpacking on dim 1 is equivalent to unpacking on dim x.
+ SmallVector<int64_t> projectedInnerDimsPos =
+ projectToInnerMostNonUnitDimsPos(innerDimsPos, reassocIndices, dstShape);
+
+ if (!isDimsDivisibleByTileSizes(projectedInnerDimsPos, dstShape,
+ innerTileSizes)) {
+ return failure();
+ }
+ // Expand the outer dims permutation with the associated expanded dims for the
+ // new permutation after pushing. This is because moving a source dim is
+ // equivalent to moving the associated expanded dims together.
+ SmallVector<int64_t> newOuterDimsPerm;
+ for (auto outerPos : outerDimsPerm) {
+ newOuterDimsPerm.insert(newOuterDimsPerm.end(),
+ reassocIndices[outerPos].begin(),
+ reassocIndices[outerPos].end());
+ }
+
+ SmallVector<ReassociationIndices> newReassocIndices = reassocIndices;
+ // First apply the permutation on the reassociations of the outer dims.
+ // For example given the permutation [1, 0], the reassociations [[0, 1], [2]]
+ // -> [[0], [1, 2]]
+ int64_t nextPos =
+ applyPermutationAndReindexReassoc(newReassocIndices, outerDimsPerm);
+ // Then add direct mapping for the inner tile dims.
+ for (size_t i = 0; i < innerDimsPos.size(); ++i) {
+ newReassocIndices.push_back({nextPos});
+ nextPos += 1;
+ }
+
+ RankedTensorType newExpandType =
+ tensor::PackOp::inferPackedType(expandOp.getType(), innerTileSizes,
+ projectedInnerDimsPos, newOuterDimsPerm);
+ auto newExpandOp = rewriter.create<tensor::ExpandShapeOp>(
+ expandOp.getLoc(), newExpandType, unPackOp.getSource(),
+ newReassocIndices);
+
+ auto emptyOp = tensor::UnPackOp::createDestinationTensor(
+ rewriter, unPackOp.getLoc(), newExpandOp, unPackOp.getMixedTiles(),
+ projectedInnerDimsPos, newOuterDimsPerm);
+ auto newUnPackOp = rewriter.create<tensor::UnPackOp>(
+ unPackOp.getLoc(), newExpandOp.getResult(), emptyOp,
+ projectedInnerDimsPos, unPackOp.getMixedTiles(), newOuterDimsPerm);
+ rewriter.replaceOp(expandOp, newUnPackOp);
+
+ return success();
+}
+
+class PushDownUnPackOpThroughReshapeOp final
+ : public OpRewritePattern<tensor::UnPackOp> {
+public:
+ PushDownUnPackOpThroughReshapeOp(MLIRContext *context,
+ ControlPropagationFn fun)
+ : OpRewritePattern<tensor::UnPackOp>(context), controlFn(std::move(fun)) {
+ }
+
+ LogicalResult matchAndRewrite(tensor::UnPackOp unPackOp,
+ PatternRewriter &rewriter) const override {
+ Value result = unPackOp.getResult();
+ // Currently only support unpack op with the single user.
+ if (!result.hasOneUse()) {
+ return failure();
+ }
+ // Currently only support static inner tile sizes.
+ if (llvm::any_of(unPackOp.getStaticTiles(), [](int64_t size) {
+ return ShapedType::isDynamic(size);
+ })) {
+ return failure();
+ }
+
+ Operation *consumerOp = *result.user_begin();
+ // User controlled propagation function.
+ if (!controlFn(consumerOp))
+ return failure();
+
+ return TypeSwitch<Operation *, LogicalResult>(consumerOp)
+ .Case([&](tensor::ExpandShapeOp op) {
+ return pushDownUnPackOpThroughExpandShape(unPackOp, op, rewriter);
+ })
+ .Default([](Operation *) { return failure(); });
+ }
+
+private:
+ ControlPropagationFn controlFn;
+};
+
// TODO: Relax this restriction. We should unpack a generic op also
// in the presence of multiple unpack ops as producers.
/// Return the unpacked operand, if present, for the current generic op.
@@ -774,6 +1074,7 @@ void mlir::linalg::populateDataLayoutPropagationPatterns(
const ControlPropagationFn &controlPackUnPackPropagation) {
patterns
.insert<BubbleUpPackOpThroughGenericOpPattern, BubbleUpPackThroughPadOp,
- PushDownUnPackOpThroughGenericOp, PushDownUnPackThroughPadOp>(
+ BubbleUpPackOpThroughReshapeOp, PushDownUnPackOpThroughGenericOp,
+ PushDownUnPackThroughPadOp, PushDownUnPackOpThroughReshapeOp>(
patterns.getContext(), controlPackUnPackPropagation);
}
diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index e036695a2ac9fd..79d61ab757e327 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -905,3 +905,163 @@ func.func @unpack_
diff erent_destination_shape(%arg0: tensor<1x1x1080x1920x16xi32
// CHECK-SAME: inner_dims_pos = [0] inner_tiles = [16]
// CHECK-SAME: into %[[UNPACK_NEW_DEST]]
// CHECK: return %[[UNPACK]] : tensor<16x540x960xi32>
+
+// -----
+
+func.func @bubble_up_pack_through_collapse(%1: tensor<?x16x4xf32>, %dim : index) -> tensor<?x4x8x1xf32> {
+ %collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<?x16x4xf32> into tensor<?x4xf32>
+ %2 = tensor.empty(%dim) : tensor<?x4x8x1xf32>
+ %pack = tensor.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %2 : tensor<?x4xf32> -> tensor<?x4x8x1xf32>
+ func.return %pack : tensor<?x4x8x1xf32>
+}
+// CHECK-LABEL: func.func @bubble_up_pack_through_collapse
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x16x4xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x2x4x8x1xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<?x16x4xf32> -> tensor<?x2x4x8x1xf32>
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1], [2], [3], [4]] : tensor<?x2x4x8x1xf32> into tensor<?x4x8x1xf32>
+// CHECK: return %[[COLLAPSED]] : tensor<?x4x8x1xf32>
+
+// -----
+
+func.func @bubble_up_permuted_pack_through_collapse(%1: tensor<4x192x16x256xf32>) -> tensor<4x32x3072x8x1xf32> {
+ %collapsed = tensor.collapse_shape %1 [[0], [1, 2], [3]] : tensor<4x192x16x256xf32> into tensor<4x3072x256xf32>
+ %2 = tensor.empty() : tensor<4x32x3072x8x1xf32>
+ %pack = tensor.pack %collapsed outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [8, 1] into %2 : tensor<4x3072x256xf32> -> tensor<4x32x3072x8x1xf32>
+ func.return %pack : tensor<4x32x3072x8x1xf32>
+}
+// CHECK-LABEL: func.func @bubble_up_permuted_pack_through_collapse
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x32x192x16x8x1xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3, 2] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<4x192x16x256xf32> -> tensor<4x32x192x16x8x1xf32>
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %pack {{\[}}[0], [1], [2, 3], [4], [5]] : tensor<4x32x192x16x8x1xf32> into tensor<4x32x3072x8x1xf32>
+// CHECK: return %[[COLLAPSED]] : tensor<4x32x3072x8x1xf32>
+
+// -----
+
+func.func @bubble_up_pack_through_unit_collapse(%1: tensor<1x64x1x4xf32>) -> tensor<8x4x8x1xf32> {
+ %collapsed = tensor.collapse_shape %1 [[0, 1, 2], [3]] : tensor<1x64x1x4xf32> into tensor<64x4xf32>
+ %2 = tensor.empty() : tensor<8x4x8x1xf32>
+ %pack = tensor.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %2 : tensor<64x4xf32> -> tensor<8x4x8x1xf32>
+ func.return %pack : tensor<8x4x8x1xf32>
+}
+// CHECK-LABEL: func.func @bubble_up_pack_through_unit_collapse
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<1x8x1x4x8x1xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [1, 3] inner_tiles = [8, 1] into %[[EMPTY]] : tensor<1x64x1x4xf32> -> tensor<1x8x1x4x8x1xf32>
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1, 2], [3], [4], [5]] : tensor<1x8x1x4x8x1xf32> into tensor<8x4x8x1xf32>
+// CHECK: return %[[COLLAPSED]] : tensor<8x4x8x1xf32>
+
+// -----
+
+func.func @bubble_up_pack_through_collapse_on_outer_dims(%1: tensor<?x16x4xf32>, %dim : index) -> tensor<?x1x4xf32> {
+ %collapsed = tensor.collapse_shape %1 [[0, 1], [2]] : tensor<?x16x4xf32> into tensor<?x4xf32>
+ %2 = tensor.empty(%dim) : tensor<?x1x4xf32>
+ %pack = tensor.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [4] into %2 : tensor<?x4xf32> -> tensor<?x1x4xf32>
+ func.return %pack : tensor<?x1x4xf32>
+}
+// CHECK-LABEL: func.func @bubble_up_pack_through_collapse_on_outer_dims
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x16x4xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x16x1x4xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[ARG0]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [2] inner_tiles = [4] into %[[EMPTY]] : tensor<?x16x4xf32> -> tensor<?x16x1x4xf32>
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PACK]] {{\[}}[0, 1], [2], [3]] : tensor<?x16x1x4xf32> into tensor<?x1x4xf32>
+// CHECK: return %[[COLLAPSED]] : tensor<?x1x4xf32>
+
+// -----
+
+func.func @no_bubble_up_pack_through_non_divisible_collapse(%1: tensor<3072x64x4xf32>) -> tensor<384x32x8x8xf32> {
+ %collapsed = tensor.collapse_shape %1 [[0], [1, 2]] : tensor<3072x64x4xf32> into tensor<3072x256xf32>
+ %2 = tensor.empty() : tensor<384x32x8x8xf32>
+ %pack = tensor.pack %collapsed outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %2 : tensor<3072x256xf32> -> tensor<384x32x8x8xf32>
+ func.return %pack : tensor<384x32x8x8xf32>
+}
+// CHECK-LABEL: func.func @no_bubble_up_pack_through_non_divisible_collapse
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<3072x64x4xf32> into tensor<3072x256xf32>
+// CHECK: %[[PACK:.+]] = tensor.pack %[[COLLAPSED]]
+// CHECK: return %[[PACK]] : tensor<384x32x8x8xf32>
+
+// -----
+
+func.func @push_down_unpack_through_expand(%5: tensor<?x32x8x8xf32>, %dim: index) -> tensor<?x256x256xf32> {
+ %6 = tensor.empty(%dim) : tensor<?x256xf32>
+ %unpack = tensor.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<?x32x8x8xf32> -> tensor<?x256xf32>
+ %expanded = tensor.expand_shape %unpack [[0, 1], [2]] : tensor<?x256xf32> into tensor<?x256x256xf32>
+ func.return %expanded : tensor<?x256x256xf32>
+}
+// CHECK-LABEL: func.func @push_down_unpack_through_expand
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3], [4]] : tensor<?x32x8x8xf32> into tensor<?x32x32x8x8xf32>
+// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]] : tensor<?x32x32x8x8xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x256x256xf32>
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[EXPANDED:.+]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<?x32x32x8x8xf32> -> tensor<?x256x256xf32>
+// CHECK: return %[[UNPACK]] : tensor<?x256x256xf32>
+
+// -----
+
+func.func @push_down_permuted_unpack_through_expand(%5: tensor<4x32x384x8x8xf32>) -> tensor<4x12x256x256xf32> {
+ %6 = tensor.empty() : tensor<4x3072x256xf32>
+ %unpack = tensor.unpack %5 outer_dims_perm = [0, 2, 1] inner_dims_pos = [2, 1] inner_tiles = [8, 8] into %6 : tensor<4x32x384x8x8xf32> -> tensor<4x3072x256xf32>
+ %expanded = tensor.expand_shape %unpack [[0], [1, 2], [3]] : tensor<4x3072x256xf32> into tensor<4x12x256x256xf32>
+ func.return %expanded : tensor<4x12x256x256xf32>
+}
+// CHECK-LABEL: @push_down_permuted_unpack_through_expand
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1], [2, 3], [4], [5]] : tensor<4x32x384x8x8xf32> into tensor<4x32x12x32x8x8xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x12x256x256xf32>
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[EXPANDED]] outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3, 2] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<4x32x12x32x8x8xf32> -> tensor<4x12x256x256xf32>
+// CHECK: return %[[UNPACK]] : tensor<4x12x256x256xf32>
+
+// -----
+
+func.func @push_down_unpack_through_unit_expand(%5: tensor<6x32x8x8xf32>) -> tensor<3x16x1x256xf32> {
+ %6 = tensor.empty() : tensor<48x256xf32>
+ %unpack = tensor.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<6x32x8x8xf32> -> tensor<48x256xf32>
+ %expanded = tensor.expand_shape %unpack [[0, 1, 2], [3]] : tensor<48x256xf32> into tensor<3x16x1x256xf32>
+ func.return %expanded : tensor<3x16x1x256xf32>
+}
+// CHECK-LABEL: func.func @push_down_unpack_through_unit_expand
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3], [4], [5]] : tensor<6x32x8x8xf32> into tensor<3x2x1x32x8x8xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x16x1x256xf32>
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[EXPANDED]] outer_dims_perm = [0, 1, 2, 3] inner_dims_pos = [1, 3] inner_tiles = [8, 8] into %[[EMPTY]] : tensor<3x2x1x32x8x8xf32> -> tensor<3x16x1x256xf32>
+// CHECK: return %[[UNPACK]] : tensor<3x16x1x256xf32>
+
+// -----
+
+func.func @push_down_unpack_through_expand_on_outer_dims(%5: tensor<?x32x8xf32>, %dim: index) -> tensor<?x256x256xf32> {
+ %6 = tensor.empty(%dim) : tensor<?x256xf32>
+ %unpack = tensor.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [8] into %6 : tensor<?x32x8xf32> -> tensor<?x256xf32>
+ %expanded = tensor.expand_shape %unpack [[0, 1], [2]] : tensor<?x256xf32> into tensor<?x256x256xf32>
+ func.return %expanded : tensor<?x256x256xf32>
+}
+// CHECK-LABEL: func.func @push_down_unpack_through_expand_on_outer_dims
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3]] : tensor<?x32x8xf32> into tensor<?x256x32x8xf32>
+// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]] : tensor<?x256x32x8xf32>
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%[[DIM]]) : tensor<?x256x256xf32>
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[EXPANDED:.+]] outer_dims_perm = [0, 1, 2] inner_dims_pos = [2] inner_tiles = [8] into %[[EMPTY]] : tensor<?x256x32x8xf32> -> tensor<?x256x256xf32>
+// CHECK: return %[[UNPACK]] : tensor<?x256x256xf32>
+
+// -----
+
+func.func @no_push_down_unpack_through_non_divisible_expand(%5: tensor<384x32x8x8xf32>) -> tensor<256x12x256xf32> {
+ %6 = tensor.empty() : tensor<3072x256xf32>
+ %unpack = tensor.unpack %5 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [8, 8] into %6 : tensor<384x32x8x8xf32> -> tensor<3072x256xf32>
+ %expanded = tensor.expand_shape %unpack [[0, 1], [2]] : tensor<3072x256xf32> into tensor<256x12x256xf32>
+ func.return %expanded : tensor<256x12x256xf32>
+}
+// CHECK-LABEL: func.func @no_push_down_unpack_through_non_divisible_expand
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[ARG0]]
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[UNPACK]] {{\[}}[0, 1], [2]] : tensor<3072x256xf32> into tensor<256x12x256xf32>
+// CHECK: return %[[EXPANDED]] : tensor<256x12x256xf32>
More information about the Mlir-commits
mailing list