[Mlir-commits] [mlir] [mlir][affine] Reject affine.for with step <= 0 in parser and verifier (PR #184158)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/184158

An affine.for loop with step 0 caused a crash in affine-loop-unroll due to a division-by-zero assertion in getTripCountMapAndOperands:

  llvm::divideCeilSigned(loopSpan, step) -- step == 0

The parser already rejected negative steps but allowed zero through.

Fix this by:
- Updating the parser to reject any step <= 0 (not only negative).
- Adding a verifier check in AffineForOp::verifyRegions() so that loops constructed programmatically with an invalid step are caught before any analysis pass runs.
- Adding a defensive early-return guard in getTripCountMapAndOperands for step <= 0, consistent with getTrivialConstantTripCount which already performs this guard.

Fixes #107812

>From 7b5b19ffdd83551def22df10c182511fe1e2ed59 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 2 Mar 2026 07:39:09 -0800
Subject: [PATCH] [mlir][affine] Reject affine.for with step <= 0 in parser and
 verifier

An affine.for loop with step 0 caused a crash in affine-loop-unroll
due to a division-by-zero assertion in getTripCountMapAndOperands:

  llvm::divideCeilSigned(loopSpan, step) -- step == 0

The parser already rejected negative steps but allowed zero through.

Fix this by:
- Updating the parser to reject any step <= 0 (not only negative).
- Adding a verifier check in AffineForOp::verifyRegions() so that
  loops constructed programmatically with an invalid step are caught
  before any analysis pass runs.
- Adding a defensive early-return guard in getTripCountMapAndOperands
  for step <= 0, consistent with getTrivialConstantTripCount which
  already performs this guard.

Fixes #107812
---
 .../Dialect/Affine/Analysis/LoopAnalysis.cpp  |  4 +++
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp      |  7 ++++-
 mlir/test/Dialect/Affine/invalid.mlir         | 27 +++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index 166d39e88d41e..1d2d5c01337c4 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -168,6 +168,10 @@ void mlir::affine::getTripCountMapAndOperands(
     SmallVectorImpl<Value> *tripCountOperands) {
   MLIRContext *context = forOp.getContext();
   int64_t step = forOp.getStepAsInt();
+  if (step <= 0) {
+    *tripCountMap = AffineMap();
+    return;
+  }
   int64_t loopSpan;
   if (forOp.hasConstantBounds()) {
     int64_t lb = forOp.getConstantLowerBound();
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 561fcc2ee5f6a..17107b52d2571 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2204,6 +2204,11 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result, int64_t lb,
 }
 
 LogicalResult AffineForOp::verifyRegions() {
+  // Step must be a strictly positive integer.
+  if (getStepAsInt() <= 0)
+    return emitOpError("expected step to be a positive integer, got ")
+           << getStepAsInt();
+
   // Check that the body defines as single block argument for the induction
   // variable.
   auto *body = getBody();
@@ -2370,7 +2375,7 @@ ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
                               result.attributes))
       return failure();
 
-    if (stepAttr.getValue().isNegative())
+    if (!stepAttr.getValue().isStrictlyPositive())
       return parser.emitError(
           stepLoc,
           "expected step to be representable as a positive signed integer");
diff --git a/mlir/test/Dialect/Affine/invalid.mlir b/mlir/test/Dialect/Affine/invalid.mlir
index c4bb22f9a8dae..e31f2f00f96e2 100644
--- a/mlir/test/Dialect/Affine/invalid.mlir
+++ b/mlir/test/Dialect/Affine/invalid.mlir
@@ -605,3 +605,30 @@ func.func @invalid_symbol() {
   }
   return
 }
+
+
+// -----
+
+// Regression test: affine.for with step 0 must be rejected by the parser.
+// https://github.com/llvm/llvm-project/issues/107812
+func.func @affine_for_zero_step_parser(%mem : memref<8xf32>) {
+  // expected-error at +1 {{expected step to be representable as a positive signed integer}}
+  affine.for %i = 0 to 8 step 0 {
+    affine.load %mem[%i] : memref<8xf32>
+  }
+  return
+}
+
+// -----
+
+// Regression test: affine.for with step 0 constructed via generic syntax must
+// be rejected by the verifier.
+// https://github.com/llvm/llvm-project/issues/107812
+func.func @affine_for_zero_step_verifier() {
+  // expected-error at +1 {{'affine.for' op expected step to be a positive integer, got 0}}
+  "affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 0 : index, upperBoundMap = affine_map<() -> (8)>}> ({
+  ^bb0(%i : index):
+    "affine.yield"() : () -> ()
+  }) : () -> ()
+  return
+}