Skip to content

Commit ead2724

Browse files
authored
[mlir][scf] Fix a div-by-zero bug when step of scf.for is zero (#131079)
Fixes #130095.
1 parent 08dda4d commit ead2724

File tree

2 files changed

+20
-0
lines changed

2 files changed

+20
-0
lines changed

mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,10 @@ bool LoopPipelinerInternal::initializeLoopInfo(
119119
int64_t ubImm = upperBoundCst.value();
120120
int64_t lbImm = lowerBoundCst.value();
121121
int64_t stepImm = stepCst.value();
122+
if (stepImm <= 0) {
123+
LDBG("--invalid loop step -> BAIL");
124+
return false;
125+
}
122126
int64_t numIteration = llvm::divideCeilSigned(ubImm - lbImm, stepImm);
123127
if (numIteration > maxStage) {
124128
dynamicLoop = false;

mlir/test/Dialect/SCF/loop-pipelining.mlir

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -985,3 +985,19 @@ func.func @invalid_schedule3(%A: memref<?xf32>, %result: memref<?xf32>, %ext: in
985985
} { __test_pipelining_loop__ }
986986
return
987987
}
988+
989+
// -----
990+
991+
// Ensure this case not crash when step is zero.
992+
993+
// CHECK-LABEL: @invalid_loop_step
994+
func.func @invalid_loop_step(%A: memref<?xf32>, %result: memref<?xf32>) {
995+
%c0 = arith.constant 0 : index
996+
%cf = arith.constant 1.0 : f32
997+
scf.for %i0 = %c0 to %c0 step %c0 {
998+
%A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
999+
%A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : f32
1000+
memref.store %A1_elem, %result[%i0] { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : memref<?xf32>
1001+
} { __test_pipelining_loop__ }
1002+
return
1003+
}

0 commit comments

Comments
 (0)