[Mlir-commits] [mlir] [MLIR][Linalg] Prevent PackOp canonicalization crash on zero tile factor (PR #185624)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Ayush Kumar Gaur (Ayush3941)

<details>
<summary>Changes</summary>

What this Pr does: Prevent linalg.pack canonicalization from crashing when --inline --sccp folds a tile factor to zero.

Why it happened: Zero tile factors are invalid, but the pass should not crash before the verifier emits the normal invalid zero tile factor diagnostic.

whats the Fix: Add a guard in paddingIsNotNeeded() to bail out when any static inner tile is zero, avoiding the bad path into requirePaddingValue().

Fixes  #<!-- -->185352

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+3) 
- (added) mlir/test/Dialect/Linalg/pack-zero-tile-after-sccp-invalid.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index ad2909f656eea..da84c2cc1d16e 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5806,6 +5806,9 @@ static bool paddingIsNotNeeded(PackOp op) {
   if (llvm::any_of(op.getInnerDimsPos(),
                    [&](int64_t pos) { return srcType.isDynamicDim(pos); }))
     return false;
+  if (llvm::any_of(op.getStaticInnerTiles(),
+                   [](int64_t tile) { return tile == 0; }))
+    return false;
   if (ShapedType::isDynamicShape(op.getStaticInnerTiles()))
     return false;
   return !PackOp::requirePaddingValue(
diff --git a/mlir/test/Dialect/Linalg/pack-zero-tile-after-sccp-invalid.mlir b/mlir/test/Dialect/Linalg/pack-zero-tile-after-sccp-invalid.mlir
new file mode 100644
index 0000000000000..9913cd7663743
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/pack-zero-tile-after-sccp-invalid.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s -split-input-file -verify-diagnostics
+
+func.func @get_tile_size() -> index {
+  %c0 = arith.constant 0 : index
+  return %c0 : index
+}
+
+func.func private @use(%A: tensor<?x16x?x1xi32>)
+
+func.func @pack(%A: tensor<7x16xi32>) {
+  %c1 = arith.constant 1 : index
+  %pad_val = arith.constant 123 : i32
+  %tile_size = func.call @get_tile_size() : () -> index
+  %empty = tensor.empty(%c1, %tile_size) : tensor<?x16x?x1xi32>
+  // expected-error @below {{invalid zero tile factor}}
+  %pack = linalg.pack %A
+    padding_value(%pad_val : i32)
+    inner_dims_pos = [0, 1]
+    inner_tiles = [%tile_size, 1]
+    into %empty : tensor<7x16xi32> -> tensor<?x16x?x1xi32>
+  func.call @use(%pack) : (tensor<?x16x?x1xi32>) -> ()
+  return
+}

``````````

</details>


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