[InstCombine] visitTrunc - trunc (*shr (trunc A), C) --> trunc(*shr A, C)

[RKSimon]

Summary

Attempt to fold trunc (*shr (trunc A), C) --> trunc(*shr A, C) iff the shift amount if small enough that all zero/sign bits created by the shift are removed by the last trunc.

Helps fix the regressions encountered in D88316.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

862

Do you want to deal with non-splats from the getgo?
See m_SpecificInt_ICMP()

867

I think you can preserve exactness of the shift.

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

862

OMG - I never knew that existed - beautiful!!!!

867

Yes - alive2 says so https://alive2.llvm.org/ce/z/FLH5SX

[RKSimon]

Use m_SpecificInt_ICMP to support non-uniform constants

[lebedev.ri]

Please can you post the alive proof you're basing this on?
I don't understand how you arrived at that precondition,
i think it should be much simpler: https://rise4fun.com/Alive/zpJn

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

859

I do not understand why this isn't just unsigned MaxShiftAmt = AWidth - DestWidth;

[RKSimon]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

859

I tried that and hit issues with cases like this: https://alive2.llvm.org/ce/z/qV29xP truncating i64 -> i9 -> i8 - so I think there's a typo and it should be std::min<unsigned>(DestWidth, SrcWidth - DestWidth)

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

859

It is correct that https://alive2.llvm.org/ce/z/qV29x9 is a micompile, it should not be getting folded:
https://rise4fun.com/Alive/RTe7

However to me it still looks like that the precondition i proposed is correct:
https://rise4fun.com/Alive/saBR

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

859

Err, i guess i'm just not reading what i'm writing.
Of course unsigned MaxShiftAmt = AWidth - DestWidth; won't work, because i'm prooving that C1 u<= width(%b)-width(%R).
So it should be unsigned MaxShiftAmt = C0->getType()->getScalarSizeInBits() - DestWidth;
:/

[RKSimon]

Fix MaxShiftAmt value

[lebedev.ri]

LGTM, but please see https://lists.llvm.org/pipermail/llvm-commits/Week-of-Mon-20200928/835878.html

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

859

Please do mark done comments as such

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (C->getZExtValue() <= MaxShiftAmt) {
         unsigned ShAmt = std::min((unsigned)C->getZExtValue(), AWidth - 1);
         Value *Shift = Builder.CreateAShr(A, ShAmt);
         return CastInst::CreateIntegerCast(Shift, DestTy, true);
       }
     }
     // TODO: Mask high bits with 'and'.
   }
 
+  // trunc (*shr (trunc A), C) --> trunc(*shr A, C)
+  Constant *C0;
+  if (match(Src, m_OneUse(m_Shr(m_Trunc(m_Value(A)), m_Constant(C0))))) {
+    Type *AType = A->getType();
+    unsigned AWidth = AType->getScalarSizeInBits();
+    unsigned MaxShiftAmt = std::min<unsigned>(DestWidth, AWidth - DestWidth);
+
+    // If the shift is small enough, all zero/sign bits created by the shift are
+    // removed by the trunc.
+    if (match(C0, m_SpecificInt_ICMP(ICmpInst::ICMP_ULE,
+                                     APInt(SrcWidth, MaxShiftAmt)))) {
+      auto *OldShift = cast<Instruction>(Src);
+      auto *ShAmt = ConstantExpr::getIntegerCast(C0, AType, true);
+      bool IsExact = OldShift->isExact();
+      Value *Shift =
+          OldShift->getOpcode() == Instruction::AShr
+              ? Builder.CreateAShr(A, ShAmt, OldShift->getName(), IsExact)
+              : Builder.CreateLShr(A, ShAmt, OldShift->getName(), IsExact);
+      return CastInst::CreateTruncOrBitCast(Shift, DestTy);
+    }
+  }
+
   if (Instruction *I = narrowBinOp(Trunc))
     return I;
 
   if (Instruction *I = shrinkSplatShuffle(Trunc, Builder))
     return I;
 
   if (Instruction *I = shrinkInsertElt(Trunc, Builder))
     return I;

