From b060661da8b3b53db55644e5e358bb2dca8b56d7 Mon Sep 17 00:00:00 2001 From: Florian Hahn Date: Thu, 17 Oct 2024 13:55:20 -0700 Subject: [PATCH] [SCEVExpander] Expand UDiv avoiding UB when in seq_min/max. (#92177) Update SCEVExpander to introduce an SafeUDivMode, which is set when expanding operands of SCEVSequentialMinMaxExpr. In this mode, the expander will make sure that the divisor of the expanded UDiv is neither 0 nor poison. Fixes https://github.com/llvm/llvm-project/issues/89958. PR https://github.com/llvm/llvm-project/pull/92177 --- .../Utils/ScalarEvolutionExpander.h | 5 +++ .../Utils/ScalarEvolutionExpander.cpp | 20 ++++++++++- .../trip-count-expansion-may-introduce-ub.ll | 36 ++++++++++--------- 3 files changed, 44 insertions(+), 17 deletions(-) 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]])