[InstCombine] Optimize lshr+shl+and conversion pattern

[bcl5980]

Summary

if C1 and C3 are pow2 and Log2(C3) >= C2:

((C1 >> X) << C2) & C3 -> X == (Log2(C1)+C2-Log2(C3)) ? C3 : 0

https://alive2.llvm.org/ce/z/zvrkKF

[bcl5980]

update test

[bcl5980]

C3 support non-uniform constant for shl+lshr and lshr+shl
update tests.

[bcl5980]

add nullptr check

[spatel]

I don't like increasing risk by combining what are really 2 independent changes into 1 patch.
Please go back to the earlier revision that only adds the new fold. We can ease the vector restrictions as a second step (I haven't seen any evidence that those cases matter, but it opens up more possibilities to miscompile in ways that are very hard to detect).

[bcl5980]

revert to diff2

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

if (match(Op0, m_OneUse(m_BinOp(BO))) && isSuitableBinOpcode(BO)) {
           C->isMask(X->getType()->getScalarSizeInBits())) {
         Y = BO->getOperand(0);
         Value *TrY = Builder.CreateTrunc(Y, X->getType(), Y->getName() + ".tr");
         Value *NewBO =
             Builder.CreateBinOp(BOpcode, TrY, X, BO->getName() + ".narrow");
         return new ZExtInst(NewBO, Ty);
       }
     }
+  }
 
