diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp --- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp +++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp @@ -609,6 +609,44 @@ 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(C2->getZExtValue())) + return nullptr; + + auto *FI = dyn_cast(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 @@ -1743,6 +1781,9 @@ 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; diff --git a/llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll b/llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll --- a/llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll +++ b/llvm/test/Transforms/InstCombine/select-icmp-and-zero-shl.ll @@ -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 @@ -20,11 +17,8 @@ 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:%.*]], -; CHECK-NEXT: [[TOBOOL:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer -; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], -; 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:%.*]], +; CHECK-NEXT: ret <2 x i32> [[MUL]] ; %shl.mask = and <2 x i32> %x, %tobool = icmp eq <2 x i32> %shl.mask, zeroinitializer @@ -35,11 +29,8 @@ 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 @@ -50,11 +41,8 @@ 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:%.*]], -; CHECK-NEXT: [[TOBOOL_NOT_NOT:%.*]] = icmp eq <2 x i32> [[SHL_MASK]], zeroinitializer -; CHECK-NEXT: [[MUL:%.*]] = shl <2 x i32> [[X]], -; 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:%.*]], +; CHECK-NEXT: ret <2 x i32> [[MUL]] ; %shl.mask = and <2 x i32> %x, %tobool.not = icmp ne <2 x i32> %shl.mask, zeroinitializer @@ -65,11 +53,8 @@ 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 @@ -80,11 +65,8 @@ 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 @@ -101,9 +83,8 @@ ; 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 @@ -117,10 +98,9 @@ ; 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