[Mlir-commits] [mlir] [mlir][affine] Fix crash in affine.linearize_index fold with non-integer operands (PR #183713)

[Mehdi Amini]

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

>From dfac7065583046990b8ca63081ed40ba391f7049 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Feb 2026 17:45:00 -0800
Subject: [PATCH] [mlir][affine] Fix crash in affine.linearize_index fold with
 non-integer operands

The fold for affine.linearize_index checked for null attributes in the
multi-index operands but did not guard against non-null attributes that
are not IntegerAttr (e.g. ub.poison, which folds to PoisonAttr). The
unchecked cast<IntegerAttr> then crashed at runtime.

Fix this by replacing the separate null check with a combined predicate
that returns early if any multi-index operand is absent or not an
IntegerAttr.

Fixes #180106
Fixes #179405
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   |  8 ++++++--
 mlir/test/Dialect/Affine/canonicalize.mlir | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 08ce00128e481..e772e2d63731e 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5298,10 +5298,14 @@ OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
   if (getMultiIndex().size() == 1)
     return getMultiIndex().front();
 
-  if (llvm::is_contained(adaptor.getMultiIndex(), nullptr))
+  if (!adaptor.getDynamicBasis().empty())
     return nullptr;
 
-  if (!adaptor.getDynamicBasis().empty())
+  // Fold only when all multi-index operands folded to integer constants.
+  // Non-null but non-integer attributes (e.g. ub.poison) must not be cast.
+  if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute attr) {
+        return !attr || !isa<IntegerAttr>(attr);
+      }))
     return nullptr;
 
   int64_t result = 0;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1a1e549346a28..19ad8db226fa3 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2415,3 +2415,19 @@ func.func @linearize_dont_fold_poison_basis(%arg0: index) -> index {
   %ret = affine.linearize_index [%arg0] by (%poison) : index
   return %ret : index
 }
+
+// -----
+
+// Ensure affine.linearize_index fold doesn't crash with ub.poison operands.
+// ub.poison folds to a PoisonAttr (not IntegerAttr), so the fold must not
+// attempt to cast it to IntegerAttr.
+// CHECK-LABEL: func @linearize_index_with_poison
+func.func @linearize_index_with_poison() -> index {
+  // CHECK: %[[POISON:.*]] = ub.poison : index
+  // CHECK-NOT: affine.linearize_index
+  // CHECK: return %[[POISON]]
+  %0 = ub.poison : index
+  %c0 = arith.constant 0 : index
+  %1 = affine.linearize_index [%0, %c0] by (1) : index
+  return %1 : index
+}