[Mlir-commits] [mlir] [mlir][linalg] Fix UBSan division-by-zero in PackOp folding (PR #186271)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/186271

When tensor-cast folding propagates a zero constant into a dynamic tile size, `FoldTensorCastPackOp` would proceed to fold the pack into an `expand_shape` using that zero tile — causing undefined behaviour (integer division/modulo by zero, caught by UBSan as SIGFPE).

Two fixes:

1. Guard `FoldTensorCastPackOp`: bail out early if any of the resolved tile sizes is zero, preventing the invalid fold entirely.

2. Restrict the `hasZeros` check in `commonVerifierPackAndUnPackOp` to only inspect `Attribute` operands (statically-known zeros), not dynamic `Value` operands.  The verifier can only meaningfully reject zero tiles that are statically visible; dynamic zeros are an inherently runtime condition.

Add a regression test that ensures a pack with a zero tile size is not folded into `tensor.expand_shape`.

Fixes #185352

Assisted-by: Claude Code

>From 9f6a6f3c884635836200b61a23dcf878ec9c9f98 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 12 Mar 2026 14:44:06 -0700
Subject: [PATCH] [mlir][linalg] Fix UBSan division-by-zero in PackOp folding
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When tensor-cast folding propagates a zero constant into a dynamic tile
size, `FoldTensorCastPackOp` would proceed to fold the pack into an
`expand_shape` using that zero tile — causing undefined behaviour (integer
division/modulo by zero, caught by UBSan as SIGFPE).

Two fixes:

1. Guard `FoldTensorCastPackOp`: bail out early if any of the resolved
   tile sizes is zero, preventing the invalid fold entirely.

2. Restrict the `hasZeros` check in `commonVerifierPackAndUnPackOp` to
   only inspect `Attribute` operands (statically-known zeros), not dynamic
   `Value` operands.  The verifier can only meaningfully reject zero tiles
   that are statically visible; dynamic zeros are an inherently runtime
   condition.

Add a regression test that ensures a pack with a zero tile size is not
folded into `tensor.expand_shape`.

Fixes #185352

Assisted-by: Claude Code
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp       |  6 +++++-
 .../Dialect/Linalg/simplify-pack-unpack.mlir   | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index ad2909f656eea..f908827b255f2 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5112,7 +5112,9 @@ static LogicalResult commonVerifierPackAndUnPackOp(OpTy packOrUnPack) {
 
   // Return true if we have a zero-value tile.
   auto hasZeros = [&](ArrayRef<OpFoldResult> tiles) {
-    return llvm::any_of(tiles, isZeroInteger);
+    return llvm::any_of(tiles, [](OpFoldResult tile) {
+      return isa<Attribute>(tile) && isZeroInteger(tile);
+    });
   };
 
   // Verify that the source and destination are ranked types.
@@ -5997,6 +5999,8 @@ struct FoldTensorCastPackOp : public OpRewritePattern<PackOp> {
     // Get the updated mixed-tile-sizes attribute.
     SmallVector<OpFoldResult> newMixedTileSizes =
         getNewMixedTileSizes(rewriter, newResultTypes[0], op.getMixedTiles());
+    if (llvm::any_of(newMixedTileSizes, isZeroInteger))
+      return failure();
 
     // Clone op.
     // TODO: Strictly speaking, discardable attributes should be _discarded_ at
diff --git a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
index 6979770154bab..afd3a04a8996f 100644
--- a/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Linalg/simplify-pack-unpack.mlir
@@ -134,6 +134,24 @@ func.func @pack_32x1_to_16x1x1x2(%arg0 : tensor<32x1xf32>) -> tensor<16x1x1x2xf3
 
 // -----
 
+// Zero tile size: pack must not be folded into expand_shape.
+// CHECK-LABEL: func.func @pack_zero_tile_not_folded
+// CHECK-NOT:     tensor.expand_shape
+// CHECK:         linalg.pack
+func.func @pack_zero_tile_not_folded(%A: tensor<7x16xi32>) -> tensor<1x16x?x1xi32> {
+  %pad_val = arith.constant 123 : i32
+  %tile_size = arith.constant 0 : index
+  %empty = tensor.empty(%tile_size) : tensor<1x16x?x1xi32>
+  %pack = linalg.pack %A
+    padding_value(%pad_val : i32)
+    inner_dims_pos = [0, 1]
+    inner_tiles = [%tile_size, 1]
+    into %empty : tensor<7x16xi32> -> tensor<1x16x?x1xi32>
+  return %pack : tensor<1x16x?x1xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func.func @unpack_1d_to_collapse
 // CHECK-SAME:    %[[ARG0:.+]]: tensor<8x32xf32>)
 // CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<8x32xf32> into tensor<256xf32>