[InstCombine] fold low-bit mask of sext-in-reg

[spatel]

Summary

This is an alternative to D91343. I do not have stats for how often the shl+ashr pattern (sext-in-reg) occurs, but we have an SDAG node for that operation, so it seems common enough to warrant a dedicated fold for this pattern:
https://rise4fun.com/Alive/vI9

[lebedev.ri]

I mildly prefer the other approach (D91343), because it doesn't hardcode the and instruction.

[spatel]

In D91364#2391743, @lebedev.ri wrote:

I mildly prefer the other approach (D91343), because it doesn't hardcode the and instruction.

Right, that's more flexible with potentially more compute known bits cost.
@Bhramar.vatsa - I committed the baseline tests shown here with 8a1e636. Please update/rebase your patch and confirm that you get the same diffs. I think you can reduce the code a bit too based on the match code shown here.

[lebedev.ri]

In D91364#2391784, @spatel wrote:

In D91364#2391743, @lebedev.ri wrote:

I mildly prefer the other approach (D91343), because it doesn't hardcode the and instruction.

Right, that's more flexible with potentially more compute known bits cost.

Ack.
But i think the pattern isn't *that* common, i would be surprised if this fires more than ~1k times as of vanilla test-suite.

@Bhramar.vatsa - I committed the baseline tests shown here with 8a1e636. Please update/rebase your patch and confirm that you get the same diffs. I think you can reduce the code a bit too based on the match code shown here.

[spatel]

Abandoning in favor of the more general D91343.

llvm/lib/Transforms/InstCombine/InstCombineAndOrXor.cpp

if (match(Op0, m_OneUse(m_Xor(m_Value(X), m_Value(Y)))) ||
       }
       if (!isa<Constant>(Y) && MaskedValueIsZero(Y, NotAndMask, 0, &I)) {
         // Not masking anything out for the RHS, move mask to LHS.
         // and ({x}or X, Y), C --> {x}or (and X, C), Y
         Value *NewLHS = Builder.CreateAnd(X, Op1, X->getName() + ".masked");
         return BinaryOperator::Create(BinOp, NewLHS, Y);
       }
     }
