[Mlir-commits] [mlir] [mlir][linalg] Fix crash in FoldTensorCastUnPackOp with dynamic non-constant tile size (PR #189071)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/189071

>From 78d86686d214c37ce74eeb21c8f9619b8b960451 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Mar 2026 10:46:30 -0700
Subject: [PATCH] [mlir][linalg] Fix crash in FoldTensorCastUnPackOp with
 dynamic non-constant tile size

`getNewMixedTileSizes` was calling `getConstantIntValue(tile).value()`
unconditionally when a tile OpFoldResult was a dynamic SSA Value (not an
Attribute). When the tile is a runtime-dynamic Value that is not a
compile-time constant, `getConstantIntValue` returns `std::nullopt` and
`.value()` aborts with `bad_optional_access`.

The fix replaces the assertion with the correct behavior: when the packed
type statically provides `dimSize`, we can use that value directly as the
new tile size, regardless of whether the original tile Value happens to be
constant.

Fixes #187975

Assisted-by: Claude Code
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp   |  6 +-
 mlir/test/Dialect/Linalg/canonicalize.mlir | 65 ++++++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index f9c8589683ba7..8b588a97b4363 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5037,10 +5037,10 @@ getNewMixedTileSizes(PatternRewriter &rewriter, Type newPackedTy,
       // Already a constant
       newMixedTileSizes.push_back(tile);
     } else {
-      assert(getConstantIntValue(tile).value() == dimSize &&
-             "tile size and dim size don't match!");
+      // The tile is a dynamic Value. Since the packed type tells us the
+      // dimension is statically dimSize, we can use that static value directly.
       newMixedTileSizes.push_back(
-          (rewriter.getIntegerAttr(rewriter.getIndexType(), dimSize)));
+          rewriter.getIntegerAttr(rewriter.getIndexType(), dimSize));
     }
   }
 
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 0c5a1c6108ae3..ae74bb2ed9e30 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1952,6 +1952,71 @@ func.func @fold_cast_unpack_dynamic_tile_size(
 
 // -----
 
+// Regression test: FoldTensorCastUnPackOp must not crash when a tile size is a
+// dynamic (non-constant) SSA value and the cast makes the packed dim static.
+// The static dim value from the cast should be used as the new tile size.
+// CHECK-LABEL: func.func @fold_cast_unpack_nonconstant_dynamic_tile(
+// CHECK-SAME:     %[[SRC:.*]]: tensor<1x3x8x1xi32>,
+// CHECK-SAME:     %[[TILE:.*]]: index,
+// CHECK-SAME:     %[[DEST:.*]]: tensor<7x3xi32>) -> tensor<7x3xi32> {
+// CHECK:          %[[RES:.*]] = linalg.unpack %[[SRC]] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %[[DEST]] : tensor<1x3x8x1xi32> -> tensor<7x3xi32>
+// CHECK:          return %[[RES]] : tensor<7x3xi32>
+func.func @fold_cast_unpack_nonconstant_dynamic_tile(
+  %src: tensor<1x3x8x1xi32>,
+  %tile_size: index,
+  %dest: tensor<7x3xi32>) -> tensor<7x3xi32> {
+    %cast = tensor.cast %src : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
+    %unpack = linalg.unpack %cast
+      inner_dims_pos = [0, 1]
+      inner_tiles = [%tile_size, 1]
+      into %dest : tensor<?x3x?x1xi32> -> tensor<7x3xi32>
+    return %unpack : tensor<7x3xi32>
+}
+
+// -----
+
+// When the cast reveals a static dim but the tile is a constant SSA value that
+// doesn't match.
+// CHECK-LABEL: func.func @fold_cast_unpack_constant_tile_mismatch(
+// CHECK:          linalg.unpack %[[CAST]] inner_dims_pos = [0, 1] inner_tiles = [8, 1]
+func.func @fold_cast_unpack_constant_tile_mismatch(
+  %src: tensor<1x3x8x1xi32>,
+  %dest: tensor<7x3xi32>) -> tensor<7x3xi32> {
+    %cast = tensor.cast %src : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
+    %c4 = arith.constant 4 : index
+    %unpack = linalg.unpack %cast
+      inner_dims_pos = [0, 1]
+      inner_tiles = [%c4, 1]
+      into %dest : tensor<?x3x?x1xi32> -> tensor<7x3xi32>
+    return %unpack : tensor<7x3xi32>
+}
+
+// -----
+
+// Regression test: FoldTensorCastPackOp must not crash when a tile size is a
+// dynamic (non-constant) SSA value and the cast makes the packed dim static.
+// The static dim value from the cast should be used as the new tile size.
+// CHECK-LABEL: func.func @fold_cast_pack_nonconstant_dynamic_tile(
+// CHECK-SAME:     %[[SRC:.*]]: tensor<8x3xi32>,
+// CHECK-SAME:     %[[TILE:.*]]: index,
+// CHECK-SAME:     %[[DEST:.*]]: tensor<1x3x8x1xi32>) -> tensor<1x3x8x1xi32> {
+// CHECK:          %[[RES:.*]] = linalg.pack %[[SRC]] inner_dims_pos = [0, 1] inner_tiles = [8, 1] into %[[DEST]] : tensor<8x3xi32> -> tensor<1x3x8x1xi32>
+// CHECK:          return %[[RES]] : tensor<1x3x8x1xi32>
+func.func @fold_cast_pack_nonconstant_dynamic_tile(
+  %src: tensor<8x3xi32>,
+  %tile_size: index,
+  %dest: tensor<1x3x8x1xi32>) -> tensor<1x3x8x1xi32> {
+    %cast = tensor.cast %dest : tensor<1x3x8x1xi32> to tensor<?x3x?x1xi32>
+    %pack = linalg.pack %src
+      inner_dims_pos = [0, 1]
+      inner_tiles = [%tile_size, 1]
+      into %cast : tensor<8x3xi32> -> tensor<?x3x?x1xi32>
+    %res = tensor.cast %pack : tensor<?x3x?x1xi32> to tensor<1x3x8x1xi32>
+    return %res : tensor<1x3x8x1xi32>
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // linalg.unpack + tensor.extract_slice
 //===----------------------------------------------------------------------===//