diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp index e795f3f0b019b..9f569ea05967c 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..d08d2bf79c781 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) { // CHECK: %[[INDEX_CAST_0:.*]] = arith.index_cast %[[APPLY_3]] : index to i64 // CHECK: memref.store %[[INDEX_CAST_0]], %[[ALLOC]]{{\[}}%[[REMUI_0]]] : memref // 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 +}