Skip to content

Commit b060661

Browse files
fhahnfhahn
authored
[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 #89958. PR #92177
1 parent e913a33 commit b060661

File tree

3 files changed

+44
-17
lines changed

3 files changed

+44
-17
lines changed

‎llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,11 @@ class SCEVExpander : public SCEVVisitor<SCEVExpander, Value *> {
126126
/// "expanded" form.
127127
bool LSRMode;
128128

129+
/// When true, rewrite any divisors of UDiv expressions that may be 0 to
130+
/// umax(Divisor, 1) to avoid introducing UB. If the divisor may be poison,
131+
/// freeze it first.
132+
bool SafeUDivMode = false;
133+
129134
typedef IRBuilder<InstSimplifyFolder, IRBuilderCallbackInserter> BuilderType;
130135
BuilderType Builder;
131136

‎llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -681,7 +681,21 @@ Value *SCEVExpander::visitUDivExpr(const SCEVUDivExpr *S) {
681681
SCEV::FlagAnyWrap, /*IsSafeToHoist*/ true);
682682
}
683683

684-
Value *RHS = expand(S->getRHS());
684+
const SCEV *RHSExpr = S->getRHS();
685+
Value *RHS = expand(RHSExpr);
686+
if (SafeUDivMode) {
687+
bool GuaranteedNotPoison =
688+
ScalarEvolution::isGuaranteedNotToBePoison(RHSExpr);
689+
if (!GuaranteedNotPoison)
690+
RHS = Builder.CreateFreeze(RHS);
691+
692+
// We need an umax if either RHSExpr is not known to be zero, or if it is
693+
// not guaranteed to be non-poison. In the later case, the frozen poison may
694+
// be 0.
695+
if (!SE.isKnownNonZero(RHSExpr) || !GuaranteedNotPoison)
696+
RHS = Builder.CreateIntrinsic(RHS->getType(), Intrinsic::umax,
697+
{RHS, ConstantInt::get(RHS->getType(), 1)});
698+
}
685699
return InsertBinop(Instruction::UDiv, LHS, RHS, SCEV::FlagAnyWrap,
686700
/*IsSafeToHoist*/ SE.isKnownNonZero(S->getRHS()));
687701
}
@@ -1376,11 +1390,14 @@ Value *SCEVExpander::visitSignExtendExpr(const SCEVSignExtendExpr *S) {
13761390
Value *SCEVExpander::expandMinMaxExpr(const SCEVNAryExpr *S,
13771391
Intrinsic::ID IntrinID, Twine Name,
13781392
bool IsSequential) {
1393+
bool PrevSafeMode = SafeUDivMode;
1394+
SafeUDivMode |= IsSequential;
13791395
Value *LHS = expand(S->getOperand(S->getNumOperands() - 1));
13801396
Type *Ty = LHS->getType();
13811397
if (IsSequential)
13821398
LHS = Builder.CreateFreeze(LHS);
13831399
for (int i = S->getNumOperands() - 2; i >= 0; --i) {
1400+
SafeUDivMode = (IsSequential && i != 0) || PrevSafeMode;
13841401
Value *RHS = expand(S->getOperand(i));
13851402
if (IsSequential && i != 0)
13861403
RHS = Builder.CreateFreeze(RHS);
@@ -1395,6 +1412,7 @@ Value *SCEVExpander::expandMinMaxExpr(const SCEVNAryExpr *S,
13951412
}
13961413
LHS = Sel;
13971414
}
1415+
SafeUDivMode = PrevSafeMode;
13981416
return LHS;
13991417
}
14001418

‎llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll‎

Lines changed: 20 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -459,12 +459,12 @@ exit:
459459
ret i64 %p
460460
}
461461

462-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
463462
define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch(ptr %dst, i64 %N) {
464463
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch(
465464
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) {
466465
; CHECK-NEXT: entry:
467-
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[N]]
466+
; CHECK-NEXT: [[TMP10:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1)
467+
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP10]]
468468
; CHECK-NEXT: [[TMP8:%.*]] = freeze i64 [[TMP0]]
469469
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
470470
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP8]], i64 [[SMAX]])
@@ -529,13 +529,14 @@ exit:
529529

530530
declare void @foo()
531531

532-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
533532
define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_call_before_loop(ptr %dst, i64 %N) {
534533
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_call_before_loop(
535534
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) {
536535
; CHECK-NEXT: entry:
537536
; CHECK-NEXT: call void @foo()
538-
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[N]]
537+
; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[N]]
538+
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1)
539+
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]]
539540
; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]]
540541
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
541542
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]])
@@ -599,14 +600,15 @@ exit:
599600
ret i64 %p
600601
}
601602

602-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
603603
define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_loop_may_not_execute(ptr %dst, i64 %N, i1 %c) {
604604
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_loop_may_not_execute(
605605
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i1 [[C:%.*]]) {
606606
; CHECK-NEXT: entry:
607607
; CHECK-NEXT: br i1 [[C]], label [[LOOP_HEADER_PREHEADER:%.*]], label [[EXIT:%.*]]
608608
; CHECK: loop.header.preheader:
609-
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[N]]
609+
; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[N]]
610+
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1)
611+
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]]
610612
; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]]
611613
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
612614
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]])
@@ -672,12 +674,13 @@ exit:
672674
ret i64 %p
673675
}
674676

