[InstCombine] eliminate sext and/or trunc if value has enough signbits

[spatel]

Summary

If we have enough signbits in a source value, we can skip an intermediate cast for a trunc+sext pair:
https://alive2.llvm.org/ce/z/A_mQt-

This is the original problem shown in:
https://llvm.org/PR49543

There's a test that shows we transformed what used to be a pair of shifts, so that suggests we could add another ComputeNumSignBits fold starting from a shift.

There does not appear to be any change in compile-time from the extra analysis:
https://llvm-compile-time-tracker.com/compare.php?from=3d2c9069dcafd0cbb641841aa3dd6e851fb7d760&to=b9513cdf2419704c7bb0c3a02a9ca06aae13d902&stat=instructions

[lebedev.ri]

As per 5a654bfeab588715fa6d266a7789484850605a18, did you add the test where we'd form trunc?

[spatel]

In D103617#2796286, @lebedev.ri wrote:

As per 5a654bfeab588715fa6d266a7789484850605a18, did you add the test where we'd form trunc?

I didn't know you had tried this one already. :)
Yes - see tests "wide_source_matching_signbits" and "wide_source_not_matching_signbits".

[lebedev.ri]

LG

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (ComputeNumSignBits(Res, 0, &CI) > DestBitSize - SrcBitSize)
       return replaceInstUsesWith(CI, Res);
 
     // We need to emit a shl + ashr to do the sign extend.
     Value *ShAmt = ConstantInt::get(DestTy, DestBitSize-SrcBitSize);
     return BinaryOperator::CreateAShr(Builder.CreateShl(Res, ShAmt, "sext"),
                                       ShAmt);
   }
 
-  // If the input is a trunc from the destination type, then turn sext(trunc(x))
-  // into shifts.
   Value *X;
