[InstCombine] Fold (x >> y) << y -> x & (-1 << y)

[lebedev.ri]

Summary

We already do it for matching splat constants, but not just values.

Further improvements for non-matching splat constants, as noted in
https://reviews.llvm.org/D46760#1123713 will be needed,
but i'd prefer to do that as a follow-up.

https://bugs.llvm.org/show_bug.cgi?id=37603
https://rise4fun.com/Alive/cplX
https://rise4fun.com/Alive/0HF

[lebedev.ri]

Rebase ontop of more tests, simplify matcher, support ashr too.

[spatel]

LGTM.

lib/Transforms/InstCombine/InstCombineShifts.cpp

608

Remove the duplicate/shadow variable.

[lebedev.ri]

In D47981#1127520, @spatel wrote:

LGTM.

Thank you for the review.

Even in light of https://reviews.llvm.org/D46760#1127287, i believe we still can do *this* transform, ...

lib/Transforms/InstCombine/InstCombineShifts.cpp

608

Oh, hm, i wonder why there was no -Wshadow warning.

620–624

... because we already do it here for when the y is constant, regardless of exactness of shr.
Also, maybe this block should be dropped, now that the more general fold will be done, one not limited to constants?

[spatel]

In D47981#1127526, @lebedev.ri wrote:

Even in light of https://reviews.llvm.org/D46760#1127287, i believe we still can do *this* transform, ...

Agreed. I don't think a variable shift amount is a problem, but if it is, then maybe we should enhance SCEV rather than cripple instcombine? (And I think the same could be said for the constant shift amount).

lib/Transforms/InstCombine/InstCombineShifts.cpp

620–624

