[InstCombine] Fold lshr/ashr(or(neg(x),x),bw-1) --> zext/sext(icmp_ne(x,0)) (PR50816)

[RKSimon]

Summary

Handle the missing fold reported in PR50816, which is a variant of the existing ashr(sub_nsw(X,Y),bw-1) --> sext(icmp_sgt(X,Y)) fold.

We also handle the lshr(or(neg(x),x),bw-1) --> zext(icmp_ne(x,0)) equivalent - https://alive2.llvm.org/ce/z/SnZmSj

We still allow multi uses of the neg(x) - as this is likely to let us further simplify other uses of the neg - but not multi uses of the or() which would increase instruction count.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1328

Might as well handle the lshr -> zext variant at the same time?
https://alive2.llvm.org/ce/z/SnZmSj

llvm/test/Transforms/InstCombine/sub-ashr-or-to-icmp-select.ll

115

This isn't testing the path you intended. We always commute the or operands before we reach the matcher in this patch. Search for "thwart complexity" in the instcombine test directory for examples of how to work around that.

(Best to pre-commit the tests, so we can verify the patterns are as expected.)

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1330

[RKSimon]

Use m_Deferred and add lshr(or(neg(x),x),bw-1) --> zext(icmp_ne(x,0)) handling

[lebedev.ri]

LG

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

1144–1145

Could use Builder.CreateIsNotNull()

1339–1340

Same

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op0, m_SExt(m_Value(X))) &&
         return new ZExtInst(AShr, Ty);
       }
     }
 
     Value *Y;
     if (ShAmt == BitWidth - 1) {
       // lshr i32 (X -nsw Y), 31 --> zext (X < Y)
       if (match(Op0, m_OneUse(m_NSWSub(m_Value(X), m_Value(Y)))))
         return new ZExtInst(Builder.CreateICmpSLT(X, Y), Ty);
 
       // Check if a number is negative and odd:
       // lshr i32 (srem X, 2), 31 --> and (X >> 31), X
       if (match(Op0, m_OneUse(m_SRem(m_Value(X), m_SpecificInt(2))))) {
         Value *Signbit = Builder.CreateLShr(X, ShAmt);
         return BinaryOperator::CreateAnd(Signbit, X);
       }
     }
 
if (match(Op0, m_OneUse(m_SExt(m_Value(X)))) &&
         (Ty->isVectorTy() || shouldChangeType(Ty, X->getType()))) {
       // ashr (sext X), C --> sext (ashr X, C')
       Type *SrcTy = X->getType();
       ShAmt = std::min(ShAmt, SrcTy->getScalarSizeInBits() - 1);
       Value *NewSh = Builder.CreateAShr(X, ConstantInt::get(SrcTy, ShAmt));
       return new SExtInst(NewSh, Ty);
     }
 
-    // ashr i32 (X -nsw Y), 31 --> sext (X < Y)
     Value *Y;
+
+    // ashr i32 or(X,-X), 31 --> sext (X != 0)
+    if (ShAmt == BitWidth - 1 &&
+        match(Op0, m_OneUse(m_c_Or(m_Neg(m_Value(X)), m_Value(Y)))) && X == Y)
+      return new SExtInst(Builder.CreateICmpNE(X, ConstantInt::get(Ty, 0)), Ty);
+
+    // ashr i32 (X -nsw Y), 31 --> sext (X < Y)
     if (ShAmt == BitWidth - 1 &&
         match(Op0, m_OneUse(m_NSWSub(m_Value(X), m_Value(Y)))))
       return new SExtInst(Builder.CreateICmpSLT(X, Y), Ty);
 
     // If the shifted-out value is known-zero, then this is an exact shift.
     if (!I.isExact() &&
         MaskedValueIsZero(Op0, APInt::getLowBitsSet(BitWidth, ShAmt), 0, &I)) {
       I.setIsExact();
       return &I;
     }
   }
 
   if (Instruction *R = foldVariableSignZeroExtensionOfVariableHighBitExtract(I))
     return R;
 

