[mlir][affine] Fix crash in linearize/delinearize fold when basis is ub.poison

[joker-eph]

foldCstValueToCstAttrBasis iterates the folded dynamic basis values and erases any operand whose folded attribute is non-null (i.e., was constant- folded). When an operand folds to ub.PoisonAttr, the attribute is non-null so the operand was erased from the dynamic operand list. However getConstantIntValue on the corresponding OpFoldResult in mixedBasis returns std::nullopt for poison (it is not an integer constant), so the position was left as ShapedType::kDynamic in the returned static basis.

This left the op in an inconsistent state: the static basis claimed one more dynamic entry than actually existed. A subsequent call to getMixedBasis() triggered the assertion inside getMixedValues.

Fix this by skipping poison attributes in the erasure loop, treating them like non-constant values. This keeps the dynamic operand and its matching kDynamic entry in the static basis consistent.

Fixes #179942

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

foldCstValueToCstAttrBasis iterates the folded dynamic basis values and erases any operand whose folded attribute is non-null (i.e., was constant- folded). When an operand folds to ub.PoisonAttr, the attribute is non-null so the operand was erased from the dynamic operand list. However getConstantIntValue on the corresponding OpFoldResult in mixedBasis returns std::nullopt for poison (it is not an integer constant), so the position was left as ShapedType::kDynamic in the returned static basis.

This left the op in an inconsistent state: the static basis claimed one more dynamic entry than actually existed. A subsequent call to getMixedBasis() triggered the assertion inside getMixedValues.

Fix this by skipping poison attributes in the erasure loop, treating them like non-constant values. This keeps the dynamic operand and its matching kDynamic entry in the static basis consistent.

Fixes #179942

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+6-2)
• (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+12)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a5e4b1c4d0a6f..5fab05cb6d5c4 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4932,8 +4932,12 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                            MutableOperandRange mutableDynamicBasis,
                            ArrayRef<Attribute> dynamicBasis) {
   uint64_t dynamicBasisIndex = 0;
-  for (OpFoldResult basis : dynamicBasis) {
-    if (basis) {
+  for (Attribute basis : dynamicBasis) {
+    // Skip poison values: they don't have a concrete integer value, so erasing
+    // them from the dynamic operands would create an inconsistency between
+    // the static basis (which would still hold kDynamic) and the dynamic
+    // operand list (which would be one element shorter).
+    if (basis && !isa<ub::PoisonAttr>(basis)) {
       mutableDynamicBasis.erase(dynamicBasisIndex);
     } else {
       ++dynamicBasisIndex;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..0ffb081df1a42 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1691,6 +1691,18 @@ func.func @linearize_dont_fold_dynamic_basis(%arg0: index) -> index {
 
 // -----
 
+// Folding a linearize_index with a ub.poison basis must not crash.
+// CHECK-LABEL: @linearize_dont_fold_poison_basis
+// CHECK: %[[RET:.+]] = affine.linearize_index
+// CHECK: return %[[RET]]
+func.func @linearize_dont_fold_poison_basis(%arg0: index, %arg1: index) -> index {
+  %poison = ub.poison : index
+  %ret = affine.linearize_index [%arg0, %arg1] by (%poison) : index
+  return %ret : index
+}
+
+// -----
+
 // CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_exact(
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,

[llvmbot]

@llvm/pr-subscribers-mlir-affine

Author: Mehdi Amini (joker-eph)

Changes

foldCstValueToCstAttrBasis iterates the folded dynamic basis values and erases any operand whose folded attribute is non-null (i.e., was constant- folded). When an operand folds to ub.PoisonAttr, the attribute is non-null so the operand was erased from the dynamic operand list. However getConstantIntValue on the corresponding OpFoldResult in mixedBasis returns std::nullopt for poison (it is not an integer constant), so the position was left as ShapedType::kDynamic in the returned static basis.

This left the op in an inconsistent state: the static basis claimed one more dynamic entry than actually existed. A subsequent call to getMixedBasis() triggered the assertion inside getMixedValues.

Fix this by skipping poison attributes in the erasure loop, treating them like non-constant values. This keeps the dynamic operand and its matching kDynamic entry in the static basis consistent.

Fixes #179942

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+6-2)
• (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+12)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a5e4b1c4d0a6f..5fab05cb6d5c4 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4932,8 +4932,12 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                            MutableOperandRange mutableDynamicBasis,
                            ArrayRef<Attribute> dynamicBasis) {
   uint64_t dynamicBasisIndex = 0;
-  for (OpFoldResult basis : dynamicBasis) {
-    if (basis) {
+  for (Attribute basis : dynamicBasis) {
+    // Skip poison values: they don't have a concrete integer value, so erasing
+    // them from the dynamic operands would create an inconsistency between
+    // the static basis (which would still hold kDynamic) and the dynamic
+    // operand list (which would be one element shorter).
+    if (basis && !isa<ub::PoisonAttr>(basis)) {
       mutableDynamicBasis.erase(dynamicBasisIndex);
     } else {
       ++dynamicBasisIndex;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..0ffb081df1a42 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1691,6 +1691,18 @@ func.func @linearize_dont_fold_dynamic_basis(%arg0: index) -> index {
 
 // -----
 
+// Folding a linearize_index with a ub.poison basis must not crash.
+// CHECK-LABEL: @linearize_dont_fold_poison_basis
+// CHECK: %[[RET:.+]] = affine.linearize_index
+// CHECK: return %[[RET]]
+func.func @linearize_dont_fold_poison_basis(%arg0: index, %arg1: index) -> index {
+  %poison = ub.poison : index
+  %ret = affine.linearize_index [%arg0, %arg1] by (%poison) : index
+  return %ret : index
+}
+
+// -----
+
 // CHECK-LABEL: func @cancel_delinearize_linearize_disjoint_exact(
 //  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9]+]]: index,
 //  CHECK-SAME:     %[[ARG1:[a-zA-Z0-9]+]]: index,

[jpienaar]

This makes sense to me (its conservative but results in it being consistent).

[joker-eph]

Actually I fixed it in 9882590 already, this bug is a duplicate (processed in parallel)