shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

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

Browse Source 
`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
This commit is contained in:
Mehdi Amini 2026-02-28 12:15:28 +01:00 committed by GitHub
parent f05b705dd3
commit b3be782c4d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
View File

@ -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())

14
mlir/test/Dialect/Affine/canonicalize.mlir
View File

@ -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
}