[Mlir-commits] [mlir] ede2b5f - [MLIR][Linalg] Fix crash decomposing padded pack with non-unit un-tiled outer dim (#218141)

[llvmlistbot at llvm.org]

Author: prometheusfma-llvm
Date: 2026-08-24T19:03:04Z
New Revision: ede2b5fe3ce66534b7c68f05df4a693432582b25

URL: https://github.com/llvm/llvm-project/commit/ede2b5fe3ce66534b7c68f05df4a693432582b25
DIFF: https://github.com/llvm/llvm-project/commit/ede2b5fe3ce66534b7c68f05df4a693432582b25.diff

LOG: [MLIR][Linalg] Fix crash decomposing padded pack with non-unit un-tiled outer dim (#218141)

DecomposeOuterUnitDimsPackOpPattern only checked that the *tiled* outer
dims of a linalg.pack are all 1. When the op has a padding value,
getPackOpSourceOrPaddedSource additionally requires every outer dim
(including un-tiled ones) to be 1 and asserts otherwise. A pack with a
non-unit un-tiled outer dim therefore passed the guard but tripped the
assertion, crashing the compiler.

Bail out of the pattern via notifyMatchFailure when a padding value is
set and any outer dim is not 1, turning the crash into a graceful
no-match.

Fixes #218109.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
    mlir/test/Dialect/Linalg/decompose-pack.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 0b6d067d13a16..aa021cbad24a7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1172,6 +1172,19 @@ LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite(
         packOp, "not all outer dimensions of the result are 1s");
   }
 
+  // When a padding value is set, getPackOpSourceOrPaddedSource only supports
+  // the case where every outer dim (including un-tiled ones) is 1. Bail out
+  // instead of hitting an assertion on a non-unit un-tiled outer dim.
+  // FIXME: Handle this case by decomposing the padded pack instead of bailing
+  // out; a non-unit un-tiled outer dim should be supported here.
+  if (packOp.getPaddingValue() &&
+      llvm::any_of(packOp.getAllOuterDims(),
+                   [](int64_t dim) { return dim != 1; })) {
+    return rewriter.notifyMatchFailure(
+        packOp, "cannot decompose padded pack with a non-unit un-tiled outer "
+                "dimension");
+  }
+
   ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
   auto outerDimsPerm = packOp.getOuterDimsPerm();
 

diff  --git a/mlir/test/Dialect/Linalg/decompose-pack.mlir b/mlir/test/Dialect/Linalg/decompose-pack.mlir
index 12292ee573cee..d6780053f529d 100644
--- a/mlir/test/Dialect/Linalg/decompose-pack.mlir
+++ b/mlir/test/Dialect/Linalg/decompose-pack.mlir
@@ -356,3 +356,22 @@ func.func @negative_non_unit_tiled_outer_dim(%dest: tensor<1x126x1x1x8xf32>, %sr
 }
 // CHECK-LABEL: @negative_non_unit_tiled_outer_dim(
 // CHECK: linalg.pack
+
+// -----
+
+/// Note "2" for dim 0, a non-unit un-tiled outer dim (only dims 1 and 2 are
+/// tiled). This is not supported and must not crash.
+
+func.func @negative_non_unit_untiled_outer_dim(%src: tensor<2x2x1xi32>, %dst: tensor<2x1x1x2x2xi32>) -> tensor<2x1x1x2x2xi32> {
+  %c0 = arith.constant 0 : i32
+  %pack = linalg.pack %src
+    padding_value(%c0 : i32)
+    inner_dims_pos = [1, 2]
+    inner_tiles = [2, 2]
+    into %dst
+    : tensor<2x2x1xi32> -> tensor<2x1x1x2x2xi32>
+
+  return %pack : tensor<2x1x1x2x2xi32>
+}
+// CHECK-LABEL: @negative_non_unit_untiled_outer_dim(
+// CHECK: linalg.pack