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

[Mehdi Amini]

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

>From 71565da7dc7d709e614214e485fbb2812ce0f109 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Feb 2026 11:33:34 -0800
Subject: [PATCH 1/2] [mlir][affine] Fix crash in linearize_index fold when
 multi-index is ub.poison

`AffineLinearizeIndexOp::fold` guarded the constant-folding path with
`llvm::is_contained(adaptor.getMultiIndex(), nullptr)`, which only catches
operands that have not been evaluated at all. When an operand folds to
`ub.PoisonAttr`, the attribute is non-null so the guard passed, and the
subsequent `cast<IntegerAttr>(indexAttr)` call crashed with an assertion
failure.

Fix by replacing the null-only check with one that requires every
multi-index attribute to be a concrete `IntegerAttr`, returning `nullptr`
for any other attribute (including null and PoisonAttr).

Fixes #178204
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  7 ++++++-
 mlir/test/Dialect/Affine/canonicalize.mlir | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 08ce00128e481..e249839f76926 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5298,7 +5298,12 @@ OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
   if (getMultiIndex().size() == 1)
     return getMultiIndex().front();
 
-  if (llvm::is_contained(adaptor.getMultiIndex(), nullptr))
+  // Return nullptr if any multi-index attribute has not been folded to a
+  // concrete integer (e.g. it is still a runtime value or has folded to a
+  // non-integer such as ub.PoisonAttr).
+  if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute a) {
+        return !isa_and_nonnull<IntegerAttr>(a);
+      }))
     return nullptr;
 
   if (!adaptor.getDynamicBasis().empty())
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1a1e549346a28..5a0a2b004433e 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2415,3 +2415,17 @@ func.func @linearize_dont_fold_poison_basis(%arg0: index) -> index {
   %ret = affine.linearize_index [%arg0] by (%poison) : index
   return %ret : index
 }
+
+// -----
+
+// Regression test: ensure constant folding doesn't crash when a multi-index
+// element of affine.linearize_index is ub.poison
+// (https://github.com/llvm/llvm-project/issues/178204).
+// CHECK-LABEL: @linearize_dont_fold_poison_index
+// CHECK: affine.linearize_index
+func.func @linearize_dont_fold_poison_index(%arg0: index) -> index {
+  %poison = ub.poison : index
+  %c4 = arith.constant 4 : index
+  %ret = affine.linearize_index [%poison, %arg0] by (%c4) : index
+  return %ret : index
+}

>From 805718b9525b161cd32009f8a4a15b7a6ca5589e Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 12:00:33 +0100
Subject: [PATCH 2/2] Update mlir/lib/Dialect/Affine/IR/AffineOps.cpp

Co-authored-by: Matthias Springer <me at matthiasspringer.de>
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index e249839f76926..561fcc2ee5f6a 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5300,7 +5300,7 @@ OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
 
   // Return nullptr if any multi-index attribute has not been folded to a
   // concrete integer (e.g. it is still a runtime value or has folded to a
-  // non-integer such as ub.PoisonAttr).
+  // non-integer attribute such as #ub.poison).
   if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute a) {
         return !isa_and_nonnull<IntegerAttr>(a);
       }))