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

[llvmlistbot at llvm.org]

https://github.com/prometheusfma-llvm updated https://github.com/llvm/llvm-project/pull/218141

>From 874d968b6815709773cb24969c6b491b39a23645 Mon Sep 17 00:00:00 2001
From: Prometheus <prometheus.f.ma at gmail.com>
Date: Sat, 22 Aug 2026 08:52:45 -0700
Subject: [PATCH 1/2] [MLIR][Linalg] Fix crash decomposing padded pack with
 non-unit un-tiled outer dim

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.
---
 .../Dialect/Linalg/Transforms/Transforms.cpp  | 11 +++++++++++
 mlir/test/Dialect/Linalg/decompose-pack.mlir  | 19 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 0b6d067d13a16..35d0ce1cef1da 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1172,6 +1172,17 @@ 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.
+  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

>From 792e975df63206d211787fdc8ad5ff906f6d9635 Mon Sep 17 00:00:00 2001
From: Prometheus <prometheus.f.ma at gmail.com>
Date: Mon, 24 Aug 2026 00:58:27 -0700
Subject: [PATCH 2/2] fixup! [MLIR][Linalg] Fix crash decomposing padded pack
 with non-unit un-tiled outer dim

---
 mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index 35d0ce1cef1da..aa021cbad24a7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -1175,6 +1175,8 @@ LogicalResult DecomposeOuterUnitDimsPackOpPattern::matchAndRewrite(
   // 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; })) {