[InstCombine] fold lshr (sext X), C1 --> zext (lshr X, C2)

[spatel]

Summary

This was discussed in D33338. We have larger pattern-matching ending in a truncate that we can reduce or remove by handling these smaller patterns first. Further motivation is that narrower shift ops are easier for value tracking and zext is better than sext.

http://rise4fun.com/Alive/rhh

Name: boolshift
%sext = sext i1 %x to i8
%r = lshr i8 %sext, 7

=>

%r = zext i1 %x to i8

Name: noboolshift
%sext = sext i3 %x to i8
%r = lshr i8 %sext, 7

=>

%sh = lshr i3 %x, 2
%r = zext i3 %sh to i8

[efriedma]

lib/Transforms/InstCombine/InstCombineShifts.cpp

693

This isn't calling shouldChangeType; is that intentional? (Not sure how much this matters.)

[spatel]

lib/Transforms/InstCombine/InstCombineShifts.cpp

693

No, that's an oversight. I'll add that and a test case.

[spatel]

Patch updated:

1. Use shouldChangeType() to guard against adding a shift of an illegal type.
2. Add/modify tests to show that difference.

[efriedma]

LGTM.

lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op0, m_Shl(m_Value(X), m_APInt(ShOp1)))) {
         return BinaryOperator::CreateAnd(NewShl, ConstantInt::get(Ty, Mask));
       }
       assert(ShlAmt == ShAmt);
       // (X << C) >>u C --> X & (-1 >>u C)
       APInt Mask(APInt::getLowBitsSet(BitWidth, BitWidth - ShAmt));
       return BinaryOperator::CreateAnd(X, ConstantInt::get(Ty, Mask));
     }
 
+    if (match(Op0, m_SExt(m_Value(X)))) {
+      // Are we moving the sign bit to the low bit and widening with high zeros?
+      unsigned SrcTyBitWidth = X->getType()->getScalarSizeInBits();
+      if (ShAmt == BitWidth - 1 &&
+          (!Ty->isIntegerTy() || shouldChangeType(Ty, X->getType()))) {
+        // lshr (sext i1 X to iN), N-1 --> zext X to iN
+        if (SrcTyBitWidth == 1)
+          return new ZExtInst(X, Ty);
+
+        // lshr (sext iM X to iN), N-1 --> zext (lshr X, M-1) to iN
+        if (Op0->hasOneUse()) {
+          Value *NewLShr = Builder->CreateLShr(X, SrcTyBitWidth - 1);
+          return new ZExtInst(NewLShr, Ty);
+        }
+      }
+
+      // TODO: Convert to ashr+zext if the shift equals the extension amount.
+    }
+
     if (match(Op0, m_LShr(m_Value(X), m_APInt(ShOp1)))) {
       unsigned AmtSum = ShAmt + ShOp1->getZExtValue();
       // Oversized shifts are simplified to zero in InstSimplify.
       if (AmtSum < BitWidth)
         // (X >>u C1) >>u C2 --> X >>u (C1 + C2)
         return BinaryOperator::CreateLShr(X, ConstantInt::get(Ty, AmtSum));
     }
 

test/Transforms/InstCombine/lshr.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine -S < %s | FileCheck %s
 
+target datalayout = "e-m:e-i64:64-n8:16:32:64"
+
 declare i32 @llvm.cttz.i32(i32, i1) nounwind readnone
 declare i32 @llvm.ctlz.i32(i32, i1) nounwind readnone
 declare i32 @llvm.ctpop.i32(i32) nounwind readnone
 declare <2 x i8> @llvm.cttz.v2i8(<2 x i8>, i1) nounwind readnone
 declare <2 x i8> @llvm.ctlz.v2i8(<2 x i8>, i1) nounwind readnone
 declare <2 x i8> @llvm.ctpop.v2i8(<2 x i8>) nounwind readnone
 
 define i32 @lshr_ctlz_zero_is_not_undef(i32 %x) {

 ; CHECK-NEXT:    ret <2 x i8> [[LSHR]]
 ;
   %shl = shl <2 x i8> %x, <i8 2, i8 2>
   %add = add <2 x i8> %shl, <i8 4, i8 4>
   %lshr = lshr <2 x i8> %add, <i8 2, i8 2>
   ret <2 x i8> %lshr
 }
 
