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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

`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

---
Full diff: https://github.com/llvm/llvm-project/pull/183816.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+5-1) 
- (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+14) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 08ce00128e481..5ed7838166b6b 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5298,7 +5298,11 @@ 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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/183816