[InstCombine] Support arbitrary const shift amount for `lshr (sext i1 ...)`

[anton-afanasyev]

Summary

Add lshr (sext i1 X to iN), C --> select (X, -1 >> C, 0) case. This expands
C == N-1 case to arbitrary C.

Fixes PR52078

[spatel]

Can you add the instcombine tests to the file(s) where the existing transform's tests are ("lshr.ll")? It's easier to see the variations that we handle in one look that way.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1069–1070

We don't need the shouldChangeType restriction for the fold to select / zext, do we?
(I'm not sure that we need it for the fold that creates a shift either; that was conservatively added back in D33879.)

1074–1075

For a wide type, this won't work the way you are expecting. Use getAllOnesValue or APInt::getLowBitsSet().

Please add a test like this to verify we don't miscompile:

define i128 @foo(i1 %a) {
  %s = sext i1 %a to i128
  %lshr = lshr i128 %s, 42
  ret i128 %lshr
}

This might need a data layout adjustment to trigger (or the suggestion above to remove shouldChangeType).

llvm/test/Transforms/InstCombine/pr52078.ll

14 ↗

(On Diff #377940)

This should already be covered by an existing test, so it shouldn't be necessary.
We do need other tests though: extra use on the sext, vector type, etc.

[anton-afanasyev]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1069–1070

Actually that's not a restriction since i1 type is desirable so shouldChangeType() check is always true.
But ok, moved that check down for shifts only.

1074–1075

Thanks, fixed and added test.

llvm/test/Transforms/InstCombine/pr52078.ll

14 ↗

(On Diff #377940)

Ok, removed and added new tests

[anton-afanasyev]

Address comments

[spatel]

LGTM - see inline test comments.

llvm/test/Transforms/PhaseOrdering/pr52078.ll

2–3 ↗

(On Diff #378939)

1. Better to create a common prefix if output is always expected to be the same (just the default "CHECK" would do?).
2. Add a -O2 and/or -O3 run, so we know this behaves as expected with the normal pass pipelines (we still want to adjust AIC in those pipelines in another patch?).

[anton-afanasyev]

llvm/test/Transforms/PhaseOrdering/pr52078.ll

2–3 ↗

(On Diff #378939)

1. There's still a small difference in variable names (%lshr and %trunc) after different folding ways.
2. Ok, added -O2, but have to add one more test to show difference here. (And yes, we're still going to adjust pipeline for AIC later).

[anton-afanasyev]

Update tests

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op1, m_APInt(C))) {
     if (match(Op0, m_OneUse(m_ZExt(m_Value(X)))) &&
         (!Ty->isIntegerTy() || shouldChangeType(Ty, X->getType()))) {
       assert(ShAmtC < X->getType()->getScalarSizeInBits() &&
              "Big shift not simplified to zero?");
       // lshr (zext iM X to iN), C --> zext (lshr X, C) to iN
       Value *NewLShr = Builder.CreateLShr(X, ShAmtC);
       return new ZExtInst(NewLShr, Ty);
     }
 
-    if (match(Op0, m_SExt(m_Value(X))) &&
-        (!Ty->isIntegerTy() || shouldChangeType(Ty, X->getType()))) {
-      // Are we moving the sign bit to the low bit and widening with high zeros?
+    if (match(Op0, m_SExt(m_Value(X)))) {
       unsigned SrcTyBitWidth = X->getType()->getScalarSizeInBits();
-      if (ShAmtC == BitWidth - 1) {
-        // lshr (sext i1 X to iN), N-1 --> zext X to iN
-        if (SrcTyBitWidth == 1)
-          return new ZExtInst(X, Ty);
+      // lshr (sext i1 X to iN), C --> select (X, -1 >> C, 0)
+      if (SrcTyBitWidth == 1) {
+        auto *NewC = ConstantInt::get(
+            Ty, APInt::getLowBitsSet(BitWidth, BitWidth - ShAmtC));
+        return SelectInst::Create(X, NewC, ConstantInt::getNullValue(Ty));
+      }
 
-        // lshr (sext iM X to iN), N-1 --> zext (lshr X, M-1) to iN
-        if (Op0->hasOneUse()) {
+      if ((!Ty->isIntegerTy() || shouldChangeType(Ty, X->getType())) &&
+          Op0->hasOneUse()) {
+        // Are we moving the sign bit to the low bit and widening with high
+        // zeros? lshr (sext iM X to iN), N-1 --> zext (lshr X, M-1) to iN
+        if (ShAmtC == BitWidth - 1) {
           Value *NewLShr = Builder.CreateLShr(X, SrcTyBitWidth - 1);
           return new ZExtInst(NewLShr, Ty);
         }
-      }
 
         // lshr (sext iM X to iN), N-M --> zext (ashr X, min(N-M, M-1)) to iN
-      if (ShAmtC == BitWidth - SrcTyBitWidth && Op0->hasOneUse()) {
+        if (ShAmtC == BitWidth - SrcTyBitWidth) {
           // The new shift amount can't be more than the narrow source type.
           unsigned NewShAmt = std::min(ShAmtC, SrcTyBitWidth - 1);
           Value *AShr = Builder.CreateAShr(X, NewShAmt);
           return new ZExtInst(AShr, Ty);
         }
       }
+    }
 
     Value *Y;
     if (ShAmtC == BitWidth - 1) {
       // lshr i32 or(X,-X), 31 --> zext (X != 0)
       if (match(Op0, m_OneUse(m_c_Or(m_Neg(m_Value(X)), m_Deferred(X)))))
         return new ZExtInst(Builder.CreateIsNotNull(X), Ty);
 
       // lshr i32 (X -nsw Y), 31 --> zext (X < Y)

llvm/test/Transforms/InstCombine/lshr.ll

;
   call void @use(i32 %sh)
   %tr = trunc i32 %sh to i12
   %r = lshr i12 %tr, 1
   ret i12 %r
 }
 
 define i16 @lshr_sext_i1_to_i16(i1 %a) {
 ; CHECK-LABEL: @lshr_sext_i1_to_i16(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[A:%.*]] to i16
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i16 [[SEXT]], 4
+; CHECK-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i16 4095, i16 0
 ; CHECK-NEXT:    ret i16 [[LSHR]]
 ;
   %sext = sext i1 %a to i16
   %lshr = lshr i16 %sext, 4
   ret i16 %lshr
 }
 
 define i128 @lshr_sext_i1_to_i128(i1 %a) {
 ; CHECK-LABEL: @lshr_sext_i1_to_i128(
-; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[A:%.*]] to i128
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i128 [[SEXT]], 42
+; CHECK-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i128 77371252455336267181195263, i128 0
 ; CHECK-NEXT:    ret i128 [[LSHR]]
 ;
   %sext = sext i1 %a to i128
   %lshr = lshr i128 %sext, 42
   ret i128 %lshr
 }
 
 define i32 @lshr_sext_i1_to_i32_use(i1 %a) {
 ; CHECK-LABEL: @lshr_sext_i1_to_i32_use(
 ; CHECK-NEXT:    [[SEXT:%.*]] = sext i1 [[A:%.*]] to i32
 ; CHECK-NEXT:    call void @use(i32 [[SEXT]])
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[SEXT]], 14
+; CHECK-NEXT:    [[LSHR:%.*]] = select i1 [[A]], i32 262143, i32 0
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
   %sext = sext i1 %a to i32
   call void @use(i32 %sext)
   %lshr = lshr i32 %sext, 14
   ret i32 %lshr
 }
 
 define <3 x i14> @lshr_sext_i1_to_i14_splat_vec_use1(<3 x i1> %a) {
 ; CHECK-LABEL: @lshr_sext_i1_to_i14_splat_vec_use1(
 ; CHECK-NEXT:    [[SEXT:%.*]] = sext <3 x i1> [[A:%.*]] to <3 x i14>
 ; CHECK-NEXT:    call void @usevec(<3 x i14> [[SEXT]])
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr <3 x i14> [[SEXT]], <i14 4, i14 4, i14 4>
+; CHECK-NEXT:    [[LSHR:%.*]] = select <3 x i1> [[A]], <3 x i14> <i14 1023, i14 1023, i14 1023>, <3 x i14> zeroinitializer
 ; CHECK-NEXT:    ret <3 x i14> [[LSHR]]
 ;
   %sext = sext <3 x i1> %a to <3 x i14>
   call void @usevec(<3 x i14> %sext)
   %lshr = lshr <3 x i14> %sext, <i14 4, i14 4, i14 4>
   ret <3 x i14> %lshr
 }

llvm/test/Transforms/PhaseOrdering/X86/pr52078.ll

 ; CHECK-NEXT:    [[TRUNC:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; CHECK-NEXT:    ret i16 [[TRUNC]]
 ;
 ; IC-LABEL: @foo(
 ; IC-NEXT:    [[TRUNC:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; IC-NEXT:    ret i16 [[TRUNC]]
 ;
 ; AIC_AND_IC-LABEL: @foo(
-; AIC_AND_IC-NEXT:    [[SEXT:%.*]] = sext i1 [[A:%.*]] to i16
-; AIC_AND_IC-NEXT:    [[LSHR:%.*]] = lshr i16 [[SEXT]], 1
+; AIC_AND_IC-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; AIC_AND_IC-NEXT:    ret i16 [[LSHR]]
 ;
   %sext = sext i1 %a to i16
   %zext = zext i16 %sext to i32
   %lshr = lshr i32 %zext, 1
   %trunc = trunc i32 %lshr to i16
   ret i16 %trunc
 }
 
 define i16 @foo2(i1 %a) {
 ; CHECK-LABEL: @foo2(
-; CHECK-NEXT:    [[S:%.*]] = sext i1 [[A:%.*]] to i16
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i16 [[S]], 1
+; CHECK-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; CHECK-NEXT:    ret i16 [[LSHR]]
 ;
 ; IC-LABEL: @foo2(
-; IC-NEXT:    [[S:%.*]] = sext i1 [[A:%.*]] to i16
-; IC-NEXT:    [[LSHR:%.*]] = lshr i16 [[S]], 1
+; IC-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; IC-NEXT:    ret i16 [[LSHR]]
 ;
 ; AIC_AND_IC-LABEL: @foo2(
-; AIC_AND_IC-NEXT:    [[S:%.*]] = sext i1 [[A:%.*]] to i16
-; AIC_AND_IC-NEXT:    [[LSHR:%.*]] = lshr i16 [[S]], 1
+; AIC_AND_IC-NEXT:    [[LSHR:%.*]] = select i1 [[A:%.*]], i16 32767, i16 0
 ; AIC_AND_IC-NEXT:    ret i16 [[LSHR]]
 ;
   %s = sext i1 %a to i16
   %lshr = lshr i16 %s, 1
   ret i16 %lshr
 }