llvm/test/Transforms/InstCombine/trunc-shift-trunc.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64"
 
 ; Perform the shift first and merge the truncs as long as all zero/sign
 ; bits created by the shift are removed by the trunc.
 
 declare void @use(i32)
 
 define i8 @trunc_lshr_trunc(i64 %a) {
 ; CHECK-LABEL: @trunc_lshr_trunc(
-; CHECK-NEXT:    [[B:%.*]] = trunc i64 [[A:%.*]] to i32
-; CHECK-NEXT:    [[C:%.*]] = lshr i32 [[B]], 8
-; CHECK-NEXT:    [[D:%.*]] = trunc i32 [[C]] to i8
+; CHECK-NEXT:    [[C1:%.*]] = lshr i64 [[A:%.*]], 8
+; CHECK-NEXT:    [[D:%.*]] = trunc i64 [[C1]] to i8
 ; CHECK-NEXT:    ret i8 [[D]]
 ;
   %b = trunc i64 %a to i32
   %c = lshr i32 %b, 8
   %d = trunc i32 %c to i8
   ret i8 %d
 }
 
 define <2 x i8> @trunc_lshr_trunc_uniform(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_lshr_trunc_uniform(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = lshr <2 x i32> [[B]], <i32 8, i32 8>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; CHECK-NEXT:    [[C1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 8, i64 8>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 8, i32 8>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_lshr_trunc_nonuniform(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_lshr_trunc_nonuniform(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = lshr <2 x i32> [[B]], <i32 8, i32 2>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; CHECK-NEXT:    [[C1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 8, i64 2>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 8, i32 2>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_lshr_trunc_uniform_undef(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_lshr_trunc_uniform_undef(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = lshr <2 x i32> [[B]], <i32 8, i32 undef>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; CHECK-NEXT:    [[C1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 8, i64 0>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 8, i32 undef>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = lshr <2 x i32> %b, <i32 8, i32 11>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define i8 @trunc_ashr_trunc(i64 %a) {
 ; CHECK-LABEL: @trunc_ashr_trunc(
-; CHECK-NEXT:    [[B:%.*]] = trunc i64 [[A:%.*]] to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr i32 [[B]], 8
-; CHECK-NEXT:    [[D:%.*]] = trunc i32 [[TMP1]] to i8
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i64 [[A:%.*]], 8
+; CHECK-NEXT:    [[D:%.*]] = trunc i64 [[TMP1]] to i8
 ; CHECK-NEXT:    ret i8 [[D]]
 ;
   %b = trunc i64 %a to i32
   %c = ashr i32 %b, 8
   %d = trunc i32 %c to i8
   ret i8 %d
 }
 
+define i8 @trunc_ashr_trunc_exact(i64 %a) {
+; CHECK-LABEL: @trunc_ashr_trunc_exact(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr exact i64 [[A:%.*]], 8
+; CHECK-NEXT:    [[D:%.*]] = trunc i64 [[TMP1]] to i8
+; CHECK-NEXT:    ret i8 [[D]]
+;
+  %b = trunc i64 %a to i32
+  %c = ashr exact i32 %b, 8
+  %d = trunc i32 %c to i8
+  ret i8 %d
+}
+
 define <2 x i8> @trunc_ashr_trunc_uniform(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_ashr_trunc_uniform(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i32> [[B]], <i32 8, i32 8>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[TMP1]] to <2 x i8>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i64> [[A:%.*]], <i64 8, i64 8>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[TMP1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = ashr <2 x i32> %b, <i32 8, i32 8>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_ashr_trunc_nonuniform(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_ashr_trunc_nonuniform(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i32> [[B]], <i32 0, i32 2>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; CHECK-NEXT:    [[C1:%.*]] = ashr <2 x i64> [[A:%.*]], <i64 0, i64 2>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = ashr <2 x i32> %b, <i32 0, i32 2>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }
 
 define <2 x i8> @trunc_ashr_trunc_uniform_undef(<2 x i64> %a) {
 ; CHECK-LABEL: @trunc_ashr_trunc_uniform_undef(
-; CHECK-NEXT:    [[B:%.*]] = trunc <2 x i64> [[A:%.*]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = ashr <2 x i32> [[B]], <i32 8, i32 undef>
-; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i32> [[C]] to <2 x i8>
+; CHECK-NEXT:    [[C1:%.*]] = ashr <2 x i64> [[A:%.*]], <i64 8, i64 0>
+; CHECK-NEXT:    [[D:%.*]] = trunc <2 x i64> [[C1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[D]]
 ;
   %b = trunc <2 x i64> %a to <2 x i32>
   %c = ashr <2 x i32> %b, <i32 8, i32 undef>
   %d = trunc <2 x i32> %c to <2 x i8>
   ret <2 x i8> %d
 }