llvm/test/Transforms/InstCombine/sub-ashr-or-to-icmp-select.ll

 ; CHECK-NEXT:    ret i64 [[OR]]
 ;
   %sub = sub nsw i64 %y, %x
   %shr = ashr i64 %sub, 63
   %or = or i64 %shr, %x
   ret i64 %or
 }
 
+define i32 @neg_or_ashr_i32(i32 %x) {
+; CHECK-LABEL: @neg_or_ashr_i32(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = sext i1 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[SHR]]
+;
+  %neg = sub i32 0, %x
+  %or = or i32 %neg, %x
+  %shr = ashr i32 %or, 31
+  ret i32 %shr
+}
+
 ; nuw nsw
 
 define i32 @sub_ashr_or_i32_nuw_nsw(i32 %x, i32 %y) {
 ; CHECK-LABEL: @sub_ashr_or_i32_nuw_nsw(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt i32 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    [[OR:%.*]] = select i1 [[TMP1]], i32 -1, i32 [[X]]
 ; CHECK-NEXT:    ret i32 [[OR]]
 ;

 ; CHECK-NEXT:    ret i32 [[OR]]
 ;
   %sub = sub nsw i32 %y, %x
   %shr = ashr i32 %sub, 31
   %or = or i32 %x, %shr  ; commute %shr and %x
   ret i32 %or
 }
 
+define i32 @neg_or_ashr_i32_commute(i32 %x) {
+; CHECK-LABEL: @neg_or_ashr_i32_commute(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = sext i1 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[SHR]]
+;
+  %neg = sub i32 0, %x
+  %or = or i32 %x, %neg
+  %shr = ashr i32 %or, 31
+  ret i32 %shr
+}
+
 ; Vector Types
 
 define <4 x i32> @sub_ashr_or_i32_vec(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @sub_ashr_or_i32_vec(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <4 x i32> [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    [[OR:%.*]] = select <4 x i1> [[TMP1]], <4 x i32> <i32 -1, i32 -1, i32 -1, i32 -1>, <4 x i32> [[X]]
 ; CHECK-NEXT:    ret <4 x i32> [[OR]]
 ;

 ; CHECK-NEXT:    ret <4 x i32> [[OR]]
 ;
   %sub = sub nuw nsw <4 x i32> %y, %x
   %shr = ashr <4 x i32> %sub, <i32 31, i32 31, i32 31, i32 31>
   %or = or <4 x i32> %shr, %x
   ret <4 x i32> %or
 }
 
+define <4 x i32> @neg_or_ashr_i32_vec(<4 x i32> %x) {
+; CHECK-LABEL: @neg_or_ashr_i32_vec(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne <4 x i32> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[SHR:%.*]] = sext <4 x i1> [[TMP1]] to <4 x i32>
+; CHECK-NEXT:    ret <4 x i32> [[SHR]]
+;
+  %neg = sub <4 x i32> zeroinitializer, %x
+  %or = or <4 x i32> %neg, %x
+  %shr = ashr <4 x i32> %or, <i32 31, i32 31, i32 31, i32 31>
+  ret <4 x i32> %shr
+}
+
 define <4 x i32> @sub_ashr_or_i32_vec_commute(<4 x i32> %x, <4 x i32> %y) {
 ; CHECK-LABEL: @sub_ashr_or_i32_vec_commute(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt <4 x i32> [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    [[OR:%.*]] = select <4 x i1> [[TMP1]], <4 x i32> <i32 -1, i32 -1, i32 -1, i32 -1>, <4 x i32> [[X]]
 ; CHECK-NEXT:    ret <4 x i32> [[OR]]
 ;
   %sub = sub nsw <4 x i32> %y, %x
   %shr = ashr <4 x i32> %sub, <i32 31, i32 31, i32 31, i32 31>
   %or = or <4 x i32> %x, %shr
   ret <4 x i32> %or
 }
 
+define <4 x i32> @neg_or_ashr_i32_vec_commute(<4 x i32> %x) {
+; CHECK-LABEL: @neg_or_ashr_i32_vec_commute(
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne <4 x i32> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[SHR:%.*]] = sext <4 x i1> [[TMP1]] to <4 x i32>
+; CHECK-NEXT:    ret <4 x i32> [[SHR]]
+;
+  %neg = sub <4 x i32> zeroinitializer, %x
+  %or = or <4 x i32> %x, %neg
+  %shr = ashr <4 x i32> %or, <i32 31, i32 31, i32 31, i32 31>
+  ret <4 x i32> %shr
+}
+
 ; Extra uses
 
 define i32 @sub_ashr_or_i32_extra_use_sub(i32 %x, i32 %y, i32* %p) {
 ; CHECK-LABEL: @sub_ashr_or_i32_extra_use_sub(
 ; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i32 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    store i32 [[SUB]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt i32 [[Y]], [[X]]
 ; CHECK-NEXT:    [[OR:%.*]] = select i1 [[TMP1]], i32 -1, i32 [[X]]

 ;
   %sub = sub nsw i32 %y, %x
   %shr = ashr i32 %sub, 31
   %or = or i32 %shr, %x
   store i32 %or, i32* %p
   ret i32 %or
 }
 
+define i32 @neg_extra_use_or_ashr_i32(i32 %x, i32* %p) {
+; CHECK-LABEL: @neg_extra_use_or_ashr_i32(
+; CHECK-NEXT:    [[NEG:%.*]] = sub i32 0, [[X:%.*]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = sext i1 [[TMP1]] to i32
+; CHECK-NEXT:    store i32 [[NEG]], i32* [[P:%.*]], align 4
+; CHECK-NEXT:    ret i32 [[SHR]]
+;
+  %neg = sub i32 0, %x
+  %or = or i32 %neg, %x
+  %shr = ashr i32 %or, 31
+  store i32 %neg, i32* %p
+  ret i32 %shr
+}
+
 ; Negative Tests
 
 define i32 @sub_ashr_or_i32_extra_use_ashr(i32 %x, i32 %y, i32* %p) {
 ; CHECK-LABEL: @sub_ashr_or_i32_extra_use_ashr(
 ; CHECK-NEXT:    [[TMP1:%.*]] = icmp slt i32 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    [[SHR:%.*]] = sext i1 [[TMP1]] to i32
 ; CHECK-NEXT:    store i32 [[SHR]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[OR:%.*]] = or i32 [[SHR]], [[X]]

 ; CHECK-NEXT:    ret i32 [[OR]]
 ;
   %sub = sub i32 %y, %x
   %shr = ashr i32 %sub, 31
   %or = or i32 %shr, %x
   ret i32 %or
 }
 
+define i32 @neg_or_extra_use_ashr_i32(i32 %x, i32* %p) {
+; CHECK-LABEL: @neg_or_extra_use_ashr_i32(
+; CHECK-NEXT:    [[NEG:%.*]] = sub i32 0, [[X:%.*]]
+; CHECK-NEXT:    [[OR:%.*]] = or i32 [[NEG]], [[X]]
+; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[OR]], 31
+; CHECK-NEXT:    store i32 [[OR]], i32* [[P:%.*]], align 4
+; CHECK-NEXT:    ret i32 [[SHR]]
+;
+  %neg = sub i32 0, %x
+  %or = or i32 %neg, %x
+  %shr = ashr i32 %or, 31
+  store i32 %or, i32* %p
+  ret i32 %shr
+}
+
 define <4 x i32> @sub_ashr_or_i32_vec_undef1(<4 x i32> %x) {
 ; CHECK-LABEL: @sub_ashr_or_i32_vec_undef1(
 ; CHECK-NEXT:    [[SUB:%.*]] = sub <4 x i32> <i32 255, i32 255, i32 undef, i32 255>, [[X:%.*]]
 ; CHECK-NEXT:    [[SHR:%.*]] = ashr <4 x i32> [[SUB]], <i32 31, i32 31, i32 31, i32 31>
 ; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[SHR]], [[X]]
 ; CHECK-NEXT:    ret <4 x i32> [[OR]]
 ;
   %sub = sub <4 x i32> <i32 255, i32 255, i32 undef, i32 255>, %x