[Mlir-commits] [mlir] [mlir][Affine] Fix crash in linearize/delinearize_index fold on non-IntegerAttr constants (PR #181144)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jueon Park (JueonPark)

<details>
<summary>Changes</summary>

Fixes #<!-- -->180106.

`AffineLinearizeIndexOp::fold` and `AffineDelinearizeIndexOp::fold` crash when an operand is a non-`IntegerAttr` constant such as `ub.poison`. The fold methods checked for `nullptr` (non-constant operands) but then performed an unchecked `cast<IntegerAttr>`, which triggers an assertion failure on attributes like `PoisonAttr`.

The fix replaces the null checks with `dyn_cast_or_null<IntegerAttr>`, which correctly bails out of folding for both null and non-IntegerAttr attributes.

In my view, linear_index ought to be an integer by the usage of `AffineLinearizeIndexOp` and `AffineLinearizeIndexOp`. However, I’ve implemented this addition as the current codebase does not appear to provide sufficient safeguards for it.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+4-2) 
- (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+30) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 871c7df2f71e5..f2b97c62eac44 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -5018,7 +5018,7 @@ AffineDelinearizeIndexOp::fold(FoldAdaptor adaptor,
     return success();
   }
 
-  if (adaptor.getLinearIndex() == nullptr)
+  if (!dyn_cast_or_null<IntegerAttr>(adaptor.getLinearIndex()))
     return failure();
 
   if (!adaptor.getDynamicBasis().empty())
@@ -5339,7 +5339,9 @@ OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
   if (getMultiIndex().size() == 1)
     return getMultiIndex().front();
 
-  if (llvm::is_contained(adaptor.getMultiIndex(), nullptr))
+  if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute attr) {
+        return !dyn_cast_or_null<IntegerAttr>(attr);
+      }))
     return nullptr;
 
   if (!adaptor.getDynamicBasis().empty())
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..abfa4e8a18144 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2401,3 +2401,33 @@ func.func @for_empty_body_folder_iv_yield() -> index {
   }
   return %10 : index
 }
+
+// -----
+
+// Verify that `affine.linearize_index` fold does not crash on non-IntegerAttr
+// constant operands such as `ub.poison`.
+
+// CHECK-LABEL: @linearize_index_poison_no_crash
+// CHECK:         %[[POISON:.*]] = ub.poison : index
+// CHECK:         return %[[POISON]]
+func.func @linearize_index_poison_no_crash() -> index {
+  %0 = index.constant 0
+  %1 = ub.poison : index
+  %2 = affine.linearize_index [%0, %1] by (1) : index
+  return %2 : index
+}
+
+// -----
+
+// Verify that `affine.delinearize_index` fold does not crash on non-IntegerAttr
+// constant operands such as `ub.poison`.
+
+// CHECK-LABEL: @delinearize_index_poison_no_crash
+// CHECK:         %[[POISON:.*]] = ub.poison : index
+// CHECK:         %[[DELIN:.*]]:2 = affine.delinearize_index %[[POISON]] into (2, 3) : index, index
+// CHECK:         return %[[DELIN]]#0, %[[DELIN]]#1
+func.func @delinearize_index_poison_no_crash() -> (index, index) {
+  %0 = ub.poison : index
+  %1:2 = affine.delinearize_index %0 into (2, 3) : index, index
+  return %1#0, %1#1 : index, index
+}

``````````

</details>


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