675-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
676677
define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_different_bounds(ptr %dst, i64 %N, i64 %M) {
677678
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_different_bounds(
678679
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i64 [[M:%.*]]) {
679680
; CHECK-NEXT: entry:
680-
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[M]]
681+
; CHECK-NEXT: [[TMP0:%.*]] = freeze i64 [[M]]
682+
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP0]], i64 1)
683+
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]]
681684
; CHECK-NEXT: [[TMP3:%.*]] = freeze i64 [[TMP2]]
682685
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
683686
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP3]], i64 [[SMAX]])
@@ -740,13 +743,13 @@ exit:
740743
ret i64 %p
741744
}
742745

743-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
744746
define i64 @multi_exit_4_exit_count_with_udiv_by_frozen_value_in_latch(ptr %dst, i64 %N) {
745747
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_frozen_value_in_latch(
746748
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) {
747749
; CHECK-NEXT: entry:
748750
; CHECK-NEXT: [[FR_N:%.*]] = freeze i64 [[N]]
749-
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[FR_N]]
751+
; CHECK-NEXT: [[TMP1:%.*]] = call i64 @llvm.umax.i64(i64 [[FR_N]], i64 1)
752+
; CHECK-NEXT: [[TMP2:%.*]] = udiv i64 42, [[TMP1]]
750753
; CHECK-NEXT: [[TMP10:%.*]] = freeze i64 [[TMP2]]
751754
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
752755
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP10]], i64 [[SMAX]])
@@ -931,12 +934,12 @@ exit:
931934
ret void
932935
}
933936

934-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
935937
define i64 @multi_exit_4_exit_count_with_urem_by_value_in_latch(ptr %dst, i64 %N) {
936938
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_urem_by_value_in_latch(
937939
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) {
938940
; CHECK-NEXT: entry:
939-
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[N]]
941+
; CHECK-NEXT: [[TMP12:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1)
942+
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP12]]
940943
; CHECK-NEXT: [[TMP1:%.*]] = mul nuw i64 [[N]], [[TMP0]]
941944
; CHECK-NEXT: [[TMP2:%.*]] = sub i64 42, [[TMP1]]
942945
; CHECK-NEXT: [[SMAX1:%.*]] = call i64 @llvm.smax.i64(i64 [[TMP2]], i64 0)
@@ -1002,7 +1005,6 @@ exit:
10021005
ret i64 %p
10031006
}
10041007

1005-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
10061008
define i64 @multi_exit_4_exit_count_with_urem_by_constant_in_latch(ptr %dst, i64 %N) {
10071009
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_urem_by_constant_in_latch(
10081010
; 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 %
11561158
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch1(
11571159
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]]) {
11581160
; CHECK-NEXT: entry:
1159-
; CHECK-NEXT: [[TMP9:%.*]] = udiv i64 42, [[N]]
1161+
; CHECK-NEXT: [[TMP8:%.*]] = call i64 @llvm.umax.i64(i64 [[N]], i64 1)
1162+
; CHECK-NEXT: [[TMP9:%.*]] = udiv i64 42, [[TMP8]]
11601163
; CHECK-NEXT: [[TMP10:%.*]] = freeze i64 [[TMP9]]
11611164
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
11621165
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP10]], i64 [[SMAX]])
@@ -1262,13 +1265,14 @@ exit:
12621265
ret i64 %p
12631266
}
12641267

1265-
; FIXME: currently the expansion of the loop bounds may introduce UB through the division.
12661268
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) {
12671269
; CHECK-LABEL: define i64 @multi_exit_count_with_udiv_by_value_in_latch_different_bounds_divisor_non_zero_may_be_poison(
12681270
; CHECK-SAME: ptr [[DST:%.*]], i64 [[N:%.*]], i64 [[M:%.*]]) {
12691271
; CHECK-NEXT: entry:
12701272
; CHECK-NEXT: [[M_1:%.*]] = call i64 @llvm.umax.i64(i64 [[M]], i64 1)
1271-
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[M_1]]
1273+
; CHECK-NEXT: [[TMP9:%.*]] = freeze i64 [[M_1]]
1274+
; CHECK-NEXT: [[TMP10:%.*]] = call i64 @llvm.umax.i64(i64 [[TMP9]], i64 1)
1275+
; CHECK-NEXT: [[TMP0:%.*]] = udiv i64 42, [[TMP10]]
12721276
; CHECK-NEXT: [[TMP1:%.*]] = freeze i64 [[TMP0]]
12731277
; CHECK-NEXT: [[SMAX:%.*]] = call i64 @llvm.smax.i64(i64 [[N]], i64 0)
12741278
; CHECK-NEXT: [[UMIN:%.*]] = call i64 @llvm.umin.i64(i64 [[TMP1]], i64 [[SMAX]])

0 commit comments

Comments
 (0)