[InstCombine] Consider more cases where SimplifyDemandedUseBits do not converting AShr to LShr

[aaboud]

Summary

There are cases where AShr have better chance to be optimized than LShr, especially when the demanded bits are not known to be Zero, and even known to be similar to the sign bit.

Prevent converting AShr to LShr in SimplifyDemandedUseBits when sign bit is not known to be zero and some of the demanded bits are known to be equal to the sign bit.

[craig.topper]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

524

Should we use Depth+1 instead of 0 here?

[aaboud]

Updated patch to address Craig's comment.

[aaboud]

Any more comments?

[craig.topper]

LGTM

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (match(I->getOperand(1), m_APInt(SA))) {
       // 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), 0, CxtI);
+
       assert(!Known.hasConflict() && "Bits known to be one AND zero?");
-      // Compute the new bits that are at the top now.
-      APInt HighBits(APInt::getHighBitsSet(BitWidth, ShiftAmt));
+      // 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
       // are demanded, turn this into an unsigned shift right.
       assert(BitWidth > ShiftAmt && "Shift amount not saturated?");
       if (Known.Zero[BitWidth-ShiftAmt-1] ||
           !DemandedMask.intersects(HighBits)) {

test/Transforms/InstCombine/trunc.ll

 ; CHECK-NEXT:    ret i32 [[D]]
 ;
   %B = zext i64 %A to i128
   %C = lshr i128 %B, 32
   %D = trunc i128 %C to i32
   ret i32 %D
 }
 
+define i16 @test6_ashr_mul(i8 %X, i8 %Y) {
+; CHECK-LABEL: @test6_ashr_mul(
+; CHECK-NEXT:    [[A:%.*]] = sext i8 %X to i16
+; CHECK-NEXT:    [[B:%.*]] = sext i8 %Y to i16
+; CHECK-NEXT:    [[C:%.*]] = mul nsw i16 [[A]], [[B]]
+; CHECK-NEXT:    [[D:%.*]] = ashr i16 %C, 15
+; CHECK-NEXT:    ret i16 %D
+  %A = sext i8 %X to i32
+  %B = sext i8 %Y to i32
+  %C = mul i32 %A, %B
+  %D = ashr i32 %C, 15
+  %E = trunc i32 %D to i16
+  ret i16 %E
+}
+
 define i16 @ashr_mul(i8 %X, i8 %Y) {
 ; CHECK-LABEL: @ashr_mul(
 ; CHECK-NEXT:    [[A:%.*]] = sext i8 %X to i16
 ; CHECK-NEXT:    [[B:%.*]] = sext i8 %Y to i16
 ; CHECK-NEXT:    [[C:%.*]] = mul nsw i16 [[A]], [[B]]
 ; CHECK-NEXT:    [[D:%.*]] = ashr i16 [[C]], 8
 ; CHECK-NEXT:    ret i16 [[D]]
   %A = sext i8 %X to i20