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

[Mehdi Amini]

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

`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

>From 598b9f353d6154330f0ca3027ba96e315f7d05b5 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 | 23 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index e9698365765e7..f3c88379cbade 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5039,10 +5039,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..27b52d0a0ca29 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1952,6 +1952,29 @@ 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>
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // linalg.unpack + tensor.extract_slice
 //===----------------------------------------------------------------------===//