[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_Xor(m_Value(X), m_Value(Y)))) ||
         Value *NewLHS = Builder.CreateAnd(X, Op1, X->getName() + ".masked");
         return BinaryOperator::Create(BinOp, NewLHS, Y);
       }
     }
 
     Constant *C1, *C2;
     const APInt *C3 = C;
     Value *X;
-    if (C3->isPowerOf2() &&
-        match(Op0, m_OneUse(m_LShr(m_Shl(m_ImmConstant(C1), m_Value(X)),
+    if (C3->isPowerOf2()) {
+      Constant *Log2C3 = ConstantInt::get(Ty, C3->countTrailingZeros());
+      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);
         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
+      }
+
+      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->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, ConstantInt::get(Ty, *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

;
   %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
 }
 
;
   %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
 }