Skip to content

Commit ab73bd3

Browse files
committed
[InstCombine] Enhance select icmp and folding
This folds (a << k) ? 2^k * a : 0 to 2^k * a. https://alive2.llvm.org/ce/z/_dDRjo Fix #62155. Reviewed By: nikic Differential Revision: https://reviews.llvm.org/D148420
1 parent 226896c commit ab73bd3

File tree

2 files changed

+56
-35
lines changed

2 files changed

+56
-35
lines changed

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -609,6 +609,44 @@ static Instruction *foldSelectICmpAndAnd(Type *SelType, const ICmpInst *Cmp,
609609
return new ZExtInst(ICmpNeZero, SelType);
610610
}
611611

612+
/// We want to turn:
613+
/// (select (icmp eq (and X, C1), 0), 0, (shl [nsw/nuw] X, C2));
614+
/// iff C1 is a mask and the number of its leading zeros is equal to C2
615+
/// into:
616+
/// shl X, C2
617+
static Value *foldSelectICmpAndZeroShl(const ICmpInst *Cmp, Value *TVal,
618+
Value *FVal,
619+
InstCombiner::BuilderTy &Builder) {
620+
ICmpInst::Predicate Pred;
621+
Value *AndVal;
622+
if (!match(Cmp, m_ICmp(Pred, m_Value(AndVal), m_Zero())))
623+
return nullptr;
624+
625+
if (Pred == ICmpInst::ICMP_NE) {
626+
Pred = ICmpInst::ICMP_EQ;
627+
std::swap(TVal, FVal);
628+
}
629+
630+
Value *X;
631+
const APInt *C2, *C1;
632+
if (Pred != ICmpInst::ICMP_EQ ||
633+
!match(AndVal, m_And(m_Value(X), m_APInt(C1))) ||
634+
!match(TVal, m_Zero()) || !match(FVal, m_Shl(m_Specific(X), m_APInt(C2))))
635+
return nullptr;
636+
637+
if (!C1->isMask() ||
638+
C1->countLeadingZeros() != static_cast<unsigned>(C2->getZExtValue()))
639+
return nullptr;
640+
641+
auto *FI = dyn_cast<Instruction>(FVal);
642+
if (!FI)
643+
return nullptr;
644+
645+
FI->setHasNoSignedWrap(false);
646+
FI->setHasNoUnsignedWrap(false);
647+
return FVal;
648+
}
649+
612650
/// We want to turn:
613651
/// (select (icmp sgt x, C), lshr (X, Y), ashr (X, Y)); iff C s>= -1
614652
/// (select (icmp slt x, C), ashr (X, Y), lshr (X, Y)); iff C s>= 0
@@ -1743,6 +1781,9 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
17431781
foldSelectICmpAndAnd(SI.getType(), ICI, TrueVal, FalseVal, Builder))
17441782
return V;
17451783

1784+
if (Value *V = foldSelectICmpAndZeroShl(ICI, TrueVal, FalseVal, Builder))
1785+
return replaceInstUsesWith(SI, V);
1786+
17461787
if (Instruction *V = foldSelectCtlzToCttz(ICI, TrueVal, FalseVal, Builder))
17471788
return V;
17481789

llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll

Lines changed: 15 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,8 @@
55

