[Mlir-commits] [mlir] [mlir][linalg] Refine pack/unpack simplification checks (NFC) (PR #209522)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Jul 14 08:47:35 PDT 2026
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/209522
Update `isPackOn1D`, which is used by both
`SimplifyPackToExpandShape` and `SimplifyUnPackToCollapseShape`:
* Rename it to `isPackOnEffectively1D` to better reflect its
functionality: the underlying pack can be multi-dimensional.
* Add checks to ensure that the unique non-unit inner tile is used to tile
the unique non-unit unpacked dimension (that was previously left as an
unchecked assumption).
* Add comments to the test file for these patterns, grouping the tests
according to the functionality/cases being tested.
>From 72293635a5bfdd02eb5903ea7b152401b8378ecb Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 14 Jul 2026 13:51:42 +0100
Subject: [PATCH] [mlir][linalg] Refine pack/unpack simplification checks (NFC)
Update `isPackOn1D`, which is used by both
`SimplifyPackToExpandShape` and `SimplifyUnPackToCollapseShape`:
* Rename it to `isPackOnEffectively1D` to better reflect its
functionality: the underlying pack can be multi-dimensional.
* Add checks to ensure that the unique non-unit inner tile is used to tile
the unique non-unit unpacked dimension (that was previously left as an
unchecked assumption).
* Add comments to the test file for these patterns, grouping the tests
according to the functionality/cases being tested.
---
.../Transforms/PackAndUnpackPatterns.cpp | 77 ++++++++++++++-----
.../Dialect/Linalg/simplify-pack-unpack.mlir | 58 +++++++++++---
2 files changed, 105 insertions(+), 30 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Linalg/Transforms/PackAndUnpackPatterns.cpp
index 993eae62535c3..c8e9f021938e5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/PackAndUnpackPatterns.cpp
@@ -24,24 +24,67 @@ static int64_t getNumGtOneDims(ArrayRef<int64_t> shape) {
shape, [](int64_t v) { return ShapedType::isDynamic(v) || v > 1; });
}
-/// Returns success() if there is only 1 dimension size in non-packed domain
-/// being greater than 1 and packing only happens on the dimension.
-/// Note: this method should only be used by pack/unpack to reshape conversion.
-/// It assumes that non-unit inner tile size must be used by the non-unit
-/// dimension.
-static LogicalResult isPackOn1D(RewriterBase &rewriter, Operation *op,
- ArrayRef<int64_t> srcShape,
- ArrayRef<int64_t> innerPackTileSize) {
- if (getNumGtOneDims(srcShape) > 1) {
+/// Returns the index of the first non-unit size in `sizes`. Returns -1 if
+/// there are no non-unit sizes.
+static int64_t getFirstNonUnitSizeIdx(ArrayRef<int64_t> sizes) {
+ const auto *it = llvm::find_if(sizes, [](int64_t dim) { return dim != 1; });
+ return (it != sizes.end()) ? std::distance(sizes.begin(), it) : -1;
+}
+
+/// Check whether `op` is effectively a 1D pack/unpack. Example:
+///
+/// %pack = linalg.pack %src
+/// inner_dims_pos = [0, 1]
+/// inner_tiles = [1, 2] into %dest
+/// : tensor<1x32xf32> -> tensor<1x16x1x2xf32>
+///
+/// Returns success() if there is:
+/// * only 1 non-unit dim in the un-packed domain,
+/// * only 1 non-unit inner tile size, and
+/// * the unique non-unit tile size is applied to the unique non-unit
+/// un-packed dim.
+template <typename PackOrUnpackOp>
+static LogicalResult isPackOnEffectively1D(RewriterBase &rewriter,
+ PackOrUnpackOp *op) {
+ // Obtain the unpacked shape.
+ auto pack = dyn_cast<linalg::PackOp>(op);
+ auto unpack = dyn_cast<linalg::UnPackOp>(op);
+
+ ArrayRef<int64_t> unpackedShape = pack ? pack->getSourceType().getShape()
+ : unpack->getDestType().getShape();
+
+ // Obtain the inner tile sizes.
+ ArrayRef<int64_t> innerTileSizes = op->getStaticInnerTiles();
+
+ // Make sure that there is exactly single non-unit unpacked dim.
+ if (getNumGtOneDims(unpackedShape) != 1) {
return rewriter.notifyMatchFailure(
- op, "expects non-packed domain to have at most one non-unit dims");
+ *op, "expects non-packed domain to have at most one non-unit dims");
}
- // Non-unit inner tile size must be used by the non-unit dimension. If not, it
- // will faill on getting reassociation maps.
- if (getNumGtOneDims(innerPackTileSize) > 1) {
+
+ // Make sure that there is at most one non-unit inner tile size.
+ auto numNonUnitInnerTiles = getNumGtOneDims(innerTileSizes);
+ if (numNonUnitInnerTiles > 1) {
return rewriter.notifyMatchFailure(
- op, "expects at most one non-unit inner tiles");
+ *op, "expects at most one non-unit inner tiles");
}
+
+ // If there are no non-unit tiles, there is nothing else to check.
+ if (numNonUnitInnerTiles == 0)
+ return success();
+
+ // Get the index of the unique non-unit unpacked dim.
+ int64_t nonUnitDimIdx = getFirstNonUnitSizeIdx(unpackedShape);
+
+ // Get the index of the dim that the unique non-unit tile is applied to.
+ int64_t nonUnitTileDestDimIdx = getFirstNonUnitSizeIdx(innerTileSizes);
+
+ // Make sure that the unique non-unit tile is applied to the unique unit dim.
+ if (nonUnitTileDestDimIdx != nonUnitDimIdx) {
+ return rewriter.notifyMatchFailure(
+ *op, "expects at most one non-unit inner tiles");
+ }
+
return success();
}
@@ -117,8 +160,7 @@ struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
ShapedType sourceType = packOp.getSourceType();
if (failed(isPackOnInnerMostDim(rewriter, packOp)) &&
- failed(isPackOn1D(rewriter, packOp, sourceType.getShape(),
- packOp.getStaticTiles())) &&
+ failed(isPackOnEffectively1D(rewriter, &packOp)) &&
!packOp.isLikePad()) {
return failure();
}
@@ -183,8 +225,7 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
ShapedType destType = unpackOp.getDestType();
if (failed(isUnpackOnInnerMostDim(rewriter, unpackOp)) &&
- failed(isPackOn1D(rewriter, unpackOp, destType.getShape(),
- unpackOp.getStaticTiles())) &&
+ failed(isPackOnEffectively1D(rewriter, &unpackOp)) &&
!unpackOp.isLikeUnPad()) {
return failure();
}
diff --git a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
index 6979770154bab..f4f3cb233d45e 100644
--- a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
@@ -1,5 +1,9 @@
// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns="test-simplify-pack-unpack-patterns" %s | FileCheck %s
+//===========================================================================//
+// Packing: 1D unpacked source
+//===========================================================================//
+
// CHECK-LABEL: func.func @single_dim_packing(
// CHECK-SAME: %[[ARG0:.+]]: tensor<256xf32>)
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] output_shape [8, 32] : tensor<256xf32> into tensor<8x32xf32>
@@ -25,18 +29,6 @@ func.func @single_dim_packing_with_padding(%arg0: tensor<255xf32>) -> tensor<8x3
// -----
-// CHECK-LABEL: func.func @single_last_inner_dim_packing(
-// CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
-// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [5, 8, 32] : tensor<5x256xf32> into tensor<5x8x32xf32>
-// CHECK: return %[[EXPANDED]] : tensor<5x8x32xf32>
-func.func @single_last_inner_dim_packing(%arg0: tensor<5x256xf32>) -> tensor<5x8x32xf32> {
- %empty = tensor.empty() : tensor<5x8x32xf32>
- %0 = linalg.pack %arg0 inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x256xf32> -> tensor<5x8x32xf32>
- return %0 : tensor<5x8x32xf32>
-}
-
-// -----
-
// CHECK-LABEL: func.func @pack_1d_with_outer_dims_perm(
// CHECK-SAME: %[[ARG0:.+]]: tensor<64xf32>)
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] output_shape [2, 32] : tensor<64xf32> into tensor<2x32xf32>
@@ -49,6 +41,22 @@ func.func @pack_1d_with_outer_dims_perm(%arg0: tensor<64xf32>) -> tensor<2x32xf3
// -----
+//===========================================================================//
+// Packing: 2D unpacked source
+//===========================================================================//
+
+// CHECK-LABEL: func.func @single_last_inner_dim_packing(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [5, 8, 32] : tensor<5x256xf32> into tensor<5x8x32xf32>
+// CHECK: return %[[EXPANDED]] : tensor<5x8x32xf32>
+func.func @single_last_inner_dim_packing(%arg0: tensor<5x256xf32>) -> tensor<5x8x32xf32> {
+ %empty = tensor.empty() : tensor<5x8x32xf32>
+ %0 = linalg.pack %arg0 inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x256xf32> -> tensor<5x8x32xf32>
+ return %0 : tensor<5x8x32xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @single_last_inner_dim_packing_with_identity_outer_dims_perm(
// CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] output_shape [5, 8, 32] : tensor<5x256xf32> into tensor<5x8x32xf32>
@@ -83,6 +91,10 @@ func.func @single_first_inner_dim_packing(%arg0: tensor<256x5xf32>) -> tensor<8x
// -----
+//===========================================================================//
+// Packing: Multi-dim unpacked source
+//===========================================================================//
+
// 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]] output_shape [1, 32, 1, 1]
@@ -134,6 +146,10 @@ func.func @pack_32x1_to_16x1x1x2(%arg0 : tensor<32x1xf32>) -> tensor<16x1x1x2xf3
// -----
+//===========================================================================//
+// Un-Packing: 1D unpacked dest
+//===========================================================================//
+
// 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>
@@ -172,6 +188,10 @@ func.func @unpack_dynamic(%arg0: tensor<?x32xf32>) -> tensor<?xf32> {
// -----
+//===========================================================================//
+// Un-Packing: 2D unpacked dest
+//===========================================================================//
+
// CHECK-LABEL: func.func @single_last_inner_dim_unpacking(
// CHECK-SAME: %[[ARG0:.+]]: tensor<5x8x32xf32>)
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x8x32xf32> into tensor<5x256xf32>
@@ -269,6 +289,13 @@ func.func @unpack_16x1x1x2_to_32x1(%arg0 : tensor<16x1x1x2xf32>) -> tensor<32x1x
// -----
+//===========================================================================//
+// Packing: Pad-like pack
+//
+// FIXME: With no single element added, there is no padding here. Improve
+// naming.
+//===========================================================================//
+
// CHECK-LABEL: func.func @pad_like_pack(
// CHECK-SAME: %[[ARG0:.+]]: tensor<32x64xf32>)
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] output_shape [1, 1, 32, 64] : tensor<32x64xf32> into tensor<1x1x32x64xf32>
@@ -333,6 +360,13 @@ func.func @pad_like_pack_with_transpose(%arg0: tensor<32x64x16xf32>) -> tensor<3
// -----
+//===========================================================================//
+// Un-Packing: Un-Pad-like unpack
+//
+// FIXME: With no single element added, there is no padding here. Improve
+// naming.
+//===========================================================================//
+
// CHECK-LABEL: func.func @unpad_like_unpack(
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x32x64xf32>)
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] : tensor<1x1x32x64xf32> into tensor<32x64xf32>
More information about the Mlir-commits
mailing list