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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-02-28T11:15:28Z
New Revision: b3be782c4d149d9f9fd4dbf0452db0af451bd1a2

URL: https://github.com/llvm/llvm-project/commit/b3be782c4d149d9f9fd4dbf0452db0af451bd1a2
DIFF: https://github.com/llvm/llvm-project/commit/b3be782c4d149d9f9fd4dbf0452db0af451bd1a2.diff

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

`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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/test/Dialect/Affine/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 08ce00128e481..561fcc2ee5f6a 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 attribute such as #ub.poison).
+  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
+}