+  if (match(Op1, m_Power2())) {
     Constant *C1, *C2;
-    const APInt *C3 = C;
+    Constant *C3 = cast<Constant>(Op1);
     Value *X;
-    if (C3->isPowerOf2() &&
-        match(Op0, m_OneUse(m_LShr(m_Shl(m_ImmConstant(C1), m_Value(X)),
+    Constant *Log2C3 = ConstantExpr::getExactLogBase2(C3);
+    if (match(Op0, m_OneUse(m_LShr(m_Shl(m_ImmConstant(C1), m_Value(X)),
                                    m_ImmConstant(C2)))) &&
         match(C1, m_Power2())) {
       Constant *Log2C1 = ConstantExpr::getExactLogBase2(C1);
-      Constant *Log2C3 = ConstantInt::get(Ty, C3->countTrailingZeros());
       Constant *LshrC = ConstantExpr::getAdd(C2, Log2C3);
       KnownBits KnownLShrc = computeKnownBits(LshrC, 0, nullptr);
+      unsigned Width = Ty->getScalarSizeInBits();
       if (KnownLShrc.getMaxValue().ult(Width)) {
         // iff C1,C3 is pow2 and C2 + cttz(C3) < BitWidth:
         // ((C1 << X) >> C2) & C3 -> X == (cttz(C3)+C2-cttz(C1)) ? C3 : 0
         Constant *CmpC = ConstantExpr::getSub(LshrC, Log2C1);
         Value *Cmp = Builder.CreateICmpEQ(X, CmpC);
-        return SelectInst::Create(Cmp, ConstantInt::get(Ty, *C3),
-                                  ConstantInt::getNullValue(Ty));
-      }
-      // TODO: Symmetrical case
+        return SelectInst::Create(Cmp, C3, ConstantInt::getNullValue(Ty));
+      }
+
+    if (match(Op0, m_OneUse(m_Shl(m_LShr(m_ImmConstant(C1), m_Value(X)),
+                                  m_ImmConstant(C2)))) &&
+        match(C1, m_Power2())) {
+      Constant *Log2C1 = ConstantExpr::getExactLogBase2(C1);
+      Constant *Cmp = ConstantExpr::getCompare(ICmpInst::ICMP_ULT, Log2C3, C2);
+      if (Cmp && Cmp->isZeroValue()) {
         // iff C1,C3 is pow2 and Log2(C3) >= C2:
         // ((C1 >> X) << C2) & C3 -> X == (cttz(C1)+C2-cttz(C3)) ? C3 : 0
+        Constant *ShlC = ConstantExpr::getAdd(C2, Log2C1);
+        Constant *CmpC = ConstantExpr::getSub(ShlC, Log2C3);
+        Value *Cmp = Builder.CreateICmpEQ(X, CmpC);
+        return SelectInst::Create(Cmp, C3, ConstantInt::getNullValue(Ty));
+      }
     }
   }
 
   if (match(&I, m_And(m_OneUse(m_Shl(m_ZExt(m_Value(X)), m_Value(Y))),
                       m_SignMask())) &&
       match(Y, m_SpecificInt_ICMP(
                    ICmpInst::Predicate::ICMP_EQ,
                    APInt(Ty->getScalarSizeInBits(),

llvm/test/Transforms/InstCombine/and.ll

;
   %shl = shl <3 x i16> <i16 4, i16 4, i16 4>, %x
   %lshr = lshr <3 x i16> %shl, <i16 6, i16 6, i16 6>
   %r = and <3 x i16> %lshr, <i16 8, i16 8, i16 8>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @shl_lshr_pow2_const_case1_non_uniform_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @shl_lshr_pow2_const_case1_non_uniform_vec(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> <i16 2, i16 8, i16 32>, [[X:%.*]]
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> [[SHL]], <i16 5, i16 6, i16 3>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[LSHR]], <i16 8, i16 4, i16 8>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 7, i16 5, i16 1>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 8, i16 4, i16 8>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %shl = shl <3 x i16> <i16 2, i16 8, i16 32>, %x
   %lshr = lshr <3 x i16> %shl, <i16 5, i16 6, i16 3>
   %r = and <3 x i16> %lshr, <i16 8, i16 4, i16 8>
   ret <3 x i16> %r
 }
 
;
   %shl = shl <3 x i16> <i16 16, i16 16, i16 16>, %x
   %lshr = lshr <3 x i16> %shl, <i16 undef, i16 5, i16 5>
   %r = and <3 x i16> %lshr, <i16 8, i16 8, i16 8>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @shl_lshr_pow2_const_case1_undef3_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @shl_lshr_pow2_const_case1_undef3_vec(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> <i16 16, i16 16, i16 16>, [[X:%.*]]
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> [[SHL]], <i16 5, i16 5, i16 5>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[LSHR]], <i16 undef, i16 8, i16 8>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 1, i16 4, i16 4>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 undef, i16 8, i16 8>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %shl = shl <3 x i16> <i16 16, i16 16, i16 16>, %x
   %lshr = lshr <3 x i16> %shl, <i16 5, i16 5, i16 5>
   %r = and <3 x i16> %lshr, <i16 undef, i16 8, i16 8>
   ret <3 x i16> %r
 }
 
;
   %r = xor i16 %and, 8
   ret i16 %r
 }
 
 ; Log2(LshrC) + ShlC >= BitWidth
 
 define i16 @lshr_shl_pow2_const_case2(i16 %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case2(
-; CHECK-NEXT:    [[LSHR1:%.*]] = lshr i16 8192, [[X:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl i16 [[LSHR1]], 4
-; CHECK-NEXT:    [[R:%.*]] = and i16 [[SHL]], 32
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i16 [[X:%.*]], 12
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP1]], i16 32, i16 0
 ; CHECK-NEXT:    ret i16 [[R]]
 ;
   %lshr1 = lshr i16 8192, %x
   %shl = shl i16 %lshr1, 4
   %r = and i16 %shl, 32
   ret i16 %r
 }
 
;
   %shl = shl i16 %lshr1, 4
   call void @use16(i16 %shl)
   %r = and i16 %shl, 32
   ret i16 %r
 }
 
 define <3 x i16> @lshr_shl_pow2_const_case1_uniform_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1_uniform_vec(
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> <i16 8192, i16 8192, i16 8192>, [[X:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> [[LSHR]], <i16 6, i16 6, i16 6>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[SHL]], <i16 128, i16 128, i16 128>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 12, i16 12, i16 12>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 128, i16 128, i16 128>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %lshr = lshr <3 x i16> <i16 8192, i16 8192, i16 8192>, %x
   %shl = shl <3 x i16> %lshr, <i16 6, i16 6, i16 6>
   %r = and <3 x i16> %shl, <i16 128, i16 128, i16 128>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @lshr_shl_pow2_const_case1_non_uniform_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1_non_uniform_vec(
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> <i16 8192, i16 16384, i16 -32768>, [[X:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> [[LSHR]], <i16 7, i16 5, i16 3>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[SHL]], <i16 128, i16 256, i16 512>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 13, i16 11, i16 9>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 128, i16 256, i16 512>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %lshr = lshr <3 x i16> <i16 8192, i16 16384, i16 32768>, %x
   %shl = shl <3 x i16> %lshr, <i16 7, i16 5, i16 3>
   %r = and <3 x i16> %shl, <i16 128, i16 256, i16 512>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @lshr_shl_pow2_const_case1_non_uniform_vec_negative(<3 x i16> %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1_non_uniform_vec_negative(
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> <i16 8192, i16 16384, i16 -32768>, [[X:%.*]]
 ; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> [[LSHR]], <i16 8, i16 5, i16 3>
 ; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[SHL]], <i16 128, i16 256, i16 512>
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %lshr = lshr <3 x i16> <i16 8192, i16 16384, i16 32768>, %x
   %shl = shl <3 x i16> %lshr, <i16 8, i16 5, i16 3>
   %r = and <3 x i16> %shl, <i16 128, i16 256, i16 512>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @lshr_shl_pow2_const_case1_undef1_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1_undef1_vec(
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> <i16 undef, i16 8192, i16 8192>, [[X:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> [[LSHR]], <i16 6, i16 6, i16 6>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[SHL]], <i16 128, i16 128, i16 128>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 -1, i16 12, i16 12>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 128, i16 128, i16 128>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %lshr = lshr <3 x i16> <i16 undef, i16 8192, i16 8192>, %x
   %shl = shl <3 x i16> %lshr, <i16 6, i16 6, i16 6>
   %r = and <3 x i16> %shl, <i16 128, i16 128, i16 128>
   ret <3 x i16> %r
 }
 
 define <3 x i16> @lshr_shl_pow2_const_case1_undef2_vec(<3 x i16> %x) {
 ; CHECK-LABEL: @lshr_shl_pow2_const_case1_undef2_vec(
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i16> <i16 8192, i16 8192, i16 8192>, [[X:%.*]]
-; CHECK-NEXT:    [[SHL:%.*]] = shl <3 x i16> [[LSHR]], <i16 undef, i16 6, i16 6>
-; CHECK-NEXT:    [[R:%.*]] = and <3 x i16> [[SHL]], <i16 128, i16 128, i16 128>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq <3 x i16> [[X:%.*]], <i16 undef, i16 12, i16 12>
+; CHECK-NEXT:    [[R:%.*]] = select <3 x i1> [[TMP1]], <3 x i16> <i16 128, i16 128, i16 128>, <3 x i16> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i16> [[R]]
 ;
   %lshr = lshr <3 x i16> <i16 8192, i16 8192, i16 8192>, %x
   %shl = shl <3 x i16> %lshr, <i16 undef, i16 6, i16 6>
   %r = and <3 x i16> %shl, <i16 128, i16 128, i16 128>
   ret <3 x i16> %r
 }