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

[llvmlistbot at llvm.org]

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

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

>From 525515d028be16e9e390f5476c37189eaebb094f Mon Sep 17 00:00:00 2001
From: Chennes Xu <xuchen359 at gmail.com>
Date: Tue, 23 Jun 2026 16:41:39 +0900
Subject: [PATCH] [mlir][affine] Fix delinearize fold crash on non-positive
 basis

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.

Fixes #205238
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp   | 21 ++++++++-------
 mlir/test/Dialect/Affine/canonicalize.mlir | 31 ++++++++++++++++++++++
 2 files changed, 43 insertions(+), 9 deletions(-)

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