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

[Mehdi Amini]

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

>From 5af2d9c00709317748aef50264a7a97bda9b713b 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

Assisted-by: Claude Code
---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp |  7 +++++-
 mlir/test/Dialect/Affine/invalid.mlir    | 27 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

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