-    const APInt *ShiftC;
-    if (match(Op0, m_OneUse(m_SExt(m_AShr(m_Value(X), m_APInt(ShiftC)))))) {
-      unsigned Width = Ty->getScalarSizeInBits();
+    const APInt *ShiftC, *ShiftC2;
+    if (match(Op0, m_OneUse(m_SExt(m_AShr(m_Value(X), m_APInt(ShiftC)))))) {
       if (*C == APInt::getLowBitsSet(Width, Width - ShiftC->getZExtValue())) {
         // We are clearing high bits that were potentially set by sext+ashr:
         // and (sext (ashr X, ShiftC)), C --> lshr (sext X), ShiftC
         Value *Sext = Builder.CreateSExt(X, Ty);
         Constant *ShAmtC = ConstantInt::get(Ty, ShiftC->zext(Width));
         return BinaryOperator::CreateLShr(Sext, ShAmtC);
       }
     }
+    // Peek through masking off the low bits of a sign-extend-in-register:
+    // and (ashr (shl X, ShiftC), ShiftC), C --> and X, C
+    if (match(Op0,
+              m_AShr(m_Shl(m_Value(X), m_APInt(ShiftC)), m_APInt(ShiftC2))) &&
+        ShiftC == ShiftC2 && (Width - ShiftC->getZExtValue()) > C->logBase2())
+      return BinaryOperator::CreateAnd(X, Op1);
   }
 
   ConstantInt *AndRHS;
   if (match(Op1, m_ConstantInt(AndRHS))) {
     const APInt &AndRHSMask = AndRHS->getValue();
 
     // Optimize a variety of ((val OP C1) & C2) combinations...
     if (BinaryOperator *Op0I = dyn_cast<BinaryOperator>(Op0)) {

llvm/test/Transforms/InstCombine/and.ll

;
   ret <2 x i59> %r
 }
 
 define i32 @lowmask_sext_in_reg(i32 %x) {
 ; CHECK-LABEL: @lowmask_sext_in_reg(
 ; CHECK-NEXT:    [[L:%.*]] = shl i32 [[X:%.*]], 20
 ; CHECK-NEXT:    [[R:%.*]] = ashr exact i32 [[L]], 20
 ; CHECK-NEXT:    call void @use32(i32 [[R]])
-; CHECK-NEXT:    [[AND:%.*]] = and i32 [[R]], 4095
+; CHECK-NEXT:    [[AND:%.*]] = and i32 [[X]], 4095
 ; CHECK-NEXT:    ret i32 [[AND]]
 ;
   %l = shl i32 %x, 20
   %r = ashr i32 %l, 20
   call void @use32(i32 %r)
   %and = and i32 %r, 4095
   ret i32 %and
 }
 
+; Negative test - mismatched shift amounts
+
 define i32 @lowmask_not_sext_in_reg(i32 %x) {
 ; CHECK-LABEL: @lowmask_not_sext_in_reg(
 ; CHECK-NEXT:    [[L:%.*]] = shl i32 [[X:%.*]], 19
 ; CHECK-NEXT:    [[R:%.*]] = ashr i32 [[L]], 20
 ; CHECK-NEXT:    call void @use32(i32 [[R]])
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[R]], 4095
 ; CHECK-NEXT:    ret i32 [[AND]]
 ;
   %l = shl i32 %x, 19
   %r = ashr i32 %l, 20
   call void @use32(i32 %r)
   %and = and i32 %r, 4095
   ret i32 %and
 }
 
+; Negative test - too much shift for mask
+
 define i32 @not_lowmask_sext_in_reg(i32 %x) {
 ; CHECK-LABEL: @not_lowmask_sext_in_reg(
 ; CHECK-NEXT:    [[L:%.*]] = shl i32 [[X:%.*]], 20
 ; CHECK-NEXT:    [[R:%.*]] = ashr exact i32 [[L]], 20
 ; CHECK-NEXT:    call void @use32(i32 [[R]])
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[R]], 4096
 ; CHECK-NEXT:    ret i32 [[AND]]
 ;
   %l = shl i32 %x, 20
   %r = ashr i32 %l, 20
   call void @use32(i32 %r)
   %and = and i32 %r, 4096
   ret i32 %and
 }
 
+; Negative test - too much shift for mask
+
 define i32 @not_lowmask_sext_in_reg2(i32 %x) {
 ; CHECK-LABEL: @not_lowmask_sext_in_reg2(
 ; CHECK-NEXT:    [[L:%.*]] = shl i32 [[X:%.*]], 21
 ; CHECK-NEXT:    [[R:%.*]] = ashr exact i32 [[L]], 21
 ; CHECK-NEXT:    call void @use32(i32 [[R]])
 ; CHECK-NEXT:    [[AND:%.*]] = and i32 [[R]], 4095
 ; CHECK-NEXT:    ret i32 [[AND]]
 ;
   %l = shl i32 %x, 21
   %r = ashr i32 %l, 21
   call void @use32(i32 %r)
   %and = and i32 %r, 4095
   ret i32 %and
 }
 
 define <2 x i32> @lowmask_sext_in_reg_splat(<2 x i32> %x, <2 x i32>* %p) {
 ; CHECK-LABEL: @lowmask_sext_in_reg_splat(
 ; CHECK-NEXT:    [[L:%.*]] = shl <2 x i32> [[X:%.*]], <i32 20, i32 20>
 ; CHECK-NEXT:    [[R:%.*]] = ashr exact <2 x i32> [[L]], <i32 20, i32 20>
 ; CHECK-NEXT:    store <2 x i32> [[R]], <2 x i32>* [[P:%.*]], align 8
-; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[R]], <i32 4095, i32 4095>
+; CHECK-NEXT:    [[AND:%.*]] = and <2 x i32> [[X]], <i32 4095, i32 4095>
 ; CHECK-NEXT:    ret <2 x i32> [[AND]]
 ;
   %l = shl <2 x i32> %x, <i32 20, i32 20>
   %r = ashr <2 x i32> %l, <i32 20, i32 20>
   store <2 x i32> %r, <2 x i32>* %p
   %and = and <2 x i32> %r, <i32 4095, i32 4095>
   ret <2 x i32> %and
 }