[llvm] b060661 - [SCEVExpander] Expand UDiv avoiding UB when in seq_min/max. (#92177)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Oct 17 13:55:25 PDT 2024
Author: Florian Hahn
Date: 2024-10-17T13:55:20-07:00
New Revision: b060661da8b3b53db55644e5e358bb2dca8b56d7
URL: https://github.com/llvm/llvm-project/commit/b060661da8b3b53db55644e5e358bb2dca8b56d7
DIFF: https://github.com/llvm/llvm-project/commit/b060661da8b3b53db55644e5e358bb2dca8b56d7.diff
LOG: [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
Added:
Modified:
llvm/include/llvm/Transforms/Utils/ScalarEvolutionExpander.h
llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp
llvm/test/Transforms/LoopVectorize/trip-count-expansion-may-introduce-ub.ll
Removed:
################################################################################
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<SCEVExpander, Value *> {
/// "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<InstSimplifyFolder, IRBuilderCallbackInserter> 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_
diff erent_bounds(ptr %dst, i64 %N, i64 %M) {
; CHECK-LABEL: define i64 @multi_exit_4_exit_count_with_udiv_by_value_in_latch_
diff erent_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_
diff erent_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_
diff erent_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]])
More information about the llvm-commits
mailing list