[InstCombine] SimplifyDemandedUseBits - remove ashr node if we only demand known sign bits

[RKSimon]

Summary

We already do this for SelectionDAG, but we're missing it here.

Noticed while re-triaging PR21929

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (match(I->getOperand(1), m_APInt(SA))) {
       if (ShiftAmt)
         Known.Zero.setHighBits(ShiftAmt);  // high bits known zero.
     } else {
       computeKnownBits(I, Known, Depth, CxtI);
     }
     break;
   }
   case Instruction::AShr: {
+    unsigned SignBits = ComputeNumSignBits(I->getOperand(0), Depth + 1, CxtI);
+
+    // If we only want bits that already match the signbit then we don't need
+    // to shift.
+    unsigned NumHiDemandedBits = BitWidth - DemandedMask.countTrailingZeros();
+    if (SignBits >= NumHiDemandedBits)
+      return I->getOperand(0);
+
     // If this is an arithmetic shift right and only the low-bit is set, we can
     // always convert this into a logical shr, even if the shift amount is
     // variable.  The low bit of the shift cannot be an input sign bit unless
     // the shift amount is >= the size of the datatype, which is undefined.
     if (DemandedMask.isOne()) {
       // Perform the logical shift right.
       Instruction *NewVal = BinaryOperator::CreateLShr(
                         I->getOperand(0), I->getOperand(1), I->getName());
       return InsertNewInstWith(NewVal, *I);
     }
 
-    // If the sign bit is the only bit demanded by this ashr, then there is no
-    // need to do it, the shift doesn't change the high bit.
-    if (DemandedMask.isSignMask())
-      return I->getOperand(0);
-
     const APInt *SA;
     if (match(I->getOperand(1), m_APInt(SA))) {
       uint32_t ShiftAmt = SA->getLimitedValue(BitWidth-1);
 
       // Signed shift right.
       APInt DemandedMaskIn(DemandedMask.shl(ShiftAmt));
       // If any of the high bits are demanded, we should set the sign bit as
       // demanded.
       if (DemandedMask.countLeadingZeros() <= ShiftAmt)
         DemandedMaskIn.setSignBit();
 
       // If the shift is exact, then it does demand the low bits (and knows that
       // they are zero).
       if (cast<AShrOperator>(I)->isExact())
         DemandedMaskIn.setLowBits(ShiftAmt);
 
       if (SimplifyDemandedBits(I, 0, DemandedMaskIn, Known, Depth + 1))
         return I;
 
-      unsigned SignBits = ComputeNumSignBits(I->getOperand(0), Depth + 1, CxtI);
-
       assert(!Known.hasConflict() && "Bits known to be one AND zero?");
       // Compute the new bits that are at the top now plus sign bits.
       APInt HighBits(APInt::getHighBitsSet(
           BitWidth, std::min(SignBits + ShiftAmt - 1, BitWidth)));
       Known.Zero.lshrInPlace(ShiftAmt);
       Known.One.lshrInPlace(ShiftAmt);
 
       // If the input sign bit is known to be zero, or if none of the top bits

llvm/test/Transforms/InstCombine/ashr-demand.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -passes=instcombine -S < %s | FileCheck %s
 
 ; If we only want bits that already match the signbit then we don't need to shift.
 
 define i32 @srem2_ashr_mask(i32 %a0) {
 ; CHECK-LABEL: @srem2_ashr_mask(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[A0:%.*]], -2147483647
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp eq i32 [[TMP1]], -2147483647
-; CHECK-NEXT:    [[MASK:%.*]] = select i1 [[ISNEG]], i32 2, i32 0
+; CHECK-NEXT:    [[SREM:%.*]] = srem i32 [[A0:%.*]], 2
+; CHECK-NEXT:    [[MASK:%.*]] = and i32 [[SREM]], 2
 ; CHECK-NEXT:    ret i32 [[MASK]]
 ;
   %srem = srem i32 %a0, 2 ; result = (1,0,-1) num signbits = 31
   %ashr = ashr i32 %srem, 31
   %mask = and i32 %ashr, 2
   ret i32 %mask
 }
 
+; Negative test - mask demands non-signbit from shift source
 define i32 @srem8_ashr_mask(i32 %a0) {
 ; CHECK-LABEL: @srem8_ashr_mask(
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[A0:%.*]], -2147483641
 ; CHECK-NEXT:    [[ISNEG:%.*]] = icmp ugt i32 [[TMP1]], -2147483648
 ; CHECK-NEXT:    [[MASK:%.*]] = select i1 [[ISNEG]], i32 2, i32 0
 ; CHECK-NEXT:    ret i32 [[MASK]]
 ;
   %srem = srem i32 %a0, 8
   %ashr = ashr i32 %srem, 31
   %mask = and i32 %ashr, 2
   ret i32 %mask
 }
 
 define <2 x i32> @srem2_ashr_mask_vector(<2 x i32> %a0) {
 ; CHECK-LABEL: @srem2_ashr_mask_vector(
-; CHECK-NEXT:    [[TMP1:%.*]] = and <2 x i32> [[A0:%.*]], <i32 -2147483647, i32 -2147483647>
-; CHECK-NEXT:    [[ISNEG:%.*]] = icmp eq <2 x i32> [[TMP1]], <i32 -2147483647, i32 -2147483647>
-; CHECK-NEXT:    [[MASK:%.*]] = select <2 x i1> [[ISNEG]], <2 x i32> <i32 2, i32 2>, <2 x i32> zeroinitializer
+; CHECK-NEXT:    [[SREM:%.*]] = srem <2 x i32> [[A0:%.*]], <i32 2, i32 2>
+; CHECK-NEXT:    [[MASK:%.*]] = and <2 x i32> [[SREM]], <i32 2, i32 2>
 ; CHECK-NEXT:    ret <2 x i32> [[MASK]]
 ;
   %srem = srem <2 x i32> %a0, <i32 2, i32 2>
   %ashr = ashr <2 x i32> %srem, <i32 31, i32 31>
   %mask = and <2 x i32> %ashr, <i32 2, i32 2>
   ret <2 x i32> %mask
 }
 
 define <2 x i32> @srem2_ashr_mask_vector_nonconstant(<2 x i32> %a0, <2 x i32> %a1) {
 ; CHECK-LABEL: @srem2_ashr_mask_vector_nonconstant(
 ; CHECK-NEXT:    [[SREM:%.*]] = srem <2 x i32> [[A0:%.*]], <i32 2, i32 2>
-; CHECK-NEXT:    [[ASHR:%.*]] = ashr <2 x i32> [[SREM]], [[A1:%.*]]
-; CHECK-NEXT:    [[MASK:%.*]] = and <2 x i32> [[ASHR]], <i32 2, i32 2>
+; CHECK-NEXT:    [[MASK:%.*]] = and <2 x i32> [[SREM]], <i32 2, i32 2>
 ; CHECK-NEXT:    ret <2 x i32> [[MASK]]
 ;
   %srem = srem <2 x i32> %a0, <i32 2, i32 2>
   %ashr = ashr <2 x i32> %srem, %a1
   %mask = and <2 x i32> %ashr, <i32 2, i32 2>
   ret <2 x i32> %mask
 }