[Mlir-commits] [mlir] [mlir] Fix a crash in SCF loop coalescing when loop step is zero (PR #181893)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Darshan Bhat (darshan-opensource)

<details>
<summary>Changes</summary>

The affine-loop-coalescing pass crashes when loop has a statistically known zero step. This fix does an early bail-out in coalesceLoops that return failure when a loop has a known zero step.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SCF/Utils/Utils.cpp (+9) 
- (modified) mlir/test/Dialect/Affine/loop-coalescing.mlir (+18) 


``````````diff
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index f8a4f057c9f0d..e0a2f01f37307 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -914,6 +914,15 @@ LogicalResult mlir::coalesceLoops(RewriterBase &rewriter,
   scf::ForOp innermost = loops.back();
   scf::ForOp outermost = loops.front();
 
+  // Bail out if any loop has a known zero step, as normalization
+  // would result in a division by zero.
+  for (auto loop : loops) {
+    if (auto step = getConstantIntValue(loop.getStep())) {
+      if (step.value() == 0) {
+        return failure();
+      }
+    }
+  }
   // 1. Make sure all loops iterate from 0 to upperBound with step 1.  This
   // allows the following code to assume upperBound is the number of iterations.
   for (auto loop : loops) {
diff --git a/mlir/test/Dialect/Affine/loop-coalescing.mlir b/mlir/test/Dialect/Affine/loop-coalescing.mlir
index 6a825320ff20f..2421e0eaff27a 100644
--- a/mlir/test/Dialect/Affine/loop-coalescing.mlir
+++ b/mlir/test/Dialect/Affine/loop-coalescing.mlir
@@ -444,3 +444,21 @@ func.func @inner_loop_has_iter_args(%alloc : memref<?xi64>) {
 // CHECK:   %[[INDEX_CAST_0:.*]] = arith.index_cast %[[APPLY_3]] : index to i64
 // CHECK:   memref.store %[[INDEX_CAST_0]], %[[ALLOC]]{{\[}}%[[REMUI_0]]] : memref<?xi64>
 // CHECK: }
+
+// -----
+
+// Verify that coalescing is not attempted when a loop has a zero step,
+// which would cause a division by zero during normalization.
+
+// CHECK-LABEL: @no_coalesce_zero_step
+func.func @no_coalesce_zero_step(%lb: index, %ub: index) {
+  %c0 = arith.constant 0 : index
+  // CHECK: scf.for
+  // CHECK: scf.for
+  scf.for %i = %lb to %ub step %c0 {
+    scf.for %j = %lb to %ub step %c0 {
+      "use"(%i,%j) : (index, index) -> ()
+    }
+  }
+  return
+}
\ No newline at end of file

``````````

</details>


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