[mlir][scf] Fix a div-by-zero bug when step of scf.for is zero

[CoTinker]

Fixes #130095.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: Longsheng Mou (CoTinker)

Changes

Fixes #130095.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp (+4)
• (modified) mlir/test/Dialect/SCF/loop-pipelining.mlir (+16)

diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index 1b458f410af60..e740564f422b2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -119,6 +119,10 @@ bool LoopPipelinerInternal::initializeLoopInfo(
     int64_t ubImm = upperBoundCst.value();
     int64_t lbImm = lowerBoundCst.value();
     int64_t stepImm = stepCst.value();
+    if (stepImm <= 0) {
+      LDBG("--invalid loop step -> BAIL");
+      return false;
+    }
     int64_t numIteration = llvm::divideCeilSigned(ubImm - lbImm, stepImm);
     if (numIteration > maxStage) {
       dynamicLoop = false;
diff --git a/mlir/test/Dialect/SCF/loop-pipelining.mlir b/mlir/test/Dialect/SCF/loop-pipelining.mlir
index c879c83275bf8..25e5703cea5c6 100644
--- a/mlir/test/Dialect/SCF/loop-pipelining.mlir
+++ b/mlir/test/Dialect/SCF/loop-pipelining.mlir
@@ -985,3 +985,19 @@ func.func @invalid_schedule3(%A: memref<?xf32>, %result: memref<?xf32>, %ext: in
   }  { __test_pipelining_loop__ }
   return
 }
+
+// -----
+
+// Ensure this case not crash when step is zero.
+
+// CHECK-LABEL: @invalid_loop_step
+func.func @invalid_loop_step(%A: memref<?xf32>, %result: memref<?xf32>) {
+  %c0 = arith.constant 0 : index
+  %cf = arith.constant 1.0 : f32
+  scf.for %i0 = %c0 to %c0 step %c0 {
+    %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
+    %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
+    memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : memref<?xf32>
+  }  { __test_pipelining_loop__ }
+  return
+}

[CoTinker]

Ping~