[X86][CodeGenPrepare] Enable zero comparison heuristic

llvm/include/llvm/CodeGen/TargetLowering.h

public:
   /// Return true if instruction generated for equality comparison is folded
   /// with instruction generated for signed comparison.
   virtual bool isEqualityCmpFoldedWithSignedCmp() const { return true; }
 
   /// Return true if the heuristic to prefer icmp eq zero should be used in code
   /// gen prepare.
   virtual bool preferZeroCompareBranch() const { return false; }
 
+  /// Return true if the heuristics to prefer icmp <any> zero should be used in
+  /// code gen prepare.
+  virtual bool preferAnyZeroCompareBranch() const { return false; }
+
   /// Return true if it is cheaper to split the store of a merged int val
   /// from a pair of smaller values into multiple stores.
   virtual bool isMultiStoresCheaperThanBitsMerge(EVT LTy, EVT HTy) const {
     return false;
   }
 
   /// Return if the target supports combining a
   /// chain like:

llvm/lib/CodeGen/CodeGenPrepare.cpp

static bool optimizeBranch(BranchInst *Branch, const TargetLowering &TLI,
   //  br %c, bla, blb
   //  %tc = lshr %x, 3
   // to
   //  %tc = lshr %x, 3
   //  %c = icmp eq %tc, 0
   //  br %c, bla, blb
   // Creating the cmp to zero can be better for the backend, especially if the
   // lshr produces flags that can be used automatically.
-  if (!TLI.preferZeroCompareBranch() || !Branch->isConditional())
+  if ((!TLI.preferZeroCompareBranch() && !TLI.preferAnyZeroCompareBranch()) ||
+      !Branch->isConditional())
     return false;
 
   ICmpInst *Cmp = dyn_cast<ICmpInst>(Branch->getCondition());
   if (!Cmp || !isa<ConstantInt>(Cmp->getOperand(1)) || !Cmp->hasOneUse())
     return false;
 
   Value *X = Cmp->getOperand(0);
   APInt CmpC = cast<ConstantInt>(Cmp->getOperand(1))->getValue();
if (CmpC.isPowerOf2() && Cmp->getPredicate() == ICmpInst::ICMP_ULT &&
         UI->moveBefore(Branch);
       Value *NewCmp = Builder.CreateCmp(ICmpInst::ICMP_EQ, UI,
                                         ConstantInt::get(UI->getType(), 0));
       LLVM_DEBUG(dbgs() << "Converting " << *Cmp << "\n");
       LLVM_DEBUG(dbgs() << " to compare on zero: " << *NewCmp << "\n");
       replaceAllUsesWith(Cmp, NewCmp, FreshBBs, IsHugeFunc);
       return true;
     }
-    if (Cmp->isEquality() &&
+    if ((Cmp->isEquality() || TLI.preferAnyZeroCompareBranch()) &&
         (match(UI, m_Add(m_Specific(X), m_SpecificInt(-CmpC))) ||
          match(UI, m_Sub(m_Specific(X), m_SpecificInt(CmpC))))) {
       IRBuilder<> Builder(Branch);
       if (UI->getParent() != Branch->getParent())
         UI->moveBefore(Branch);
       Value *NewCmp = Builder.CreateCmp(Cmp->getPredicate(), UI,
                                         ConstantInt::get(UI->getType(), 0));
       LLVM_DEBUG(dbgs() << "Converting " << *Cmp << "\n");

llvm/lib/Target/X86/X86ISelLowering.h

bool canMergeStoresTo(unsigned AddressSpace, EVT MemVT,
                           const MachineFunction &MF) const override;
 
     bool isCheapToSpeculateCttz(Type *Ty) const override;
 
     bool isCheapToSpeculateCtlz(Type *Ty) const override;
 
     bool isCtlzFast() const override;
 
+    bool preferAnyZeroCompareBranch() const override { return true; }
+
     bool isMultiStoresCheaperThanBitsMerge(EVT LTy, EVT HTy) const override {
       // If the pair to store is a mixture of float and int values, we will
       // save two bitwise instructions and one float-to-int instruction and
       // increase one store instruction. There is potentially a more
       // significant benefit because it avoids the float->int domain switch
       // for input value. So It is more likely a win.
       if ((LTy.isFloatingPoint() && HTy.isInteger()) ||
           (LTy.isInteger() && HTy.isFloatingPoint()))

llvm/test/CodeGen/X86/2006-05-11-InstrSched.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; REQUIRES: asserts
 ; RUN: llc < %s -mtriple=i386-linux-gnu -mcpu=penryn -mattr=+sse2 -stats 2>&1 | \
-; RUN:     grep "asm-printer" | grep 33
+; RUN:     grep "asm-printer" | grep 32
 
 target datalayout = "e-p:32:32"
 define void @foo(ptr %mc, ptr %bp, ptr %ms, ptr %xmb, ptr %mpp, ptr %tpmm, ptr %ip, ptr %tpim, ptr %dpp, ptr %tpdm, ptr %bpi, i32 %M) nounwind {
 entry:
   %tmp9 = icmp slt i32 %M, 5    ; <i1> [#uses=1]
   br i1 %tmp9, label %return, label %cond_true
 
 cond_true:    ; preds = %cond_true, %entry

llvm/test/CodeGen/X86/branch-on-zero.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=x86_64-unknown-unknown < %s | FileCheck %s
 
 declare void @foo(i32)
 
 define void @test1_slt(i32 %0) {
 ; CHECK-LABEL: test1_slt:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    cmpl $2, %edi
-; CHECK-NEXT:    jg .LBB0_1
-; CHECK-NEXT:  # %bb.2:
 ; CHECK-NEXT:    addl $-3, %edi
-; CHECK-NEXT:    jmp foo@PLT # TAILCALL
-; CHECK-NEXT:  .LBB0_1:
+; CHECK-NEXT:    js foo@PLT # TAILCALL
+; CHECK-NEXT:  # %bb.1:
 ; CHECK-NEXT:    retq
   %2 = icmp slt i32 %0, 3
   br i1 %2, label %3, label %5
 
 3:                                                ; preds = %1
   %4 = add nsw i32 %0, -3
   tail call void @foo(i32 %4)
   ret void
 
 5:                                                ; preds = %1
   ret void
 }
 
 define void @test2_sle(i32 %0) {
 ; CHECK-LABEL: test2_sle:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    cmpl $3, %edi
-; CHECK-NEXT:    jg .LBB1_1
-; CHECK-NEXT:  # %bb.2:
 ; CHECK-NEXT:    addl $-3, %edi
-; CHECK-NEXT:    jmp foo@PLT # TAILCALL
-; CHECK-NEXT:  .LBB1_1:
+; CHECK-NEXT:    jle foo@PLT # TAILCALL
+; CHECK-NEXT:  # %bb.1:
 ; CHECK-NEXT:    retq
   %2 = icmp sle i32 %0, 3
   br i1 %2, label %3, label %5
 
 3:                                                ; preds = %1
   %4 = add nsw i32 %0, -3
   tail call void @foo(i32 %4)
   ret void
 
 5:                                                ; preds = %1
   ret void
 }
 
 define void @test3_ugt(i32 %0) {
 ; CHECK-LABEL: test3_ugt:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    cmpl $4, %edi
-; CHECK-NEXT:    jb .LBB2_1
-; CHECK-NEXT:  # %bb.2:
 ; CHECK-NEXT:    addl $-3, %edi
-; CHECK-NEXT:    jmp foo@PLT # TAILCALL
-; CHECK-NEXT:  .LBB2_1:
+; CHECK-NEXT:    jne foo@PLT # TAILCALL
+; CHECK-NEXT:  # %bb.1:
 ; CHECK-NEXT:    retq
   %2 = icmp ugt i32 %0, 3
   br i1 %2, label %3, label %5
 
 3:                                                ; preds = %1
   %4 = add nuw nsw i32 %0, -3
   tail call void @foo(i32 noundef %4)
   ret void

llvm/test/CodeGen/X86/peep-test-5.ll

 ; Resulting assembly should share flags from single CMP instruction for both
 ; conditions!
 define void @decref(ptr %p) {
 ; CHECK-LABEL: decref:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    pushq %rax
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    movl (%rdi), %eax
-; CHECK-NEXT:    cmpl $1, %eax
-; CHECK-NEXT:    jne .LBB0_2
-; CHECK-NEXT:  # %bb.1: # %bb_free
-; CHECK-NEXT:    callq free_object@PLT
-; CHECK-NEXT:  .LBB0_4: # %end
+; CHECK-NEXT:    decl %eax
+; CHECK-NEXT:    je .LBB0_3
+; CHECK-NEXT:  # %bb.1: # %bb2
+; CHECK-NEXT:    testl %eax, %eax
+; CHECK-NEXT:    jle .LBB0_4
+; CHECK-NEXT:  # %bb.2: # %bb_dec
+; CHECK-NEXT:    movl %eax, (%rdi)
 ; CHECK-NEXT:    popq %rax
 ; CHECK-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-NEXT:    retq
-; CHECK-NEXT:  .LBB0_2: # %bb2
+; CHECK-NEXT:  .LBB0_3: # %bb_free
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    jle .LBB0_4
-; CHECK-NEXT:  # %bb.3: # %bb_dec
-; CHECK-NEXT:    decl %eax
-; CHECK-NEXT:    movl %eax, (%rdi)
+; CHECK-NEXT:    callq free_object@PLT
+; CHECK-NEXT:  .LBB0_4: # %end
 ; CHECK-NEXT:    popq %rax
 ; CHECK-NEXT:    .cfi_def_cfa_offset 8
 ; CHECK-NEXT:    retq
   %count = load i32, ptr %p, align 4
   %cmp0 = icmp eq i32 %count, 1
   br i1 %cmp0, label %bb_free, label %bb2
 
 bb2: