[InstCombine] Added support for: trunc(ashr(mul(sext(...), sext(...))) -> ashr(mul(...))

[aaboud]

Summary

Added support for the following case:

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

}

So far, InstCombine handled only LShr instruction and tried to replace AShr instruction with the LShr when applicable.
However, there are cases where LShr has no advantage over AShr, especially when the higher bits are not known to be zero, but are known to be equal to the sign bit.

Note: There is more to do to catch all the cases where AShr can be optimized, but we end up having LShr instead, e.g., we could replace LShr with AShr when applicable. However, the current implementation of InstCombine will probably lead into infinite loop if we try implement both transform direction (i.e., LShr->AShr and AShr->LShr).

[spatel]

lib/Transforms/InstCombine/InstCombineShifts.cpp

562–565 ↗

(On Diff #110960)

Use 'm_Shr' to reduce the code. This change is independent of anything else, so I think it should go in first.
http://rise4fun.com/Alive/lBi
Here's a test you can use for that patch:

define i8 @shishi(i8 %x) {
  %a = ashr i8 %x, 6
  %b = shl i8 %a, 6
  %extra_use_of_a = mul i8 %a, 5
  %r = sdiv i8 %extra_use_of_a, %b
  ret i8 %r
}

[aaboud]

Split the patch by moving: (X >>s C) << C --> X & (-1 << C) to D36743.

[spatel]

I think all of the non-multiply tests are simplified after rL310942. Can you double-check that? If so, you can remove those tests from this patch. Feel free to add them as an independent NFC commit if they cover some case that isn't covered already.

It may also be possible to simplify the code in this patch now, but I didn't check that.

[aaboud]

Updated patch to top of trunk and removed NFC tests (will commit them separately).

[spatel]

Can we have the ashr fix as a separate patch before the mul patch? If yes, make tests for just that diff and split this up?

lib/Transforms/InstCombine/InstCombineCasts.cpp

378

Over in D36763, we're going to (or maybe already have) made the other shift cases splat-vector-friendly. Can you do the same here and add a vector test?

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

468

Why change this in this patch? Seem like that might be better as just 'unsigned', but either way, I think it's an independent diff.

[aaboud]

In D36679#842794, @spatel wrote:

Can we have the ashr fix as a separate patch before the mul patch? If yes, make tests for just that diff and split this up?

I uploaded a separate patch that handles ashr including vector type handling (D36784).

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

468

I change it to uint32_t to be consistent with line 507 (the case for AShr).
Sure, we can commit this separately.
Do we need a review for this?
Do you want to take care of it? or I should just commit it now?

[spatel]

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

468

Ah, I see the LangRef says we're allowed to go up to 24 bits of width:
http://llvm.org/docs/LangRef.html#integer-type

So I guess it's ok to specify 32-bit here...or just change them all to 'unsigned'? Either way seems fine to me. Feel free to commit separately.

[aaboud]

After committing the other parts, this patch was reduced to handle only the multiply in ComputeNumSignBitsImpl.

[spatel]

Thanks for breaking this up - LGTM.
See inline for a couple of nits.

lib/Analysis/ValueTracking.cpp

2200

can at -> can be at

2201–2203

It's poor to have variables named 'Tmp'. These can be locals now, so let's give these meaningful names - SignBitsOp0/SignBitsOp1 or something like that.

lib/Analysis/ValueTracking.cpp

if (const auto *CLHS = dyn_cast<Constant>(U->getOperand(0)))
       }
 
     // Sub can have at most one carry bit.  Thus we know that the output
     // is, at worst, one more bit than the inputs.
     Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
     if (Tmp == 1) return 1;  // Early out.
     return std::min(Tmp, Tmp2)-1;
 
+  case Instruction::Mul: {
+    // The output of the Mul can at most twice the valid bits in the inputs.
+    Tmp = ComputeNumSignBits(U->getOperand(0), Depth + 1, Q);
+    if (Tmp == 1) return 1;  // Early out.
+    Tmp2 = ComputeNumSignBits(U->getOperand(1), Depth + 1, Q);
+    if (Tmp2 == 1) return 1;
+    unsigned OutValidBits = (TyBits - Tmp + 1) + (TyBits - Tmp2 + 1);
+    return OutValidBits > TyBits ? 1 : TyBits - OutValidBits + 1;
+  }
+
   case Instruction::PHI: {
     const PHINode *PN = cast<PHINode>(U);
     unsigned NumIncomingValues = PN->getNumIncomingValues();
     // Don't analyze large in-degree PHIs.
     if (NumIncomingValues > 4) break;
     // Unreachable blocks may have zero-operand PHI nodes.
     if (NumIncomingValues == 0) break;
 

lib/Transforms/InstCombine/InstCombineCasts.cpp

if (ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1))) {
       uint32_t BitWidth = Ty->getScalarSizeInBits();
       if (IC.MaskedValueIsZero(I->getOperand(0),
             APInt::getHighBitsSet(OrigBitWidth, OrigBitWidth-BitWidth), 0, CxtI) &&
           CI->getLimitedValue(BitWidth) < BitWidth) {
         return canEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
       }
     }
     break;
