Skip to content

Commit

Permalink
[InstCombine] Enhance select icmp and folding
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@PeixinQiao
PeixinQiao committed Jul 12, 2023
1 parent 226896c commit ab73bd3
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 35 deletions.
41 changes: 41 additions & 0 deletions llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
View file
Expand Up @@ -609,6 +609,44 @@ static Instruction *foldSelectICmpAndAnd(Type *SelType, const ICmpInst *Cmp,
return new ZExtInst(ICmpNeZero, SelType);
}

/// We want to turn:
/// (select (icmp eq (and X, C1), 0), 0, (shl [nsw/nuw] X, C2));
/// iff C1 is a mask and the number of its leading zeros is equal to C2
/// into:
/// shl X, C2
static Value *foldSelectICmpAndZeroShl(const ICmpInst *Cmp, Value *TVal,
Value *FVal,
InstCombiner::BuilderTy &Builder) {
ICmpInst::Predicate Pred;
Value *AndVal;
if (!match(Cmp, m_ICmp(Pred, m_Value(AndVal), m_Zero())))
return nullptr;

if (Pred == ICmpInst::ICMP_NE) {
Pred = ICmpInst::ICMP_EQ;
std::swap(TVal, FVal);
}

Value *X;
const APInt *C2, *C1;
if (Pred != ICmpInst::ICMP_EQ ||
!match(AndVal, m_And(m_Value(X), m_APInt(C1))) ||
!match(TVal, m_Zero()) || !match(FVal, m_Shl(m_Specific(X), m_APInt(C2))))
return nullptr;

if (!C1->isMask() ||
C1->countLeadingZeros() != static_cast<unsigned>(C2->getZExtValue()))
return nullptr;

auto *FI = dyn_cast<Instruction>(FVal);
if (!FI)
return nullptr;

FI->setHasNoSignedWrap(false);
FI->setHasNoUnsignedWrap(false);
return FVal;
}

/// We want to turn:
/// (select (icmp sgt x, C), lshr (X, Y), ashr (X, Y)); iff C s>= -1
/// (select (icmp slt x, C), ashr (X, Y), lshr (X, Y)); iff C s>= 0
Expand Down Expand Up @@ -1743,6 +1781,9 @@ Instruction *InstCombinerImpl::foldSelectInstWithICmp(SelectInst &SI,
foldSelectICmpAndAnd(SI.getType(), ICI, TrueVal, FalseVal, Builder))
return V;

if (Value *V = foldSelectICmpAndZeroShl(ICI, TrueVal, FalseVal, Builder))
return replaceInstUsesWith(SI, V);

if (Instruction *V = foldSelectCtlzToCttz(ICI, TrueVal, FalseVal, Builder))
return V;

Expand Down
50 changes: 15 additions & 35 deletions llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll
View file
Expand Up @@ -5,11 +5,8 @@

define i32 @test_eq(i32 %x) {
; CHECK-LABEL: @test_eq(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool = icmp eq i32 %shl.mask, 0
Expand All @@ -20,11 +17,8 @@ define i32 @test_eq(i32 %x) {

define <2 x i32> @test_eq_vect(<2 x i32> %x) {
; CHECK-LABEL: @test_eq_vect(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and <2 x i32> [[X:%.*]], <i32 1073741823, i32 1073741823>
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], <i32 2, i32 2>
; CHECK-NEXT: [[COND:%.*]] = select <2 x i1> [[TOBOOL]], <2 x i32> zeroinitializer, <2 x i32> [[MUL]]
; CHECK-NEXT: ret <2 x i32> [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X:%.*]], <i32 2, i32 2>
; CHECK-NEXT: ret <2 x i32> [[MUL]]
;
%shl.mask = and <2 x i32> %x, <i32 1073741823, i32 1073741823>
%tobool = icmp eq <2 x i32> %shl.mask, zeroinitializer
Expand All @@ -35,11 +29,8 @@ define <2 x i32> @test_eq_vect(<2 x i32> %x) {

define i32 @test_ne(i32 %x) {
; CHECK-LABEL: @test_ne(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL_NOT_NOT:%.*]] = icmp eq i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL_NOT_NOT]], i32 0, i32 [[MUL]]
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool.not = icmp ne i32 %shl.mask, 0
Expand All @@ -50,11 +41,8 @@ define i32 @test_ne(i32 %x) {

define <2 x i32> @test_ne_vect(<2 x i32> %x) {
; CHECK-LABEL: @test_ne_vect(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and <2 x i32> [[X:%.*]], <i32 1073741823, i32 1073741823>
; CHECK-NEXT: [[TOBOOL_NOT_NOT:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], <i32 2, i32 2>
; CHECK-NEXT: [[COND:%.*]] = select <2 x i1> [[TOBOOL_NOT_NOT]], <2 x i32> zeroinitializer, <2 x i32> [[MUL]]
; CHECK-NEXT: ret <2 x i32> [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X:%.*]], <i32 2, i32 2>
; CHECK-NEXT: ret <2 x i32> [[MUL]]
;
%shl.mask = and <2 x i32> %x, <i32 1073741823, i32 1073741823>
%tobool.not = icmp ne <2 x i32> %shl.mask, zeroinitializer
Expand All @@ -65,11 +53,8 @@ define <2 x i32> @test_ne_vect(<2 x i32> %x) {

define i32 @test_nuw_dropped(i32 %x) {
; CHECK-LABEL: @test_nuw_dropped(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl nuw i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool = icmp eq i32 %shl.mask, 0
Expand All @@ -80,11 +65,8 @@ define i32 @test_nuw_dropped(i32 %x) {

define i32 @test_nsw_dropped(i32 %x) {
; CHECK-LABEL: @test_nsw_dropped(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl nsw i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X:%.*]], 2
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool = icmp eq i32 %shl.mask, 0
Expand All @@ -101,9 +83,8 @@ define i32 @test_multi_use(i32 %x) {
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL_NOT:%.*]] = icmp ne i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL_NOT]], i32 [[MUL]], i32 0
; CHECK-NEXT: call void @use_multi(i32 [[SHL_MASK]], i1 [[TOBOOL_NOT]], i32 [[MUL]])
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool.not = icmp ne i32 %shl.mask, 0
Expand All @@ -117,10 +98,9 @@ define i32 @test_multi_use_nuw_dropped(i32 %x) {
; CHECK-LABEL: @test_multi_use_nuw_dropped(
; CHECK-NEXT: [[SHL_MASK:%.*]] = and i32 [[X:%.*]], 1073741823
; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq i32 [[SHL_MASK]], 0
; CHECK-NEXT: [[MUL:%.*]] = shl nuw i32 [[X]], 2
; CHECK-NEXT: [[COND:%.*]] = select i1 [[TOBOOL]], i32 0, i32 [[MUL]]
; CHECK-NEXT: [[MUL:%.*]] = shl i32 [[X]], 2
; CHECK-NEXT: call void @use_multi(i32 [[SHL_MASK]], i1 [[TOBOOL]], i32 [[MUL]])
; CHECK-NEXT: ret i32 [[COND]]
; CHECK-NEXT: ret i32 [[MUL]]
;
%shl.mask = and i32 %x, 1073741823
%tobool = icmp eq i32 %shl.mask, 0
Expand Down

0 comments on commit ab73bd3

Please sign in to comment.