-  if (match(Src, m_OneUse(m_Trunc(m_Value(X)))) && X->getType() == DestTy) {
+  if (match(Src, m_Trunc(m_Value(X)))) {
+    // If the input has more sign bits than bits truncated, then convert
+    // directly to final type.
+    unsigned XBitSize = X->getType()->getScalarSizeInBits();
+    if (ComputeNumSignBits(X, 0, &CI) > XBitSize - SrcBitSize)
+      return CastInst::CreateIntegerCast(X, DestTy, /* isSigned */ true);
+
+    // If input is a trunc from the destination type, then convert into shifts.
+    if (Src->hasOneUse() && X->getType() == DestTy) {
       // sext (trunc X) --> ashr (shl X, C), C
       Constant *ShAmt = ConstantInt::get(DestTy, DestBitSize - SrcBitSize);
       return BinaryOperator::CreateAShr(Builder.CreateShl(X, ShAmt), ShAmt);
     }
+  }
 
   if (ICmpInst *ICI = dyn_cast<ICmpInst>(Src))
     return transformSExtICmp(ICI, CI);
 
   // If the input is a shl/ashr pair of a same constant, then this is a sign
   // extension from a smaller value.  If we could trust arbitrary bitwidth
   // integers, we could turn this into a truncate to the smaller bit and then
   // use a sext for the whole extension.  Since we don't, look deeper and check

llvm/test/Transforms/InstCombine/sext-of-trunc-nsw.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; Iff we know that trunc only chops off the sign bits,
 ; and not all of them, then we can bypass said trunc
 ; for the purpose of sign extension.
 
 declare void @use8(i8)
 declare void @use4(i4)
 declare void @usevec(<2 x i8>)
 
 define i16 @t0(i8 %x) {
 ; CHECK-LABEL: @t0(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 5
 ; CHECK-NEXT:    call void @use8(i8 [[A]])
-; CHECK-NEXT:    [[B:%.*]] = trunc i8 [[A]] to i4
-; CHECK-NEXT:    [[C:%.*]] = sext i4 [[B]] to i16
+; CHECK-NEXT:    [[C:%.*]] = sext i8 [[A]] to i16
 ; CHECK-NEXT:    ret i16 [[C]]
 ;
   %a = ashr i8 %x, 5
   call void @use8(i8 %a)
   %b = trunc i8 %a to i4
   %c = sext i4 %b to i16
   ret i16 %c
 }
 
 define i16 @t1(i8 %x) {
 ; CHECK-LABEL: @t1(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 4
 ; CHECK-NEXT:    call void @use8(i8 [[A]])
-; CHECK-NEXT:    [[B:%.*]] = trunc i8 [[A]] to i4
-; CHECK-NEXT:    [[C:%.*]] = sext i4 [[B]] to i16
+; CHECK-NEXT:    [[C:%.*]] = sext i8 [[A]] to i16
 ; CHECK-NEXT:    ret i16 [[C]]
 ;
   %a = ashr i8 %x, 4
   call void @use8(i8 %a)
   %b = trunc i8 %a to i4
   %c = sext i4 %b to i16
   ret i16 %c
 }
;
   %c = sext i4 %b to i16
   ret i16 %c
 }
 
 define <2 x i16> @t3_vec(<2 x i8> %x) {
 ; CHECK-LABEL: @t3_vec(
 ; CHECK-NEXT:    [[A:%.*]] = ashr <2 x i8> [[X:%.*]], <i8 4, i8 4>
 ; CHECK-NEXT:    call void @usevec(<2 x i8> [[A]])
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i8> [[A]] to <2 x i4>
-; CHECK-NEXT:    [[C:%.*]] = sext <2 x i4> [[B]] to <2 x i16>
+; CHECK-NEXT:    [[C:%.*]] = sext <2 x i8> [[A]] to <2 x i16>
 ; CHECK-NEXT:    ret <2 x i16> [[C]]
 ;
   %a = ashr <2 x i8> %x, <i8 4, i8 4>
   call void @usevec(<2 x i8> %a)
   %b = trunc <2 x i8> %a to <2 x i4>
   %c = sext <2 x i4> %b to <2 x i16>
   ret <2 x i16> %c
 }

 }
 
 define i16 @t5_extrause(i8 %x) {
 ; CHECK-LABEL: @t5_extrause(
 ; CHECK-NEXT:    [[A:%.*]] = ashr i8 [[X:%.*]], 5
 ; CHECK-NEXT:    call void @use8(i8 [[A]])
 ; CHECK-NEXT:    [[B:%.*]] = trunc i8 [[A]] to i4
 ; CHECK-NEXT:    call void @use4(i4 [[B]])
-; CHECK-NEXT:    [[C:%.*]] = sext i4 [[B]] to i16
+; CHECK-NEXT:    [[C:%.*]] = sext i8 [[A]] to i16
 ; CHECK-NEXT:    ret i16 [[C]]
 ;
   %a = ashr i8 %x, 5
   call void @use8(i8 %a)
   %b = trunc i8 %a to i4
   call void @use4(i4 %b)
   %c = sext i4 %b to i16
   ret i16 %c
 }
 
 define i64 @narrow_source_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @narrow_source_matching_signbits(
 ; CHECK-NEXT:    [[M:%.*]] = and i32 [[X:%.*]], 7
 ; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M]]
-; CHECK-NEXT:    [[B:%.*]] = trunc i32 [[A]] to i8
-; CHECK-NEXT:    [[C:%.*]] = sext i8 [[B]] to i64
+; CHECK-NEXT:    [[C:%.*]] = sext i32 [[A]] to i64
 ; CHECK-NEXT:    ret i64 [[C]]
 ;
   %m = and i32 %x, 7
   %a = shl nsw i32 -1, %m
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i64
   ret i64 %c
 }
 
+; negative test - not enough sign-bits
+
 define i64 @narrow_source_not_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @narrow_source_not_matching_signbits(
 ; CHECK-NEXT:    [[M:%.*]] = and i32 [[X:%.*]], 8
 ; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M]]
 ; CHECK-NEXT:    [[B:%.*]] = trunc i32 [[A]] to i8
 ; CHECK-NEXT:    [[C:%.*]] = sext i8 [[B]] to i64
 ; CHECK-NEXT:    ret i64 [[C]]
 ;
   %m = and i32 %x, 8
   %a = shl nsw i32 -1, %m
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i64
   ret i64 %c
 }
 
 define i24 @wide_source_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @wide_source_matching_signbits(
 ; CHECK-NEXT:    [[M:%.*]] = and i32 [[X:%.*]], 7
 ; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M]]
