[Mlir-commits] [mlir] [mlir][affine] Fix crash in linearize_index fold when basis is ub.poison (PR #183650)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/183650

>From 6bf63a20653462ec7dc95259d9167d3d0341cb4d Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Feb 2026 11:40:37 -0800
Subject: [PATCH] [mlir][affine] Fix crash in linearize_index 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 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 #179265
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  8 ++++++--
 mlir/test/Dialect/Affine/canonicalize.mlir | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a5e4b1c4d0a6f..5ea694104979a 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4932,8 +4932,12 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                            MutableOperandRange mutableDynamicBasis,
                            ArrayRef<Attribute> dynamicBasis) {
   uint64_t dynamicBasisIndex = 0;
-  for (OpFoldResult basis : dynamicBasis) {
-    if (basis) {
+  for (Attribute basis : dynamicBasis) {
+    // Skip poison values: they don't have a concrete integer value, so erasing
+    // 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, IndexAttr>(basis)) {
       mutableDynamicBasis.erase(dynamicBasisIndex);
     } else {
       ++dynamicBasisIndex;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..1a1e549346a28 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2401,3 +2401,17 @@ func.func @for_empty_body_folder_iv_yield() -> index {
   }
   return %10 : index
 }
+
+// -----
+
+// Regression test: ensure canonicalization doesn't crash when the basis of
+// affine.linearize_index is ub.poison
+// (https://github.com/llvm/llvm-project/issues/179265).
+// CHECK-LABEL: @linearize_dont_fold_poison_basis
+// CHECK-NOT: affine.linearize_index
+// CHECK: return %arg0 : index
+func.func @linearize_dont_fold_poison_basis(%arg0: index) -> index {
+  %poison = ub.poison : index
+  %ret = affine.linearize_index [%arg0] by (%poison) : index
+  return %ret : index
+}