[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

1333

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

1335

[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()

1335–1336

Same

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (match(Op0, m_SExt(m_Value(X))) &&
         unsigned NewShAmt = std::min(ShAmt, SrcTyBitWidth - 1);
         Value *AShr = Builder.CreateAShr(X, NewShAmt);
         return new ZExtInst(AShr, Ty);
       }
     }
 
     Value *Y;
     if (ShAmt == 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))))) {
+        Constant *Zero = ConstantInt::get(Ty, 0);
+        return new ZExtInst(Builder.CreateICmpNE(X, Zero), Ty);
+      }
+
       // 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);
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);
     }
 
+    if (ShAmt == BitWidth - 1) {
+      // ashr i32 or(X,-X), 31 --> sext (X != 0)
+      if (match(Op0, m_OneUse(m_c_Or(m_Neg(m_Value(X)), m_Deferred(X))))) {
+        Constant *Zero = ConstantInt::get(Ty, 0);
+        return new SExtInst(Builder.CreateICmpNE(X, Zero), Ty);
+      }
+
       // ashr i32 (X -nsw Y), 31 --> sext (X < Y)
       Value *Y;
-    if (ShAmt == BitWidth - 1 &&
-        match(Op0, m_OneUse(m_NSWSub(m_Value(X), m_Value(Y)))))
+      if (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;
     }
   }

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

;
   %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:    [[NEG:%.*]] = sub i32 0, [[X:%.*]]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[OR]], 31
+; 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
 }
 

 ; 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 %x0) {
 ; CHECK-LABEL: @neg_or_ashr_i32_commute(
 ; CHECK-NEXT:    [[X:%.*]] = sdiv i32 42, [[X0:%.*]]
-; CHECK-NEXT:    [[NEG:%.*]] = sub nsw i32 0, [[X]]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[X]], [[NEG]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[OR]], 31
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = sext i1 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %x = sdiv i32 42, %x0 ; thwart complexity-based canonicalization
   %neg = sub i32 0, %x
   %or = or i32 %x, %neg
   %shr = ashr i32 %or, 31
   ret i32 %shr
 }
;
   %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:    [[NEG:%.*]] = sub <4 x i32> zeroinitializer, [[X:%.*]]
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr <4 x i32> [[OR]], <i32 31, i32 31, i32 31, i32 31>
+; 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> %x0) {
 ; CHECK-LABEL: @neg_or_ashr_i32_vec_commute(
 ; CHECK-NEXT:    [[X:%.*]] = sdiv <4 x i32> <i32 42, i32 42, i32 42, i32 42>, [[X0:%.*]]
-; CHECK-NEXT:    [[NEG:%.*]] = sub nsw <4 x i32> zeroinitializer, [[X]]
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[X]], [[NEG]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr <4 x i32> [[OR]], <i32 31, i32 31, i32 31, i32 31>
+; 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]]
 ;
   %x = sdiv <4 x i32> <i32 42, i32 42, i32 42, i32 42>, %x0 ; thwart complexity-based canonicalization
   %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
 }
;
   %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:    [[OR:%.*]] = or i32 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = ashr i32 [[OR]], 31
+; 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

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -instcombine %s -S -o - | FileCheck %s
 
 ; Scalar Types
 
 define i32 @neg_or_lshr_i32(i32 %x) {
 ; CHECK-LABEL: @neg_or_lshr_i32(
-; CHECK-NEXT:    [[NEG:%.*]] = sub i32 0, [[X:%.*]]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[OR]], 31
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X:%.*]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = zext i1 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %neg = sub i32 0, %x
   %or = or i32 %neg, %x
   %shr = lshr i32 %or, 31
   ret i32 %shr
 }
 
 ; Commute
 
 define i32 @neg_or_lshr_i32_commute(i32 %x0) {
 ; CHECK-LABEL: @neg_or_lshr_i32_commute(
 ; CHECK-NEXT:    [[X:%.*]] = sdiv i32 42, [[X0:%.*]]
-; CHECK-NEXT:    [[NEG:%.*]] = sub nsw i32 0, [[X]]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[X]], [[NEG]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[OR]], 31
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = zext i1 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[SHR]]
 ;
   %x = sdiv i32 42, %x0 ; thwart complexity-based canonicalization
   %neg = sub i32 0, %x
   %or = or i32 %x, %neg
   %shr = lshr i32 %or, 31
   ret i32 %shr
 }
 
 ; Vector Types
 
 define <4 x i32> @neg_or_lshr_i32_vec(<4 x i32> %x) {
 ; CHECK-LABEL: @neg_or_lshr_i32_vec(
-; CHECK-NEXT:    [[NEG:%.*]] = sub <4 x i32> zeroinitializer, [[X:%.*]]
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <4 x i32> [[OR]], <i32 31, i32 31, i32 31, i32 31>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne <4 x i32> [[X:%.*]], zeroinitializer
+; CHECK-NEXT:    [[SHR:%.*]] = zext <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 = lshr <4 x i32> %or, <i32 31, i32 31, i32 31, i32 31>
   ret <4 x i32> %shr
 }
 
 define <4 x i32> @neg_or_lshr_i32_vec_commute(<4 x i32> %x0) {
 ; CHECK-LABEL: @neg_or_lshr_i32_vec_commute(
 ; CHECK-NEXT:    [[X:%.*]] = sdiv <4 x i32> <i32 42, i32 42, i32 42, i32 42>, [[X0:%.*]]
-; CHECK-NEXT:    [[NEG:%.*]] = sub nsw <4 x i32> zeroinitializer, [[X]]
-; CHECK-NEXT:    [[OR:%.*]] = or <4 x i32> [[X]], [[NEG]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr <4 x i32> [[OR]], <i32 31, i32 31, i32 31, i32 31>
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne <4 x i32> [[X]], zeroinitializer
+; CHECK-NEXT:    [[SHR:%.*]] = zext <4 x i1> [[TMP1]] to <4 x i32>
 ; CHECK-NEXT:    ret <4 x i32> [[SHR]]
 ;
   %x = sdiv <4 x i32> <i32 42, i32 42, i32 42, i32 42>, %x0 ; thwart complexity-based canonicalization
   %neg = sub <4 x i32> zeroinitializer, %x
   %or = or <4 x i32> %x, %neg
   %shr = lshr <4 x i32> %or, <i32 31, i32 31, i32 31, i32 31>
   ret <4 x i32> %shr
 }
 
 ; Extra uses
 
 define i32 @neg_extra_use_or_lshr_i32(i32 %x, i32* %p) {
 ; CHECK-LABEL: @neg_extra_use_or_lshr_i32(
 ; CHECK-NEXT:    [[NEG:%.*]] = sub i32 0, [[X:%.*]]
-; CHECK-NEXT:    [[OR:%.*]] = or i32 [[NEG]], [[X]]
-; CHECK-NEXT:    [[SHR:%.*]] = lshr i32 [[OR]], 31
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32 [[X]], 0
+; CHECK-NEXT:    [[SHR:%.*]] = zext 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 = lshr i32 %or, 31
   store i32 %neg, i32* %p
   ret i32 %shr