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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: prometheusfma-llvm

<details>
<summary>Changes</summary>


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.

---
Full diff: https://github.com/llvm/llvm-project/pull/218141.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+11) 
- (modified) mlir/test/Dialect/Linalg/decompose-pack.mlir (+19) 


``````````diff
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

``````````

</details>


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