[InstCombine] SimplifyDemandedUseBits - remove lshr node if we only demand known sign bit

[RKSimon]

Summary

This is a lshr equivalent to D122340 - if we don't demand any of the additional sign bits introduced by the ashr, the lshr can be treated as an ashr and we can remove the shift entirely if we only demand an already known sign bit.

If people prefer, I can simplify the code further, but I figured it made sense to try to keep it similar to the ashr fold for comparison.

Another step towards PR21929

https://alive2.llvm.org/ce/z/6f3kjq

[RKSimon]

We don't have to explicitly demand the signbit for this to work.

[lebedev.ri]

LGTM, thanks.

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

556

I believe the same applies here.

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (I->getOperand(0) == I->getOperand(1) && DemandedMask.ult(4)) {
       Constant *One = ConstantInt::get(VTy, 1);
       Instruction *And1 = BinaryOperator::CreateAnd(I->getOperand(0), One);
       return InsertNewInstWith(And1, *I);
     }
 
     computeKnownBits(I, Known, Depth, CxtI);
     break;
   }
   case Instruction::Shl: {
     const APInt *SA;
     if (match(I->getOperand(1), m_APInt(SA))) {
       const APInt *ShrAmt;
       if (match(I->getOperand(0), m_Shr(m_Value(), m_APInt(ShrAmt))))
         if (Instruction *Shr = dyn_cast<Instruction>(I->getOperand(0)))
           if (Value *R = simplifyShrShlDemandedBits(Shr, *ShrAmt, I, *SA,
                                                     DemandedMask, Known))
             return R;
case Instruction::Shl: {
     }
     break;
   }
   case Instruction::LShr: {
     const APInt *SA;
     if (match(I->getOperand(1), m_APInt(SA))) {
       uint64_t ShiftAmt = SA->getLimitedValue(BitWidth-1);
 
+      // If we are just demanding the shifted sign bit and below, then this can
+      // be treated as an ASHR in disguise.
+      if (DemandedMask.countLeadingZeros() >= ShiftAmt) {
+        // If we only want bits that already match the signbit then we don't
+        // need to shift.
+        unsigned NumHiDemandedBits =
+            BitWidth - DemandedMask.countTrailingZeros();
+        unsigned SignBits =
+            ComputeNumSignBits(I->getOperand(0), Depth + 1, CxtI);
+        if (SignBits >= NumHiDemandedBits)
+          return I->getOperand(0);
+      }
+
       // Unsigned shift right.
       APInt DemandedMaskIn(DemandedMask.shl(ShiftAmt));
 
       // If the shift is exact, then it does demand the low bits (and knows that
       // they are zero).
       if (cast<LShrOperator>(I)->isExact())
         DemandedMaskIn.setLowBits(ShiftAmt);
 

llvm/test/Transforms/InstCombine/modulo.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -S -passes=instcombine | FileCheck %s
 
 define i32 @PR21929(i32 %x) {
 ; CHECK-LABEL: @PR21929(
 ; CHECK-NEXT:    [[REM_I:%.*]] = srem i32 [[X:%.*]], 2
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[REM_I]], 30
-; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], 2
-; CHECK-NEXT:    [[RET_I:%.*]] = add nsw i32 [[TMP2]], [[REM_I]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[REM_I]], 2
+; CHECK-NEXT:    [[RET_I:%.*]] = add nsw i32 [[TMP1]], [[REM_I]]
 ; CHECK-NEXT:    ret i32 [[RET_I]]
 ;
   %rem.i = srem i32 %x, 2
   %cmp.i = icmp slt i32 %rem.i, 0
   %add.i = select i1 %cmp.i, i32 2, i32 0
   %ret.i = add nsw i32 %add.i, %rem.i
   ret i32 %ret.i
 }
 
 define <2 x i32> @PR21929_vec(<2 x i32> %x) {
 ; CHECK-LABEL: @PR21929_vec(
 ; CHECK-NEXT:    [[REM_I:%.*]] = srem <2 x i32> [[X:%.*]], <i32 2, i32 2>
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[REM_I]], <i32 30, i32 30>
-; CHECK-NEXT:    [[TMP2:%.*]] = and <2 x i32> [[TMP1]], <i32 2, i32 2>
-; CHECK-NEXT:    [[RET_I:%.*]] = add nsw <2 x i32> [[TMP2]], [[REM_I]]
+; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[REM_I]], <i32 2, i32 2>
+; CHECK-NEXT:    [[RET_I:%.*]] = add nsw <2 x i32> [[TMP1]], [[REM_I]]
 ; CHECK-NEXT:    ret <2 x i32> [[RET_I]]
 ;
   %rem.i = srem <2 x i32> %x, <i32 2, i32 2>
   %cmp.i = icmp slt <2 x i32> %rem.i, zeroinitializer
   %add.i = select <2 x i1> %cmp.i, <2 x i32> <i32 2, i32 2>, <2 x i32> zeroinitializer
   %ret.i = add nsw <2 x i32> %add.i, %rem.i
   ret <2 x i32> %ret.i
 }