Right - I saw that block. I think it's a just a matter of taste because we don't need the one-use check if the shift amount is constant (so you'd have to add an 'if' for that possibility in the later block).

[lebedev.ri]

In D47981#1127538, @spatel wrote:

In D47981#1127526, @lebedev.ri wrote:

Even in light of https://reviews.llvm.org/D46760#1127287, i believe we still can do *this* transform, ...

Agreed. I don't think a variable shift amount is a problem, but if it is, then maybe we should enhance SCEV rather than cripple instcombine?

Yep, absolutely, that is why i CC'd Max. That crippling is just wrong.
To point the obvious, it just kinda ignores all the cases that already were shifts.

(And I think the same could be said for the constant shift amount).

lib/Transforms/InstCombine/InstCombineShifts.cpp

620–624

Oh, right, good point. Will keep it then.

lib/Transforms/InstCombine/InstCombineShifts.cpp

if (Value *V =
     return replaceInstUsesWith(I, V);
 
   if (Instruction *X = foldShuffledBinop(I))
     return X;
 
   if (Instruction *V = commonShiftTransforms(I))
     return V;
 
+  Type *Ty = I.getType();
   const APInt *ShAmtAPInt;
   if (match(Op1, m_APInt(ShAmtAPInt))) {
     unsigned ShAmt = ShAmtAPInt->getZExtValue();
     unsigned BitWidth = I.getType()->getScalarSizeInBits();
     Type *Ty = I.getType();
 
     // shl (zext X), ShAmt --> zext (shl X, ShAmt)
     // This is only valid if X would have zeros shifted out.
     Value *X;
     if (match(Op0, m_ZExt(m_Value(X)))) {
       unsigned SrcWidth = X->getType()->getScalarSizeInBits();
       if (ShAmt < SrcWidth &&
           MaskedValueIsZero(X, APInt::getHighBitsSet(SrcWidth, ShAmt), 0, &I))
         return new ZExtInst(Builder.CreateShl(X, ShAmt), Ty);
     }
 
     // (X >> C) << C --> X & (-1 << C)
     if (match(Op0, m_Shr(m_Value(X), m_Specific(Op1)))) {
       APInt Mask(APInt::getHighBitsSet(BitWidth, BitWidth - ShAmt));
       return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, Mask));
     }
 
     // Be careful about hiding shl instructions behind bit masks. They are used
     // to represent multiplies by a constant, and it is important that simple
     // arithmetic expressions are still recognizable by scalar evolution.
     // The inexact versions are deferred to DAGCombine, so we don't hide shl
     // behind a bit mask.
     const APInt *ShOp1;
     if (match(Op0, m_Exact(m_Shr(m_Value(X), m_APInt(ShOp1))))) {
if (match(Op1, m_APInt(ShAmtAPInt))) {
 
     // If the shifted-out value is all signbits, then this is a NSW shift.
     if (!I.hasNoSignedWrap() && ComputeNumSignBits(Op0, 0, &I) > ShAmt) {
       I.setHasNoSignedWrap();
       return &I;
     }
   }
 
+  // Transform  (x >> y) << y  to  x & (-1 << y)
+  // Valid for any type of right-shift.
+  Value *X;
+  if (match(Op0, m_OneUse(m_Shr(m_Value(X), m_Specific(Op1))))) {
+    Constant *AllOnes = ConstantInt::getAllOnesValue(Ty);
+    Value *Mask = Builder.CreateShl(AllOnes, Op1);
+    return BinaryOperator::CreateAnd(Mask, X);
+  }
+
   Constant *C1;
   if (match(Op1, m_Constant(C1))) {
     Constant *C2;
     Value *X;
     // (C2 << X) << C1 --> (C2 << C1) << X
     if (match(Op0, m_OneUse(m_Shl(m_Constant(C2), m_Value(X)))))
       return BinaryOperator::CreateShl(ConstantExpr::getShl(C2, C1), X);
 

test/Transforms/InstCombine/canonicalize-ashr-shl-to-masking.ll

 ;   x & (-1 << y)
 
 ; ============================================================================ ;
 ; Basic positive tests
 ; ============================================================================ ;
 
 define i32 @positive_samevar(i32 %x, i32 %y) {
 ; CHECK-LABEL: @positive_samevar(
-; CHECK-NEXT:    [[TMP0:%.*]] = ashr i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = shl i32 [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 -1, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and i32 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %tmp0 = ashr i32 %x, %y
   %ret = shl i32 %tmp0, %y
   ret i32 %ret
 }
 
 define i32 @positive_sameconst(i32 %x) {

 }
 
 ; ============================================================================ ;
 ; Vector
 ; ============================================================================ ;
 
 define <2 x i32> @positive_samevar_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @positive_samevar_vec(
-; CHECK-NEXT:    [[TMP0:%.*]] = ashr <2 x i32> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = shl <2 x i32> [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl <2 x i32> <i32 -1, i32 -1>, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and <2 x i32> [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[RET]]
 ;
   %tmp0 = ashr <2 x i32> %x, %y
   %ret = shl <2 x i32> %tmp0, %y
   ret <2 x i32> %ret
 }
 
 ; ============================================================================ ;

 ;
   %tmp0 = ashr <3 x i32> %x, <i32 5, i32 5, i32 5>
   %ret = shl <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <3 x i32> @positive_sameconst_vec_undef2(<3 x i32> %x) {
 ; CHECK-LABEL: @positive_sameconst_vec_undef2(
-; CHECK-NEXT:    [[TMP0:%.*]] = ashr <3 x i32> [[X:%.*]], <i32 5, i32 undef, i32 5>
-; CHECK-NEXT:    [[RET:%.*]] = shl <3 x i32> [[TMP0]], <i32 5, i32 undef, i32 5>
+; CHECK-NEXT:    [[RET:%.*]] = and <3 x i32> [[X:%.*]], <i32 -32, i32 undef, i32 -32>
 ; CHECK-NEXT:    ret <3 x i32> [[RET]]
 ;
   %tmp0 = ashr <3 x i32> %x, <i32 5, i32 undef, i32 5>
   %ret = shl <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <2 x i32> @positive_biggerashr_vec(<2 x i32> %x) {

test/Transforms/InstCombine/canonicalize-lshr-shl-to-masking.ll

 ;   x & (-1 << y)
 
 ; ============================================================================ ;
 ; Basic positive tests
 ; ============================================================================ ;
 
 define i32 @positive_samevar(i32 %x, i32 %y) {
 ; CHECK-LABEL: @positive_samevar(
-; CHECK-NEXT:    [[TMP0:%.*]] = lshr i32 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = shl i32 [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 -1, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and i32 [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %tmp0 = lshr i32 %x, %y
   %ret = shl i32 %tmp0, %y
   ret i32 %ret
 }
 
 define i32 @positive_sameconst(i32 %x) {

 }
 
 ; ============================================================================ ;
 ; Vector
 ; ============================================================================ ;
 
 define <2 x i32> @positive_samevar_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @positive_samevar_vec(
-; CHECK-NEXT:    [[TMP0:%.*]] = lshr <2 x i32> [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[RET:%.*]] = shl <2 x i32> [[TMP0]], [[Y]]
+; CHECK-NEXT:    [[TMP1:%.*]] = shl <2 x i32> <i32 -1, i32 -1>, [[Y:%.*]]
+; CHECK-NEXT:    [[RET:%.*]] = and <2 x i32> [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[RET]]
 ;
   %tmp0 = lshr <2 x i32> %x, %y
   %ret = shl <2 x i32> %tmp0, %y
   ret <2 x i32> %ret
 }
 
 ; ============================================================================ ;

 ;
   %tmp0 = lshr <3 x i32> %x, <i32 5, i32 5, i32 5>
   %ret = shl <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <3 x i32> @positive_sameconst_vec_undef2(<3 x i32> %x) {
 ; CHECK-LABEL: @positive_sameconst_vec_undef2(
-; CHECK-NEXT:    [[TMP0:%.*]] = lshr <3 x i32> [[X:%.*]], <i32 5, i32 undef, i32 5>
-; CHECK-NEXT:    [[RET:%.*]] = shl <3 x i32> [[TMP0]], <i32 5, i32 undef, i32 5>
+; CHECK-NEXT:    [[RET:%.*]] = and <3 x i32> [[X:%.*]], <i32 -32, i32 undef, i32 -32>
 ; CHECK-NEXT:    ret <3 x i32> [[RET]]
 ;
   %tmp0 = lshr <3 x i32> %x, <i32 5, i32 undef, i32 5>
   %ret = shl <3 x i32> %tmp0, <i32 5, i32 undef, i32 5>
   ret <3 x i32> %ret
 }
 
 define <2 x i32> @positive_biggerlshr_vec(<2 x i32> %x) {