[Mlir-commits] [mlir] [mlir][affine] Fix crash in linearize/delinearize fold when basis is ub.poison (PR #185058)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`foldCstValueToCstAttrBasis` iterates the folded dynamic basis values and erases any operand whose folded attribute is non-null (i.e., was constant- folded). When an operand folds to `ub.PoisonAttr`, the attribute is non-null so the operand was erased from the dynamic operand list. However `getConstantIntValue` on the corresponding `OpFoldResult` in `mixedBasis` returns `std::nullopt` for poison (it is not an integer constant), so the position was left as `ShapedType::kDynamic` in the returned static basis.

This left the op in an inconsistent state: the static basis claimed one more dynamic entry than actually existed. A subsequent call to `getMixedBasis()` triggered the assertion inside `getMixedValues`.

Fix this by skipping poison attributes in the erasure loop, treating them like non-constant values. This keeps the dynamic operand and its matching `kDynamic` entry in the static basis consistent.

Fixes #<!-- -->179942

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+1-1) 
- (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+12) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 17107b52d2571..29f506e25f323 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4942,7 +4942,7 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
     // them from the dynamic operands would create an inconsistency between
     // the static basis (which would still hold kDynamic) and the dynamic
     // operand list (which would be one element shorter).
-    if (basis && isa<IntegerAttr>(basis)) {
+    if (basis && !isa<ub::PoisonAttr>(basis)) {
       mutableDynamicBasis.erase(dynamicBasisIndex);
     } else {
       ++dynamicBasisIndex;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 5a0a2b004433e..15efcc1451039 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1691,6 +1691,18 @@ func.func @linearize_dont_fold_dynamic_basis(%arg0: index) -> index {
 
 // -----
 
+// Folding a linearize_index with a ub.poison basis must not crash.
+// CHECK-LABEL: @linearize_dont_fold_poison_basis
+// CHECK: %[[RET:.+]] = affine.linearize_index
+// CHECK: return %[[RET]]
+func.func @linearize_dont_fold_poison_basis(%arg0: index, %arg1: index) -> index {
+  %poison = ub.poison : index
+  %ret = affine.linearize_index [%arg0, %arg1] by (%poison) : index
+  return %ret : index
+}
+
+// -----
+
 // CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_exact(
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,

``````````

</details>


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