diff --git a/llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h b/llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h index 5697d983c9ad48..7dd754a2bc0def 100644 --- a/llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h +++ b/llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h @@ -126,6 +126,11 @@ class SCEVExpander : public SCEVVisitor { /// "expanded" form. bool LSRMode; + /// When true, rewrite any divisors of UDiv expressions that may be 0 to + /// umax(Divisor, 1) to avoid introducing UB. If the divisor may be poison, + /// freeze it first. + bool SafeUDivMode = false; + typedef IRBuilder BuilderType; BuilderType Builder; diff --git a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp index c412d0398b9558..39da38e4918176 100644 --- a/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp +++ b/llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp @@ -681,7 +681,21 @@ Value *SCEVExpander::visitUDivExpr(const SCEVUDivExpr *S) { SCEV::FlagAnyWrap, /*IsSafeToHoist*/ true); } - Value *RHS = expand(S->getRHS()); + const SCEV *RHSExpr = S->getRHS(); + Value *RHS = expand(RHSExpr); + if (SafeUDivMode) { + bool GuaranteedNotPoison = + ScalarEvolution::isGuaranteedNotToBePoison(RHSExpr); + if (!GuaranteedNotPoison) + RHS = Builder.CreateFreeze(RHS); + + // We need an umax if either RHSExpr is not known to be zero, or if it is + // not guaranteed to be non-poison. In the later case, the frozen poison may + // be 0. + if (!SE.isKnownNonZero(RHSExpr) || !GuaranteedNotPoison) + RHS = Builder.CreateIntrinsic(RHS->getType(), Intrinsic::umax, + {RHS, ConstantInt::get(RHS->getType(), 1)}); + } return InsertBinop(Instruction::UDiv, LHS, RHS, SCEV::FlagAnyWrap, /*IsSafeToHoist*/ SE.isKnownNonZero(S->getRHS())); } @@ -1376,11 +1390,14 @@ Value *SCEVExpander::visitSignExtendExpr(const SCEVSignExtendExpr *S) { Value *SCEVExpander::expandMinMaxExpr(const SCEVNAryExpr *S, Intrinsic::ID IntrinID, Twine Name, bool IsSequential) { + bool PrevSafeMode = SafeUDivMode; + SafeUDivMode |= IsSequential; Value *LHS = expand(S->getOperand(S->getNumOperands() - 1)); Type *Ty = LHS->getType(); if (IsSequential) LHS = Builder.CreateFreeze(LHS); for (int i = S->getNumOperands() - 2; i >= 0; --i) { + SafeUDivMode = (IsSequential && i != 0) || PrevSafeMode; Value *RHS = expand(S->getOperand(i)); if (IsSequential && i != 0) RHS = Builder.CreateFreeze(RHS); @@ -1395,6 +1412,7 @@ Value *SCEVExpander::expandMinMaxExpr(const SCEVNAryExpr *S, } LHS = Sel; } + SafeUDivMode = PrevSafeMode; return LHS; } diff --git a/llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll b/llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll index 7dfd80a688f312..1b646200e1c7c5 100644 --- a/llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll +++ b/llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll @@ -459,12 +459,12 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch(ptr %dst, i64 %N) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { ; CHECK-NEXT: entry: -; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[N]] +; CHECK-NEXT: [[TMP10:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1) +; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP10]] ; CHECK-NEXT: [[TMP8:%.*]] = freeze i64 [[TMP0]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP8]], i64 [[SMAX]]) @@ -529,13 +529,14 @@ exit: declare void @foo() -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_call_before_loop(ptr %dst, i64 %N) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_call_before_loop( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { ; CHECK-NEXT: entry: ; CHECK-NEXT: call void @foo() -; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[N]] +; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[N]] +; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1) +; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]] ; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]]) @@ -599,14 +600,15 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_loop_may_not_execute(ptr %dst, i64 %N, i1 %c) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_loop_may_not_execute( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i1 [[C:%.*]]) { ; CHECK-NEXT: entry: ; CHECK-NEXT: br i1 [[C]], label [[LOOP_HEADER_PREHEADER:%.*]], label [[EXIT:%.*]] ; CHECK: loop.header.preheader: -; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[N]] +; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[N]] +; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1) +; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]] ; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]]) @@ -672,12 +674,13 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_different_bounds(ptr %dst, i64 %N, i64 %M) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_different_bounds( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i64 [[M:%.*]]) { ; CHECK-NEXT: entry: -; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[M]] +; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[M]] +; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1) +; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]] ; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]]) @@ -740,13 +743,13 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_udiv_by_frozen_value_in_latch(ptr %dst, i64 %N) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_frozen_value_in_latch( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { ; CHECK-NEXT: entry: ; CHECK-NEXT: [[FR_N:%.*]] = freeze i64 [[N]] -; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[FR_N]] +; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[FR_N]], i64 1) +; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]] ; CHECK-NEXT: [[TMP10:%.*]] = freeze i64 [[TMP2]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP10]], i64 [[SMAX]]) @@ -931,12 +934,12 @@ exit: ret void } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_urem_by_value_in_latch(ptr %dst, i64 %N) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_urem_by_value_in_latch( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { ; CHECK-NEXT: entry: -; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[N]] +; CHECK-NEXT: [[TMP12:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1) +; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP12]] ; CHECK-NEXT: [[TMP1:%.*]] = mul nuw i64 [[N]], [[TMP0]] ; CHECK-NEXT: [[TMP2:%.*]] = sub i64 42, [[TMP1]] ; CHECK-NEXT: [[SMAX1:%.*]] = call i64 @llvm.smax.i64(i64 [[TMP2]], i64 0) @@ -1002,7 +1005,6 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_4_exit_count_with_urem_by_constant_in_latch(ptr %dst, i64 %N) { ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_urem_by_constant_in_latch( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { @@ -1156,7 +1158,8 @@ define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch1(ptr %dst, i64 % ; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch1( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) { ; CHECK-NEXT: entry: -; CHECK-NEXT: [[TMP9:%.*]] = udiv i64 42, [[N]] +; CHECK-NEXT: [[TMP8:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1) +; CHECK-NEXT: [[TMP9:%.*]] = udiv i64 42, [[TMP8]] ; CHECK-NEXT: [[TMP10:%.*]] = freeze i64 [[TMP9]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP10]], i64 [[SMAX]]) @@ -1262,13 +1265,14 @@ exit: ret i64 %p } -; FIXME: currently the expansion of the loop bounds may introduce UB through the division. define i64 @multi_exit_count_with_udiv_by_value_in_latch_different_bounds_divisor_non_zero_may_be_poison(ptr %dst, i64 %N, i64 %M) { ; CHECK-LABEL: define i64 @multi_exit_count_with_udiv_by_value_in_latch_different_bounds_divisor_non_zero_may_be_poison( ; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i64 [[M:%.*]]) { ; CHECK-NEXT: entry: ; CHECK-NEXT: [[M_1:%.*]] = call i64 @llvm.umax.i64(i64 [[M]], i64 1) -; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[M_1]] +; CHECK-NEXT: [[TMP9:%.*]] = freeze i64 [[M_1]] +; CHECK-NEXT: [[TMP10:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP9]], i64 1) +; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP10]] ; CHECK-NEXT: [[TMP1:%.*]] = freeze i64 [[TMP0]] ; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0) ; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP1]], i64 [[SMAX]])