[InstCombineCasts] Cast elimination for sext->lshr->trunc patterns

[kuhar]

Summary

There is a trunc(lshr (zext A), Cst) optimization in InstCombineCasts that removes cast by performing the lshr on smaller types. However, currently there is no trunc(lshr (sext A), Cst) variant.

This patch add such optimization by transforming trunc(lshr (sext A), Cst) to ashr A, Cst.

[aadg]

Hi Jakub,

Thanks for looking into this. This looks good to me with some fixes addressed (see below).

lib/Transforms/InstCombine/InstCombineCasts.cpp

504 ↗

(On Diff #33669)

Typo: s/comming/coming/

513 ↗

(On Diff #33669)

Typo: s/destanation/destination/

516–517 ↗

(On Diff #33669)

Isn't this comment redundant with the previous one ?

test/Transforms/InstCombine/cast.ll

1072–1073 ↗

(On Diff #33669)

Here you should pattern match the result, and reuse it when checking the return. This will make the test more robust if anything changes in the naming of the values.

1084–1085 ↗

(On Diff #33669)

Same. Use pattern matching.

1096–1099 ↗

(On Diff #33669)

Same.

[kuhar]

Typos fixed, regexes used in test-cases.

[kuhar]

Fix for PR24763

llvm/trunk/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (Src->hasOneUse() &&
     // Since we're doing an lshr and a zero extend, and know that the shift
     // amount is smaller than ASize, it is always safe to do the shift in A's
     // type, then zero extend or truncate to the result.
     Value *Shift = Builder->CreateLShr(A, Cst->getZExtValue());
     Shift->takeName(Src);
     return CastInst::CreateIntegerCast(Shift, CI.getType(), false);
   }
 
+  // Transform trunc(lshr (sext A), Cst) to ashr A, Cst to eliminate type
+  // conversion.
+  // It works because bits coming from sign extension have the same value as
+  // sign bit of the original value; performing ashr instead of lshr
+  // generates bits of the same value as the sign bit.
+  if (Src->hasOneUse() &&
+      match(Src, m_LShr(m_SExt(m_Value(A)), m_ConstantInt(Cst))) &&
+      cast<Instruction>(Src)->getOperand(0)->hasOneUse()) {
+    const unsigned ASize = A->getType()->getPrimitiveSizeInBits();
+    // This optimization can be only performed when zero bits generated by
+    // the original lshr aren't pulled into the value after truncation, so we
+    // can only shift by values smaller then the size of destination type (in
+    // bits).
+    if (Cst->getValue().ult(ASize)) {
+      Value *Shift = Builder->CreateAShr(A, Cst->getZExtValue());
+      Shift->takeName(Src);
+      return CastInst::CreateIntegerCast(Shift, CI.getType(), false);
+    }
+  }
+
   // Transform "trunc (and X, cst)" -> "and (trunc X), cst" so long as the dest
   // type isn't non-native.
   if (Src->hasOneUse() && isa<IntegerType>(Src->getType()) &&
       ShouldChangeType(Src->getType(), CI.getType()) &&
       match(Src, m_And(m_Value(A), m_ConstantInt(Cst)))) {
     Value *NewTrunc = Builder->CreateTrunc(A, CI.getType(), A->getName()+".tr");
     return BinaryOperator::CreateAnd(NewTrunc,
                                      ConstantExpr::getTrunc(Cst, CI.getType()));

llvm/trunk/test/Transforms/InstCombine/cast.ll

define i8 @test85(i32 %a) {
   ret i8 %trunc
 
 ; CHECK-LABEL: @test85(
 ; CHECK: [[ADD:%.*]] = add i32 %a, 2130706432
 ; CHECK: [[SHR:%.*]] = lshr exact i32 [[ADD]], 23
 ; CHECK: [[CST:%.*]] = trunc i32 [[SHR]] to i8
 }
 
+define i16 @test86(i16 %v) {
+  %a = sext i16 %v to i32
+  %s = ashr i32 %a, 4
+  %t = trunc i32 %s to i16
+  ret i16 %t
+
+; CHECK-LABEL: @test86(
+; CHECK:  [[ASHR:%.*]] = ashr i16 %v, 4
+; CHECK-NEXT: ret i16 [[ASHR]]
+}
+
+define i16 @test87(i16 %v) {
+  %c = sext i16 %v to i32
+  %m = mul nsw i32 %c, 16
+  %a = ashr i32 %m, 16
+  %t = trunc i32 %a to i16
+  ret i16 %t
+
+; CHECK-LABEL: @test87(
+; CHECK:  [[ASHR:%.*]] = ashr i16 %v, 12
+; CHECK-NEXT: ret i16 [[ASHR]]
+}
+
+define i16 @test88(i16 %v) {
+  %a = sext i16 %v to i32
+  %s = ashr i32 %a, 18
+  %t = trunc i32 %s to i16
+  ret i16 %t
+
+; Do not optimize to ashr i16 (shift by 18)
+; CHECK-LABEL: @test88(
+; CHECK: [[SEXT:%.*]] = sext i16 %v to i32
+; CHECK-NEXT: [[ASHR:%.*]] = ashr i32 [[SEXT]], 18
+; CHECK-NEXT: [[TRUNC:%.*]] = trunc i32 [[ASHR]] to i16
+; CHECK-NEXT: ret i16 [[TRUNC]]
+}
+
 ; Overflow on a float to int or int to float conversion is undefined (PR21130).
 
 define i8 @overflow_fptosi() {
   %i = fptosi double 1.56e+02 to i8
   ret i8 %i
 ; CHECK-LABEL: @overflow_fptosi(
 ; CHECK-NEXT: ret i8 undef 
 }