-; CHECK-NEXT:    [[B:%.*]] = trunc i32 [[A]] to i8
-; CHECK-NEXT:    [[C:%.*]] = sext i8 [[B]] to i24
+; CHECK-NEXT:    [[C:%.*]] = trunc i32 [[A]] to i24
 ; CHECK-NEXT:    ret i24 [[C]]
 ;
   %m = and i32 %x, 7
   %a = shl nsw i32 -1, %m
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i24
   ret i24 %c
 }
 
+; negative test - not enough sign-bits
+
 define i24 @wide_source_not_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @wide_source_not_matching_signbits(
 ; CHECK-NEXT:    [[M2:%.*]] = and i32 [[X:%.*]], 8
 ; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M2]]
 ; CHECK-NEXT:    [[B:%.*]] = trunc i32 [[A]] to i8
 ; CHECK-NEXT:    [[C:%.*]] = sext i8 [[B]] to i24
 ; CHECK-NEXT:    ret i24 [[C]]
 ;
   %m2 = and i32 %x, 8
   %a = shl nsw i32 -1, %m2
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i24
   ret i24 %c
 }
 
 define i32 @same_source_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @same_source_matching_signbits(
 ; CHECK-NEXT:    [[M:%.*]] = and i32 [[X:%.*]], 7
-; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 -16777216, [[M]]
-; CHECK-NEXT:    [[C:%.*]] = ashr exact i32 [[TMP1]], 24
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M]]
+; CHECK-NEXT:    ret i32 [[A]]
 ;
   %m = and i32 %x, 7
   %a = shl nsw i32 -1, %m
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i32
   ret i32 %c
 }
 
+; negative test - not enough sign-bits
+
 define i32 @same_source_not_matching_signbits(i32 %x) {
 ; CHECK-LABEL: @same_source_not_matching_signbits(
 ; CHECK-NEXT:    [[M2:%.*]] = and i32 [[X:%.*]], 8
 ; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 -16777216, [[M2]]
 ; CHECK-NEXT:    [[C:%.*]] = ashr exact i32 [[TMP1]], 24
 ; CHECK-NEXT:    ret i32 [[C]]
 ;
   %m2 = and i32 %x, 8
   %a = shl nsw i32 -1, %m2
   %b = trunc i32 %a to i8
   %c = sext i8 %b to i32
   ret i32 %c
 }
 
 define i32 @same_source_matching_signbits_extra_use(i32 %x) {
 ; CHECK-LABEL: @same_source_matching_signbits_extra_use(
 ; CHECK-NEXT:    [[M:%.*]] = and i32 [[X:%.*]], 7
 ; CHECK-NEXT:    [[A:%.*]] = shl nsw i32 -1, [[M]]
 ; CHECK-NEXT:    [[B:%.*]] = trunc i32 [[A]] to i8
 ; CHECK-NEXT:    call void @use8(i8 [[B]])
-; CHECK-NEXT:    [[C:%.*]] = sext i8 [[B]] to i32
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    ret i32 [[A]]
 ;
   %m = and i32 %x, 7
   %a = shl nsw i32 -1, %m
   %b = trunc i32 %a to i8
   call void @use8(i8 %b)
   %c = sext i8 %b to i32
   ret i32 %c
 }