[Mlir-commits] [mlir] [mlir][linalg] Avoid crash on non-dim affine expressions during pack propagation (PR #214415)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 6 00:01:56 PDT 2026
https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/214415
Data layout propagation could crash on valid IR when an outer-dimension
permutation encountered a non-dim affine expression such as `d0 + d1`.
Treat such cases as unsupported and bail out cleanly instead of asserting.
Simply replacing the assertion with `return failure()` would allow the rewrite
to fail after partially mutating the IR, violating the pattern rewriter
contract. Compute packed operand details for all inputs and outputs before
creating any new IR, ensuring failed rewrites leave the IR unchanged.
With all failure conditions handled by its callers, `packGenericOp` now returns
`GenericOp` directly.
Add regression tests covering the bubble-up and push-down paths, plus a
positive case where propagation succeeds without an outer-dimension
permutation.
Testing:
- llvm-lit mlir/test/Dialect/Linalg/data-layout-propagation.mlir
- llvm-lit mlir/test/Dialect/Linalg
>From c852496491fc3ee03784e8d41fc915358f13e5f0 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Sat, 1 Aug 2026 11:54:15 +0200
Subject: [PATCH] [mlir][linalg] Avoid crash on non-dim affine expressions
during pack propagation
---
.../Transforms/DataLayoutPropagation.cpp | 99 +++++-----
.../Linalg/data-layout-propagation.mlir | 169 ++++++++++++++++++
2 files changed, 227 insertions(+), 41 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index d36ca43a6cbb3..9f70985f107b0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -197,10 +197,11 @@ struct PackedOperandDetails {
AffineMap indexingMap;
};
-/// Helper function for getOrCreatePackedViewOfOperand that populates
-/// the details of the packedOperand that needs to be formed and also
-/// returns if the packing would require padding.
-static bool getPackedOperandDetails(
+/// Populates the packed operand details and returns whether packing requires
+/// padding. Returns failure when propagation is unsupported, such as when an
+/// outer-dimension permutation would need to remap a non-dimension,
+/// non-constant affine expression.
+static FailureOr<bool> getPackedOperandDetails(
OpBuilder &b, PackInfo packInfo, GenericOp genericOp, OpOperand *opOperand,
DenseMap<OpOperand *, PackedOperandDetails> &packedOperandMap) {
PackedOperandDetails currOperandDetails;
@@ -253,8 +254,8 @@ static bool getPackedOperandDetails(
exprs[i] = b.getAffineDimExpr(inversedOuterPerm[dimPos]);
continue;
}
- assert(isa<AffineConstantExpr>(exprs[i]) &&
- "Attempted to permute non-constant and non-affine dim expression");
+ if (!isa<AffineConstantExpr>(exprs[i]))
+ return failure();
}
// Step 2.2: Undo the transposition on `exprs` and propagate the
// transposition on the pack using outerDimsPerm.
@@ -352,12 +353,13 @@ static std::tuple<Value, AffineMap> getOrCreatePackedViewOfOperand(
/// will create a new generic op with the packed operand and the packed output
/// according to packInfo when we attempt to push down unpack or bubble up pack
/// around it. Implicitly this will only work when a packInfo can be obtained.
-/// This make sure that we are only using this function on parallel permuted
+/// This ensures that the function is only used for parallel permuted
/// dimensions.
-static FailureOr<GenericOp>
-packGenericOp(RewriterBase &rewriter, GenericOp genericOp, Value dest,
- AffineMap packedOutIndexingMap, const PackInfo &packInfo,
- bool isFoldableUnpackPack, bool poisonPaddingOk) {
+static GenericOp packGenericOp(
+ RewriterBase &rewriter, GenericOp genericOp, Value dest,
+ AffineMap packedOutIndexingMap, const PackInfo &packInfo,
+ const DenseMap<OpOperand *, PackedOperandDetails> &packedOperandMap,
+ bool isFoldableUnpackPack) {
Location loc = genericOp.getLoc();
SmallVector<Value> inputOperands;
SmallVector<Value> inputOperandsFromUnpackedSource;
@@ -367,14 +369,6 @@ packGenericOp(RewriterBase &rewriter, GenericOp genericOp, Value dest,
packOp.getInnerDimsPos() == unPackOp.getInnerDimsPos() &&
llvm::equal(packOp.getMixedTiles(), unPackOp.getMixedTiles());
};
- DenseMap<OpOperand *, PackedOperandDetails> packedOperandMap;
- bool requiresPadding = false;
- for (OpOperand *inputOperand : genericOp.getDpsInputOperands()) {
- requiresPadding |= getPackedOperandDetails(rewriter, packInfo, genericOp,
- inputOperand, packedOperandMap);
- }
- if (requiresPadding && !poisonPaddingOk)
- return failure();
for (OpOperand *inputOperand : genericOp.getDpsInputOperands()) {
auto [packedOperand, packedIndexingMap] = getOrCreatePackedViewOfOperand(
@@ -390,7 +384,7 @@ packGenericOp(RewriterBase &rewriter, GenericOp genericOp, Value dest,
indexingMaps.push_back(packedIndexingMap);
}
- // If the unpack->pack sequences can be folded, replace use the sources of
+ // If the unpack->pack sequences can be folded, use the sources of
// the unpack ops in any unpack->pack chains on the generic op operands.
if (isFoldableUnpackPack) {
inputOperands = inputOperandsFromUnpackedSource;
@@ -509,10 +503,6 @@ bubbleUpPackOpThroughGenericOp(RewriterBase &rewriter, linalg::PackOp packOp,
if (failed(packInfo))
return failure();
- // We want to move the pack not the generic.
- OpBuilder::InsertionGuard guard(rewriter);
- rewriter.setInsertionPoint(genericOp);
-
// We need to handle two cases:
// 1) The linalg.pack destination is a tensor.empty. If this is the case, we
// create a new tensor.empty to avoid breaking dominance, as we are moving the
@@ -522,23 +512,42 @@ bubbleUpPackOpThroughGenericOp(RewriterBase &rewriter, linalg::PackOp packOp,
Value packOpDest = packOp.getDest();
if (!packOpDest.hasOneUse())
return failure();
- if (auto emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>()) {
- packOpDest = tensor::EmptyOp::create(rewriter, genericOp->getLoc(),
- emptyOp.getMixedSizes(),
- emptyOp.getType().getElementType());
- } else {
+
+ tensor::EmptyOp emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>();
+ if (!emptyOp) {
DominanceInfo dom(genericOp);
if (!dom.properlyDominates(packOpDest, genericOp))
return failure();
}
- // Rebuild the indexing map for the corresponding init operand.
+ // Rebuild the indexing maps for all operands before mutating the IR.
DenseMap<OpOperand *, PackedOperandDetails> packedOperandMap;
- bool requiresPadding = getPackedOperandDetails(rewriter, *packInfo, genericOp,
- opOperand, packedOperandMap);
+ FailureOr<bool> outputRequiresPadding = getPackedOperandDetails(
+ rewriter, *packInfo, genericOp, opOperand, packedOperandMap);
+ if (failed(outputRequiresPadding))
+ return failure();
+
+ bool requiresPadding = *outputRequiresPadding;
+ for (OpOperand *inputOperand : genericOp.getDpsInputOperands()) {
+ FailureOr<bool> inputRequiresPadding = getPackedOperandDetails(
+ rewriter, *packInfo, genericOp, inputOperand, packedOperandMap);
+ if (failed(inputRequiresPadding))
+ return failure();
+ requiresPadding |= *inputRequiresPadding;
+ }
if (requiresPadding && !poisonPaddingOk)
return failure();
+ // We want to move the pack, not the generic.
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPoint(genericOp);
+
+ if (emptyOp) {
+ packOpDest = tensor::EmptyOp::create(rewriter, genericOp->getLoc(),
+ emptyOp.getMixedSizes(),
+ emptyOp.getType().getElementType());
+ }
+
auto [packedOutOperand, packedOutIndexingMap] =
getOrCreatePackedViewOfOperand(rewriter, genericOp.getLoc(), opOperand,
packedOperandMap);
@@ -556,8 +565,8 @@ bubbleUpPackOpThroughGenericOp(RewriterBase &rewriter, linalg::PackOp packOp,
// pack(unpack) isn't naively foldable because the unpack op can be from
// an arbitrary domain so we need to keep both.
return packGenericOp(rewriter, genericOp, dest, packedOutIndexingMap,
- *packInfo, /*isFoldableUnpackPack=*/false,
- poisonPaddingOk);
+ *packInfo, packedOperandMap,
+ /*isFoldableUnpackPack=*/false);
}
/// Wrapper pattern that applies bubbleUpPackOpThroughGenericOp method.
@@ -1200,11 +1209,22 @@ pushDownUnPackOpThroughGenericOp(RewriterBase &rewriter, GenericOp genericOp,
if (failed(packInfo))
return failure();
- // Rebuild the indexing map for the corresponding init operand.
+ // Rebuild the indexing maps for all operands before mutating the IR.
DenseMap<OpOperand *, PackedOperandDetails> packedOperandMap;
- bool requiresPadding =
+ FailureOr<bool> outputRequiresPadding =
getPackedOperandDetails(rewriter, *packInfo, genericOp,
genericOp.getDpsInitOperand(0), packedOperandMap);
+ if (failed(outputRequiresPadding))
+ return failure();
+
+ bool requiresPadding = *outputRequiresPadding;
+ for (OpOperand *inputOperand : genericOp.getDpsInputOperands()) {
+ FailureOr<bool> inputRequiresPadding = getPackedOperandDetails(
+ rewriter, *packInfo, genericOp, inputOperand, packedOperandMap);
+ if (failed(inputRequiresPadding))
+ return failure();
+ requiresPadding |= *inputRequiresPadding;
+ }
if (requiresPadding && !poisonPaddingOk)
return failure();
@@ -1231,12 +1251,9 @@ pushDownUnPackOpThroughGenericOp(RewriterBase &rewriter, GenericOp genericOp,
// pack(unpack) is foldable in this case. This is because in pushing down the
// unpack, by default we will populate an additional pack op after the unpack.
// This guarantees them to be foldable.
- auto maybeGenericOp =
+ GenericOp newGenericOp =
packGenericOp(rewriter, genericOp, dest, packedOutIndexingMap, *packInfo,
- /*isFoldableUnpackPack=*/true, poisonPaddingOk);
- if (failed(maybeGenericOp))
- return failure();
- GenericOp newGenericOp = *maybeGenericOp;
+ packedOperandMap, /*isFoldableUnpackPack=*/true);
Value newResult =
newGenericOp.getTiedOpResult(newGenericOp.getDpsInitOperand(0));
diff --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index af6f70637d657..1b1404168799f 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1637,3 +1637,172 @@ func.func @push_extract_through_generic_secondextract(%arg0: tensor<128x128xf32>
// CHECK-SAME: ins(%[[PAD]], %[[ARG0]]
// CHECK: %[[EXTRACT2:.+]] = tensor.extract_slice %[[GENERIC]]
// CHECK: scf.yield %[[EXTRACT2]]
+
+// -----
+
+func.func @no_propagate_pack_with_non_dim_affine_output_expr(
+ %arg0: tensor<8x8x16xf32>) -> tensor<15x4x4xf32> {
+ %init = tensor.empty() : tensor<15x16xf32>
+ %generic = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0 + d1, d2)>
+ ],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ }
+ ins(%arg0 : tensor<8x8x16xf32>)
+ outs(%init : tensor<15x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<15x16xf32>
+
+ %empty = tensor.empty() : tensor<15x4x4xf32>
+ %pack = linalg.pack %generic
+ outer_dims_perm = [0, 1]
+ inner_dims_pos = [1]
+ inner_tiles = [4]
+ into %empty
+ : tensor<15x16xf32> -> tensor<15x4x4xf32>
+
+ return %pack : tensor<15x4x4xf32>
+}
+
+// CHECK-DAG: #[[$INPUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1, d2)>
+// CHECK-LABEL: func.func @no_propagate_pack_with_non_dim_affine_output_expr
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<15x16xf32>
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$INPUT_MAP]], #[[$OUTPUT_MAP]]]
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<15x4x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[GENERIC]]
+// CHECK-SAME: outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [4]
+// CHECK-SAME: into %[[EMPTY]]
+// CHECK: return %[[PACK]] : tensor<15x4x4xf32>
+
+// -----
+
+func.func @no_propagate_pack_with_non_dim_affine_input_expr(
+ %arg0: tensor<15x16xf32>) -> tensor<4x8x8x4xf32> {
+ %init = tensor.empty() : tensor<8x8x16xf32>
+ %generic = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0 + d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+ ],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ }
+ ins(%arg0 : tensor<15x16xf32>)
+ outs(%init : tensor<8x8x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<8x8x16xf32>
+
+ %empty = tensor.empty() : tensor<4x8x8x4xf32>
+ %pack = linalg.pack %generic
+ outer_dims_perm = [2, 1, 0]
+ inner_dims_pos = [2]
+ inner_tiles = [4]
+ into %empty
+ : tensor<8x8x16xf32> -> tensor<4x8x8x4xf32>
+
+ return %pack : tensor<4x8x8x4xf32>
+}
+
+// CHECK-DAG: #[[$NON_DIM_INPUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1, d2)>
+// CHECK-DAG: #[[$IDENTITY_OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func.func @no_propagate_pack_with_non_dim_affine_input_expr
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<8x8x16xf32>
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$NON_DIM_INPUT_MAP]], #[[$IDENTITY_OUTPUT_MAP]]]
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x8x8x4xf32>
+// CHECK: %[[PACK:.+]] = linalg.pack %[[GENERIC]]
+// CHECK-SAME: outer_dims_perm = [2, 1, 0] inner_dims_pos = [2] inner_tiles = [4]
+// CHECK-SAME: into %[[EMPTY]]
+// CHECK: return %[[PACK]] : tensor<4x8x8x4xf32>
+
+// -----
+
+func.func @no_push_down_unpack_with_non_dim_affine_input_expr(
+ %packed: tensor<4x8x8x4xf32>,
+ %non_dim: tensor<15x16xf32>,
+ %init: tensor<8x8x16xf32>) -> tensor<8x8x16xf32> {
+ %empty = tensor.empty() : tensor<8x8x16xf32>
+ %unpack = linalg.unpack %packed
+ outer_dims_perm = [2, 1, 0]
+ inner_dims_pos = [2]
+ inner_tiles = [4]
+ into %empty
+ : tensor<4x8x8x4xf32> -> tensor<8x8x16xf32>
+
+ %generic = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0 + d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+ ],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ }
+ ins(%unpack, %non_dim : tensor<8x8x16xf32>, tensor<15x16xf32>)
+ outs(%init : tensor<8x8x16xf32>) {
+ ^bb0(%in: f32, %non_dim_in: f32, %out: f32):
+ %sum = arith.addf %in, %non_dim_in : f32
+ linalg.yield %sum : f32
+ } -> tensor<8x8x16xf32>
+
+ return %generic : tensor<8x8x16xf32>
+}
+
+// CHECK-DAG: #[[$PUSH_IDENTITY_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG: #[[$PUSH_NON_DIM_INPUT_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 + d1, d2)>
+// CHECK-LABEL: func.func @no_push_down_unpack_with_non_dim_affine_input_expr
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<8x8x16xf32>
+// CHECK: %[[UNPACK:.+]] = linalg.unpack
+// CHECK-SAME: %{{.+}} outer_dims_perm = [2, 1, 0]
+// CHECK-SAME: inner_dims_pos = [2] inner_tiles = [4]
+// CHECK-SAME: into %[[EMPTY]]
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$PUSH_IDENTITY_MAP]], #[[$PUSH_NON_DIM_INPUT_MAP]], #[[$PUSH_IDENTITY_MAP]]]
+// CHECK-SAME: ins(%[[UNPACK]]
+// CHECK: return %[[GENERIC]] : tensor<8x8x16xf32>
+
+// -----
+
+func.func @propagate_pack_with_non_dim_affine_expr_without_outer_perm(
+ %arg0: tensor<8x8x16xf32>) -> tensor<15x4x4xf32> {
+ %init = tensor.empty() : tensor<15x16xf32>
+ %generic = linalg.generic {
+ indexing_maps = [
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0 + d1, d2)>
+ ],
+ iterator_types = ["parallel", "parallel", "parallel"]
+ }
+ ins(%arg0 : tensor<8x8x16xf32>)
+ outs(%init : tensor<15x16xf32>) {
+ ^bb0(%in: f32, %out: f32):
+ linalg.yield %in : f32
+ } -> tensor<15x16xf32>
+
+ %empty = tensor.empty() : tensor<15x4x4xf32>
+ %pack = linalg.pack %generic
+ inner_dims_pos = [1]
+ inner_tiles = [4]
+ into %empty
+ : tensor<15x16xf32> -> tensor<15x4x4xf32>
+
+ return %pack : tensor<15x4x4xf32>
+}
+
+// CHECK-DAG: #[[$PACKED_INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG: #[[$PACKED_OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0 + d1, d2, d3)>
+// CHECK-LABEL: func.func @propagate_pack_with_non_dim_affine_expr_without_outer_perm
+// CHECK-DAG: %[[OUTPUT_EMPTY:.+]] = tensor.empty() : tensor<15x4x4xf32>
+// CHECK-DAG: %[[INPUT_EMPTY:.+]] = tensor.empty() : tensor<8x8x4x4xf32>
+// CHECK: %[[INPUT_PACK:[^ ]+]] = linalg.pack %{{[^ ]+}}
+// CHECK-SAME: inner_dims_pos = [2] inner_tiles = [4]
+// CHECK-SAME: into %[[INPUT_EMPTY]]
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: indexing_maps = [#[[$PACKED_INPUT_MAP]], #[[$PACKED_OUTPUT_MAP]]]
+// CHECK-SAME: ins(%[[INPUT_PACK]] : tensor<8x8x4x4xf32>)
+// CHECK-SAME: outs(%[[OUTPUT_EMPTY]] : tensor<15x4x4xf32>)
+// CHECK: return %[[GENERIC]] : tensor<15x4x4xf32>
More information about the Mlir-commits
mailing list