[Mlir-commits] [mlir] 846ad9d - [mlir] Fix canonicalizer crash when folding poison or kDynamic-sentinel index values (#223609)

[llvmlistbot at llvm.org]

Author: Akun
Date: 2026-09-24T10:53:57+08:00
New Revision: 846ad9de67ebea162840cbcf66b605226be3f18d

URL: https://github.com/llvm/llvm-project/commit/846ad9de67ebea162840cbcf66b605226be3f18d
DIFF: https://github.com/llvm/llvm-project/commit/846ad9de67ebea162840cbcf66b605226be3f18d.diff

LOG: [mlir] Fix canonicalizer crash when folding poison or kDynamic-sentinel index values (#223609)

Fixes two canonicalizer crashes, both in foldDynamicIndexList:

* #223333: tensor.extract_slice with a ub.poison offset/size crashes
canonicalization. ub.poison is matched as a constant but is not an
IntegerAttr, so getConstantIntValue returns nullopt and the subsequent
dereference asserts.
* #223336: a memref.subview offset/size/stride that wraps to
ShapedType::kDynamic (INT64_MIN) crashes canonicalization. Folding
INT64_MIN into the static array collides with the "dynamic" marker,
leaving dynamic-looking entries with no matching dynamic operands;
getMixedStrides() then asserts.

Only fold an operand when it is an integer constant whose value is not
ShapedType::kDynamic; poison and sentinel values stay dynamic operands.
Since foldDynamicIndexList is shared, this hardens all its callers.

Adds regression tests in mlir/test/Dialect/MemRef/canonicalize.mlir and
mlir/test/Dialect/Tensor/canonicalize.mlir.

Fixes #223333, #223336.

Added: 
    

Modified: 
    mlir/lib/Dialect/Utils/StaticValueUtils.cpp
    mlir/test/Dialect/MemRef/canonicalize.mlir
    mlir/test/Dialect/Tensor/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 3f66d86e48991..b5aac62d5c3a4 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -447,16 +447,16 @@ LogicalResult foldDynamicIndexList(SmallVectorImpl<OpFoldResult> &ofrs,
   for (OpFoldResult &ofr : ofrs) {
     if (isa<Attribute>(ofr))
       continue;
-    Attribute attr;
-    if (matchPattern(cast<Value>(ofr), m_Constant(&attr))) {
-      // Note: All ofrs have index type.
-      if (onlyNonNegative && *getConstantIntValue(attr) < 0)
-        continue;
-      if (onlyNonZero && *getConstantIntValue(attr) == 0)
-        continue;
-      ofr = attr;
-      valuesChanged = true;
-    }
+    // Note: All ofrs have index type; the kDynamic sentinel stays dynamic.
+    std::optional<int64_t> intVal = getConstantIntValue(ofr);
+    if (!intVal || *intVal == ShapedType::kDynamic)
+      continue;
+    if (onlyNonNegative && *intVal < 0)
+      continue;
+    if (onlyNonZero && *intVal == 0)
+      continue;
+    ofr = getAsIndexOpFoldResult(cast<Value>(ofr).getContext(), *intVal);
+    valuesChanged = true;
   }
   return success(valuesChanged);
 }

diff  --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index cc83e56ab6754..89efdc31f26e1 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -1591,6 +1591,25 @@ func.func @subview_rank_reduction(%arg0: memref<1x384x384xf32>, %idx: index)
 
 // -----
 
+// Ensure memref.subview doesn't crash when an offset/size/stride value
+// overflows to ShapedType::kDynamic (INT64_MIN). The static representation
+// uses that value as the "dynamic" marker, so such a value must remain a
+// dynamic operand instead of being folded into a static entry.
+// CHECK-LABEL: func @subview_wrapped_stride(
+//  CHECK-SAME:     %[[ARG0:.*]]: memref<8xf32, strided<[1]>>
+func.func @subview_wrapped_stride(%mem: memref<8xf32, strided<[1]>>) -> memref<?xf32, strided<[?], offset: ?>> {
+  // CHECK: %[[MIN:.*]] = arith.constant -9223372036854775808 : index
+  // CHECK: memref.subview %[[ARG0]][%[[MIN]]] [1] [%[[MIN]]]
+  %i64max = arith.constant 9223372036854775807 : i64
+  %maxIdx = arith.index_cast %i64max : i64 to index
+  %c1 = arith.constant 1 : index
+  %wrapped = arith.addi %maxIdx, %c1 : index
+  %sub = memref.subview %mem[%wrapped] [%c1] [%wrapped] : memref<8xf32, strided<[1]>> to memref<?xf32, strided<[?], offset: ?>>
+  return %sub : memref<?xf32, strided<[?], offset: ?>>
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_double_transpose(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<1x2x3x4x5xf32>
 func.func @fold_double_transpose(%arg0: memref<1x2x3x4x5xf32>) -> memref<5x3x2x4x1xf32, strided<[1, 20, 60, 5, 120]>> {

diff  --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index d8db9d502fb99..217ce64cb9868 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -431,6 +431,20 @@ func.func @from_elements_with_poison() -> tensor<1xindex> {
 
 // -----
 
+// Ensure tensor.extract_slice with a poison offset doesn't crash. The poison
+// value is not a constant integer, so it must not be folded into a static
+// offset.
+// CHECK-LABEL: func @extract_slice_with_poison_offset
+func.func @extract_slice_with_poison_offset(%t: tensor<8xf32>) -> tensor<4xf32> {
+  // CHECK: %[[POISON:.*]] = ub.poison : index
+  // CHECK: %[[SLICE:.*]] = tensor.extract_slice %arg0[%[[POISON]]] [4] [1]
+  %p = ub.poison : index
+  %s = tensor.extract_slice %t[%p] [4] [1] : tensor<8xf32> to tensor<4xf32>
+  return %s : tensor<4xf32>
+}
+
+// -----
+
 // Ensure tensor.from_elements with a vector element type doesn't crash
 // when the elements fold to constants (DenseElementsAttr does not support
 // non-scalar element types via the Attribute overload).