+  case Instruction::AShr:
+    // If this is a truncate of an arithmetic shr, we can truncate it to a
+    // smaller ashr iff we know that all the bits from the sign bit of the
+    // original type and the sign bit of the truncate type are similar.
+    // TODO: It is enough to check that the bits we would be shifting in are
+    //       similar to sign bit of the truncate type.
+    if (ConstantInt *CI = dyn_cast<ConstantInt>(I->getOperand(1))) {
+      uint32_t OrigBitWidth = OrigTy->getScalarSizeInBits();
+      uint32_t BitWidth = Ty->getScalarSizeInBits();
+      if (CI->getLimitedValue(BitWidth) < BitWidth &&
+          OrigBitWidth - BitWidth <
+              IC.ComputeNumSignBits(I->getOperand(0), 0, CxtI))
+        return canEvaluateTruncated(I->getOperand(0), Ty, IC, CxtI);
+    }
+    break;
   case Instruction::Trunc:
     // trunc(trunc(x)) -> trunc(x)
     return true;
   case Instruction::ZExt:
   case Instruction::SExt:
     // trunc(ext(x)) -> ext(x) if the source type is smaller than the new dest
     // trunc(ext(x)) -> trunc(x) if the source type is larger than the new dest
     return true;

lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

if (match(I->getOperand(1), m_APInt(SA))) {
       if (ShiftAmt)
         Known.Zero.setLowBits(ShiftAmt);
     }
     break;
   }
   case Instruction::LShr: {
     const APInt *SA;
     if (match(I->getOperand(1), m_APInt(SA))) {
-      uint64_t ShiftAmt = SA->getLimitedValue(BitWidth-1);
+      uint32_t ShiftAmt = SA->getLimitedValue(BitWidth-1);
 
       // 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);
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));
+      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 [[C]]
 ;
   %B = zext i32 %A to i128
   %C = lshr i128 %B, 16
   %D = trunc i128 %C to i32
   ret i32 %D
 }
 
+define i32 @test5_ashr(i32 %A) {
+; CHECK-LABEL: @test5_ashr(
+; CHECK-NEXT:    [[C:%.*]] = ashr i32 %A, 16
+; CHECK-NEXT:    ret i32 [[C]]
+;
+  %B = sext i32 %A to i128
+  %C = ashr i128 %B, 16
+  %D = trunc i128 %C to i32
+  ret i32 %D
+}
+
 define i32 @test6(i64 %A) {
 ; CHECK-LABEL: @test6(
 ; CHECK-NEXT:    [[C:%.*]] = lshr i64 %A, 32
 ; CHECK-NEXT:    [[D:%.*]] = trunc i64 [[C]] to i32
 ; 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 i32 @test6_lshr(i64 %A) {
+; CHECK-LABEL: @test6_lshr(
+; CHECK-NEXT:    [[C:%.*]] = lshr i64 %A, 32
+; CHECK-NEXT:    [[D:%.*]] = trunc i64 [[C]] to i32
+; CHECK-NEXT:    ret i32 [[D]]
+;
+  %B = sext i64 %A to i128
+  %C = ashr i128 %B, 32
+  %D = trunc i128 %C to i32
+  ret i32 %D
+}
+
+define i64 @test6_ashr(i64 %A) {
+; CHECK-LABEL: @test6_ashr(
+; CHECK-NEXT:    [[C:%.*]] = ashr i64 %A, 33
+; CHECK-NEXT:    ret i64 [[C]]
+;
+  %B = sext i64 %A to i128
+  %C = ashr i128 %B, 33
+  %D = trunc i128 %C to i64
+  ret i64 %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 i92 @test7(i64 %A) {
 ; CHECK-LABEL: @test7(
 ; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 %A, 32
 ; CHECK-NEXT:    [[D:%.*]] = zext i64 [[TMP1]] to i92
 ; CHECK-NEXT:    ret i92 [[D]]
 ;
   %B = zext i64 %A to i128
   %C = lshr i128 %B, 32
   %D = trunc i128 %C to i92
   ret i92 %D
 }
 
+define i92 @test7_ashr(i64 %A) {
+; CHECK-LABEL: @test7_ashr(
+; CHECK-NEXT:    [[B:%.*]] = sext i64 %A to i92
+; CHECK-NEXT:    [[C:%.*]] = ashr i92 [[B]], 32
+; CHECK-NEXT:    ret i92 [[C]]
+;
+  %B = sext i64 %A to i128
+  %C = ashr i128 %B, 32
+  %D = trunc i128 %C to i92
+  ret i92 %D
+}
+
 define i64 @test8(i32 %A, i32 %B) {
 ; CHECK-LABEL: @test8(
 ; CHECK-NEXT:    [[TMP38:%.*]] = zext i32 %A to i64
 ; CHECK-NEXT:    [[TMP32:%.*]] = zext i32 %B to i64
 ; CHECK-NEXT:    [[TMP33:%.*]] = shl nuw i64 [[TMP32]], 32
 ; CHECK-NEXT:    [[INS35:%.*]] = or i64 [[TMP33]], [[TMP38]]
 ; CHECK-NEXT:    ret i64 [[INS35]]
 ;