[Mlir-commits] [mlir] [mlir][tensor] Fix crash when canonicalizing extract_slice with a poi… (PR #223609)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

tensor.extract_slice with a ub.poison value in its dynamic offset/size list crashes the canonicalizer with a std::optional assertion:
```mlir
  %p = ub.poison : index
  %s = tensor.extract_slice %t[%p] [4] [1] : tensor<8xf32> to tensor<4xf32>
```

During canonicalization, OpWithOffsetSizesAndStridesConstantArgumentFolder calls foldDynamicIndexList, which matches constant operands with m_Constant (any op carrying the ConstantLike trait). ub.poison folds to a PoisonAttr, not an IntegerAttr, so getConstantIntValue returns std::nullopt and the unconditional dereference trips the "this->_M_is_engaged()" assertion.

Only fold the matched constant when it is an integer constant and skip non-integer constants such as poison. Since foldDynamicIndexList is shared, the guard also hardens its other callers (SCF/AMDGPU/Shard). 

Add a regression test in `mlir/test/Dialect/Tensor/canonicalize.mlir` that used to crash.

Fixes: https://github.com/llvm/llvm-project/issues/223333

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Utils/StaticValueUtils.cpp (+5-2) 
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
index 3f66d86e48991..aa62fdc01c79d 100644
--- a/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
+++ b/mlir/lib/Dialect/Utils/StaticValueUtils.cpp
@@ -449,10 +449,13 @@ LogicalResult foldDynamicIndexList(SmallVectorImpl<OpFoldResult> &ofrs,
       continue;
     Attribute attr;
     if (matchPattern(cast<Value>(ofr), m_Constant(&attr))) {
+      std::optional<int64_t> mayIntVal = getConstantIntValue(attr);
+      if (!mayIntVal)
+        continue;
       // Note: All ofrs have index type.
-      if (onlyNonNegative && *getConstantIntValue(attr) < 0)
+      if (onlyNonNegative && *mayIntVal < 0)
         continue;
-      if (onlyNonZero && *getConstantIntValue(attr) == 0)
+      if (onlyNonZero && *mayIntVal == 0)
         continue;
       ofr = attr;
       valuesChanged = true;
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).

``````````

</details>


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