[Mlir-commits] [mlir] [mlir][affine] Fix delinearize fold crash on non-positive basis (PR #205288)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Chennes (Chennesxu)

<details>
<summary>Changes</summary>

affine.delinearize_index verifies that no basis element is statically non-positive, but foldCstValueToCstAttrBasis could materialize a dynamic basis operand that folded to a non-positive constant (e.g. a chained delinearize result that became 0) directly into the static basis. This put the op into a state the verifier forbids and crashed the subsequent constant folding at llvm::mod (Denominator >= 1).

Fold only strictly-positive constants into the static basis: keep non-positive or non-constant dynamic operands dynamic, and preserve already-static basis elements unchanged. The helper is shared with affine.linearize_index, whose verifier does not reject a statically non-positive basis; preserving existing static elements avoids re-marking one kDynamic without a matching dynamic operand, which would desync the static and dynamic basis and crash getMixedValues.

Regression tests are added to canonicalize.mlir for both the delinearize crash (zero and negative basis) and the shared-helper consistency on linearize_index.

Fixes #<!-- -->205238

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+12-9) 
- (modified) mlir/test/Dialect/Affine/canonicalize.mlir (+31) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index f095500495f18..e6f9408ee0ca4 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -4993,11 +4993,12 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                            ArrayRef<Attribute> dynamicBasis) {
   uint64_t dynamicBasisIndex = 0;
   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<IntegerAttr>(basis)) {
+    // Only fold a dynamic basis operand into the static basis when it is a
+    // strictly positive constant. Poison and non-positive values are left
+    // dynamic: a statically non-positive basis is rejected by the verifier and
+    // would crash the constant folding below.
+    if (basis && isa<IntegerAttr>(basis) &&
+        cast<IntegerAttr>(basis).getInt() > 0) {
       mutableDynamicBasis.erase(dynamicBasisIndex);
     } else {
       ++dynamicBasisIndex;
@@ -5011,10 +5012,12 @@ foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
   SmallVector<int64_t> staticBasis;
   for (OpFoldResult basis : mixedBasis) {
     std::optional<int64_t> basisVal = getConstantIntValue(basis);
-    if (!basisVal)
-      staticBasis.push_back(ShapedType::kDynamic);
-    else
-      staticBasis.push_back(*basisVal);
+    // Keep already-static elements as-is; for dynamic operands only fold
+    // strictly positive constants in. This avoids a statically non-positive
+    // basis and keeps the kDynamic-marker count in sync with the dynamic
+    // operands retained above.
+    bool keepStatic = basisVal && (isa<Attribute>(basis) || *basisVal > 0);
+    staticBasis.push_back(keepStatic ? *basisVal : ShapedType::kDynamic);
   }
 
   return staticBasis;
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 7d236ef3c2421..8b8b8c6dc1c91 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -1522,6 +1522,23 @@ func.func @delinearize_dont_fold_constant_dynamic_basis(%arg0: index) -> (index,
 
 // -----
 
+// A dynamic basis operand that folds to a non-positive constant (zero or
+// negative) must not be folded into the static basis: doing so would create a
+// statically non-positive basis (rejected by the verifier) and crash the
+// constant folding. The op must stay unfolded.
+// CHECK-LABEL: @delinearize_dont_fold_non_positive_basis
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[CNEG:.+]] = arith.constant -3 : index
+// CHECK: affine.delinearize_index %{{.*}} into (4, %[[C0]], %[[CNEG]])
+func.func @delinearize_dont_fold_non_positive_basis(%idx: index) -> (index, index, index) {
+  %c0 = arith.constant 0 : index
+  %cneg = arith.constant -3 : index
+  %0:3 = affine.delinearize_index %idx into (4, %c0, %cneg) : index, index, index
+  return %0#0, %0#1, %0#2 : index, index, index
+}
+
+// -----
+
 func.func @drop_unit_basis_in_delinearize(%arg0 : index, %arg1 : index, %arg2 : index) ->
     (index, index, index, index, index, index) {
   %c1 = arith.constant 1 : index
@@ -1656,6 +1673,20 @@ func.func @linearize_fold_constants_no_outer_bound() -> index {
 
 // -----
 
+// Folding a dynamic basis operand to a constant must not disturb a pre-existing
+// statically non-positive basis element (which the linearize_index verifier
+// permits, unlike delinearize_index). Promoting the static 0 to a kDynamic
+// marker without a matching dynamic operand previously crashed the fold.
+// CHECK-LABEL: @linearize_fold_dynamic_keeps_static_non_positive_basis
+// CHECK: affine.linearize_index {{.*}} by (4, 0, 8)
+func.func @linearize_fold_dynamic_keeps_static_non_positive_basis(%i0: index, %i1: index, %i2: index) -> index {
+  %c8 = arith.constant 8 : index
+  %ret = affine.linearize_index [%i0, %i1, %i2] by (4, 0, %c8) : index
+  return %ret : index
+}
+
+// -----
+
 // CHECK-LABEL: @linearize_fold_empty_basis
 // CHECK-SAME: (%[[ARG0:.+]]: index)
 // CHECK-NOT: affine.linearize

``````````

</details>


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