-; FIXME: The bool bit got smeared across a wide val, but then we zero'd out those bits. This is just a zext.
-
 define i16 @bool_zext(i1 %x) {
 ; CHECK-LABEL: @bool_zext(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 %x to i16
-; CHECK-NEXT:    [[HIBIT:%.*]] = lshr i16 [[SEXT]], 15
+; CHECK-NEXT:    [[HIBIT:%.*]] = zext i1 %x to i16
 ; CHECK-NEXT:    ret i16 [[HIBIT]]
 ;
   %sext = sext i1 %x to i16
   %hibit = lshr i16 %sext, 15
   ret i16 %hibit
 }
 
 define <2 x i8> @bool_zext_splat(<2 x i1> %x) {
 ; CHECK-LABEL: @bool_zext_splat(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i1> %x to <2 x i8>
-; CHECK-NEXT:    [[HIBIT:%.*]] = lshr <2 x i8> [[SEXT]], <i8 7, i8 7>
+; CHECK-NEXT:    [[HIBIT:%.*]] = zext <2 x i1> %x to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[HIBIT]]
 ;
   %sext = sext <2 x i1> %x to <2 x i8>
   %hibit = lshr <2 x i8> %sext, <i8 7, i8 7>
   ret <2 x i8> %hibit
 }
 
 ; FIXME: The replicated sign bits are all that's left. This could be ashr+zext.

 ; CHECK-NEXT:    [[HIBIT:%.*]] = lshr <2 x i8> [[SEXT]], <i8 2, i8 2>
 ; CHECK-NEXT:    ret <2 x i8> [[HIBIT]]
 ;
   %sext = sext <2 x i6> %x to <2 x i8>
   %hibit = lshr <2 x i8> %sext, <i8 2, i8 2>
   ret <2 x i8> %hibit
 }
 
-; FIXME: All of the replicated sign bits are wiped out by the lshr. This could be lshr+zext.
-
-define i16 @fake_sext(i3 %x) {
+define i18 @fake_sext(i3 %x) {
 ; CHECK-LABEL: @fake_sext(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i3 %x to i16
-; CHECK-NEXT:    [[SH:%.*]] = lshr i16 [[SEXT]], 15
-; CHECK-NEXT:    ret i16 [[SH]]
-;
-  %sext = sext i3 %x to i16
-  %sh = lshr i16 %sext, 15
-  ret i16 %sh
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr i3 %x, 2
+; CHECK-NEXT:    [[SH:%.*]] = zext i3 [[TMP1]] to i18
+; CHECK-NEXT:    ret i18 [[SH]]
+;
+  %sext = sext i3 %x to i18
+  %sh = lshr i18 %sext, 17
+  ret i18 %sh
+}
+
+; Avoid the transform if it would change the shift from a legal to illegal type.
+
+define i32 @fake_sext_but_should_not_change_type(i3 %x) {
+; CHECK-LABEL: @fake_sext_but_should_not_change_type(
+; CHECK-NEXT:    [[SEXT:%.*]] = sext i3 %x to i32
+; CHECK-NEXT:    [[SH:%.*]] = lshr i32 [[SEXT]], 31
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %sext = sext i3 %x to i32
+  %sh = lshr i32 %sext, 31
+  ret i32 %sh
 }
 
 define <2 x i8> @fake_sext_splat(<2 x i3> %x) {
 ; CHECK-LABEL: @fake_sext_splat(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext <2 x i3> %x to <2 x i8>
-; CHECK-NEXT:    [[SH:%.*]] = lshr <2 x i8> [[SEXT]], <i8 7, i8 7>
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr <2 x i3> %x, <i3 2, i3 2>
+; CHECK-NEXT:    [[SH:%.*]] = zext <2 x i3> [[TMP1]] to <2 x i8>
 ; CHECK-NEXT:    ret <2 x i8> [[SH]]
 ;
   %sext = sext <2 x i3> %x to <2 x i8>
   %sh = lshr <2 x i8> %sext, <i8 7, i8 7>
   ret <2 x i8> %sh
 }