66
define i32 @test_eq(i32 %x) {
77
; CHECK-LABEL: @test_eq(
8-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
9-
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
10-
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
11-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
12-
; CHECK-NEXT: ret i32 [[COND]]
8+
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
9+
; CHECK-NEXT: ret i32 [[MUL]]
1310
;
1411
%shl.mask = and i32 %x, 1073741823
1512
%tobool = icmp eq i32 %shl.mask, 0
@@ -20,11 +17,8 @@ define i32 @test_eq(i32 %x) {
2017

2118
define <2 x i32> @test_eq_vect(<2 x i32> %x) {
2219
; CHECK-LABEL: @test_eq_vect(
23-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and <2 x i32> [[X:%.*]], <i32 1073741823, i32 1073741823>
24-
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer
25-
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], <i32 2, i32 2>
26-
; CHECK-NEXT: [[COND:%.*]] = select <2 x i1> [[TOBOOL]], <2 x i32> zeroinitializer, <2 x i32> [[MUL]]
27-
; CHECK-NEXT: ret <2 x i32> [[COND]]
20+
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X:%.*]], <i32 2, i32 2>
21+
; CHECK-NEXT: ret <2 x i32> [[MUL]]
2822
;
2923
%shl.mask = and <2 x i32> %x, <i32 1073741823, i32 1073741823>
3024
%tobool = icmp eq <2 x i32> %shl.mask, zeroinitializer
@@ -35,11 +29,8 @@ define <2 x i32> @test_eq_vect(<2 x i32> %x) {
3529

3630
define i32 @test_ne(i32 %x) {
3731
; CHECK-LABEL: @test_ne(
38-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
39-
; CHECK-NEXT: [[TOBOOL_NOT_NOT:%.*]] = icmp eq i32 [[SHL_MASK]], 0
40-
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
41-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL_NOT_NOT]], i32 0, i32 [[MUL]]
42-
; CHECK-NEXT: ret i32 [[COND]]
32+
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
33+
; CHECK-NEXT: ret i32 [[MUL]]
4334
;
4435
%shl.mask = and i32 %x, 1073741823
4536
%tobool.not = icmp ne i32 %shl.mask, 0
@@ -50,11 +41,8 @@ define i32 @test_ne(i32 %x) {
5041

5142
define <2 x i32> @test_ne_vect(<2 x i32> %x) {
5243
; CHECK-LABEL: @test_ne_vect(
53-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and <2 x i32> [[X:%.*]], <i32 1073741823, i32 1073741823>
54-
; CHECK-NEXT: [[TOBOOL_NOT_NOT:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer
55-
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], <i32 2, i32 2>
56-
; CHECK-NEXT: [[COND:%.*]] = select <2 x i1> [[TOBOOL_NOT_NOT]], <2 x i32> zeroinitializer, <2 x i32> [[MUL]]
57-
; CHECK-NEXT: ret <2 x i32> [[COND]]
44+
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X:%.*]], <i32 2, i32 2>
45+
; CHECK-NEXT: ret <2 x i32> [[MUL]]
5846
;
5947
%shl.mask = and <2 x i32> %x, <i32 1073741823, i32 1073741823>
6048
%tobool.not = icmp ne <2 x i32> %shl.mask, zeroinitializer
@@ -65,11 +53,8 @@ define <2 x i32> @test_ne_vect(<2 x i32> %x) {
6553

6654
define i32 @test_nuw_dropped(i32 %x) {
6755
; CHECK-LABEL: @test_nuw_dropped(
68-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
69-
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
70-
; CHECK-NEXT: [[MUL:%.*]] = shl nuw i32 [[X]], 2
71-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
72-
; CHECK-NEXT: ret i32 [[COND]]
56+
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
57+
; CHECK-NEXT: ret i32 [[MUL]]
7358
;
7459
%shl.mask = and i32 %x, 1073741823
7560
%tobool = icmp eq i32 %shl.mask, 0
@@ -80,11 +65,8 @@ define i32 @test_nuw_dropped(i32 %x) {
8065

8166
define i32 @test_nsw_dropped(i32 %x) {
8267
; CHECK-LABEL: @test_nsw_dropped(
83-
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
84-
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
85-
; CHECK-NEXT: [[MUL:%.*]] = shl nsw i32 [[X]], 2
86-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
87-
; CHECK-NEXT: ret i32 [[COND]]
68+
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
69+
; CHECK-NEXT: ret i32 [[MUL]]
8870
;
8971
%shl.mask = and i32 %x, 1073741823
9072
%tobool = icmp eq i32 %shl.mask, 0
@@ -101,9 +83,8 @@ define i32 @test_multi_use(i32 %x) {
10183
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
10284
; CHECK-NEXT: [[TOBOOL_NOT:%.*]] = icmp ne i32 [[SHL_MASK]], 0
10385
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
104-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL_NOT]], i32 [[MUL]], i32 0
10586
; CHECK-NEXT: call void @use_multi(i32 [[SHL_MASK]], i1 [[TOBOOL_NOT]], i32 [[MUL]])
106-
; CHECK-NEXT: ret i32 [[COND]]
87+
; CHECK-NEXT: ret i32 [[MUL]]
10788
;
10889
%shl.mask = and i32 %x, 1073741823
10990
%tobool.not = icmp ne i32 %shl.mask, 0
@@ -117,10 +98,9 @@ define i32 @test_multi_use_nuw_dropped(i32 %x) {
11798
; CHECK-LABEL: @test_multi_use_nuw_dropped(
11899
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
119100
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
120-
; CHECK-NEXT: [[MUL:%.*]] = shl nuw i32 [[X]], 2
121-
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
101+
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
122102
; CHECK-NEXT: call void @use_multi(i32 [[SHL_MASK]], i1 [[TOBOOL]], i32 [[MUL]])
123-
; CHECK-NEXT: ret i32 [[COND]]
103+
; CHECK-NEXT: ret i32 [[MUL]]
124104
;
125105
%shl.mask = and i32 %x, 1073741823
126106
%tobool = icmp eq i32 %shl.mask, 0

0 commit comments

Comments
 (0)