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

[Mehdi Amini]

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

`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

>From d1e98caa823c0b73347271a448d7bd08d7d91914 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Feb 2026 11:10:48 -0800
Subject: [PATCH] [mlir][affine] Fix crash in linearize/delinearize fold when
 basis is ub.poison

`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
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  2 +-
 mlir/test/Dialect/Affine/canonicalize.mlir | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

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,