[BranchFolding][RISCV] Optimize branches NE/EQ to zero or one

michaelmaitland · michaelmaitland · commit bd669d381617 · 2025-03-17T16:22:15.000-07:00
There are two changes here.

The first that we teach analyzeBranch how to
evaluate a conditional branch followed by a unconditional branch such that
we can evaluate the conditional branch statically. Often, we will see comparison
to one or zero since SelectionDAG often uses i1 for the conditional comparison.
As a result, we handle this specific case. We handle only EQ and NEQ for now,
but this can be expanded in the future. We can also expand on handling arbitrary
constants in the future.

The second change is that we pass AllowModify=false to analyzeBranch in the
tail merging code. The reason we do this is because this code is doing some
clever tricks to the branch code that it will restore later. Now that we are
actually optimizing branches in analyzeBranch, we have to be careful not to
mess up this canonical form that the tail merging code expects.
diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1107,7 +1107,7 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
 
       MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
       SmallVector<MachineOperand, 4> Cond;
-      if (!TII->analyzeBranch(*PBB, TBB, FBB, Cond, true)) {
+      if (!TII->analyzeBranch(*PBB, TBB, FBB, Cond)) {
         // Failing case: IBB is the target of a cbr, and we cannot reverse the
         // branch.
         SmallVector<MachineOperand, 4> NewCond(Cond);
diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp b/llvm/lib/Target/RISCV/RISCVInstrInfo.cpp
@@ -1005,6 +1005,104 @@ RISCVCC::CondCode RISCVCC::getOppositeBranchCondition(RISCVCC::CondCode CC) {
   }
 }
 
+bool RISCVInstrInfo::trySimplifyCondBr(MachineBasicBlock &MBB,
+                                       MachineBasicBlock *TBB,
+                                       MachineBasicBlock *FBB,
+                                       SmallVectorImpl<MachineOperand> &Cond,
+                                       MachineBasicBlock *&Folded) const {
+  if (!TBB || Cond.size() != 3)
+    return false;
+
+  // If there is no FBB, get the fall through.
+  if (!FBB) {
+    MachineFunction::iterator FallThrough = ++MBB.getIterator();
+    // There is no fallthrough
+    if (FallThrough == MBB.getParent()->end())
+      return false;
+    FBB = &*FallThrough;
+  }
+
+  RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(Cond[0].getImm());
+  auto LHS = Cond[1];
+  auto RHS = Cond[2];
+
+  // Return true if MO definitely contains the value one.
+  auto isOne = [](MachineOperand &MO) -> bool {
+    if (MO.isImm() && MO.getImm() == 1)
+      return true;
+
+    if (!MO.isReg() || !MO.getReg().isVirtual())
+      return false;
+
+    MachineRegisterInfo &MRI =
+        MO.getParent()->getParent()->getParent()->getRegInfo();
+    MachineInstr *DefMI = MRI.getUniqueVRegDef(MO.getReg());
+    if (!DefMI)
+      return false;
+
+    // For now, just check the canonical zero value.
+    if (DefMI->getOpcode() == RISCV::ADDI &&
+        DefMI->getOperand(1).getReg() == RISCV::X0 &&
+        DefMI->getOperand(2).getImm() == 1)
+      return true;
+
+    return false;
+  };
+
+  // Return true if MO definitely contains the value zero.
+  auto isZero = [](MachineOperand &MO) -> bool {
+    if (MO.isImm() && MO.getImm() == 0)
+      return true;
+    if (MO.isReg() && MO.getReg() == RISCV::X0)
+      return true;
+    return false;
+  };
+
+  switch (CC) {
+  default:
+    // TODO: Implement for more CCs
+    return false;
+  case RISCVCC::COND_EQ: {
+    // We can statically evaluate that we take the first branch
+    if ((isZero(LHS) && isZero(RHS)) || (isOne(LHS) && isOne(RHS))) {
+      Folded = TBB;
+      break;
+    }
+    // We can statically evaluate that we take the second branch
+    if ((isZero(LHS) && isOne(RHS)) || (isOne(LHS) && isZero(RHS))) {
+      Folded = FBB;
+      break;
+    }
+    return false;
+  }
+  case RISCVCC::COND_NE: {
+    // We can statically evaluate that we take the first branch
+    if ((isOne(LHS) && isZero(RHS)) || (isZero(LHS) && isOne(RHS))) {
+      Folded = TBB;
+      break;
+    }
+    // We can statically evaluate that we take the second branch
+    if ((isZero(LHS) && isZero(RHS)) || (isOne(LHS) && isOne(RHS))) {
+      Folded = FBB;
+      break;
+    }
+    return false;
+  }
+  }
+
+  // At this point, its legal to optimize.
+  DebugLoc DL = MBB.findBranchDebugLoc();
+  removeBranch(MBB);
+  insertBranch(MBB, Folded, nullptr, {}, DL);
+  while (!MBB.succ_empty())
+    MBB.removeSuccessor(MBB.succ_end() - 1);
+  MBB.addSuccessor(Folded);
+  TBB = Folded;
+  FBB = nullptr;
+  Cond.clear();
+  return true;
+}
+
 bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
                                    MachineBasicBlock *&TBB,
                                    MachineBasicBlock *&FBB,
@@ -1062,14 +1160,24 @@ bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
   // Handle a single conditional branch.
   if (NumTerminators == 1 && I->getDesc().isConditionalBranch()) {
     parseCondBranch(*I, TBB, Cond);
+    // Try to fold the branch of the conditional branch into a the fallthru.
+    MachineBasicBlock *Folded = nullptr;
+    if (AllowModify)
+      trySimplifyCondBr(MBB, TBB, FBB, Cond, Folded);
     return false;
   }
 
   // Handle a conditional branch followed by an unconditional branch.
   if (NumTerminators == 2 && std::prev(I)->getDesc().isConditionalBranch() &&
       I->getDesc().isUnconditionalBranch()) {
-    parseCondBranch(*std::prev(I), TBB, Cond);
+        MachineInstr &CondBr = *std::prev(I);
+    parseCondBranch(CondBr, TBB, Cond);
     FBB = getBranchDestBlock(*I);
+    // Try to fold the branch of the conditional branch into an unconditional
+    // branch.
+    MachineBasicBlock *Folded = nullptr;
+    if (AllowModify)
+      trySimplifyCondBr(MBB, TBB, FBB, Cond, Folded);
     return false;
   }
 
diff --git a/llvm/lib/Target/RISCV/RISCVInstrInfo.h b/llvm/lib/Target/RISCV/RISCVInstrInfo.h
@@ -319,6 +319,28 @@ class RISCVInstrInfo : public RISCVGenInstrInfo {
                                const MachineInstr &MI2) const;
   bool hasReassociableVectorSibling(const MachineInstr &Inst,
                                     bool &Commuted) const;
+  /// Return true if the branch represented by the conditional branch with
+  /// components TBB, FBB, and CurCond was folded into an unconditional branch
+  /// that branches to FoldedBB.
+  ///
+  /// If FBB is nullptr, then the the input represents a conditional branch with
+  /// a fallthrough.
+  ///
+  /// For example:
+  /// BRCOND EQ 0, 0, BB1
+  /// BR BB2
+  ///
+  /// can be simplified to BR BB1 since 0 == 0 statically. On the other hand,
+  ///
+  ///
+  /// BRCOND EQ 0, 1, BB1
+  /// BR BB2
+  ///
+  /// can be simplified to BR BB2 because 0 != 1 statically.
+  bool trySimplifyCondBr(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
+                         MachineBasicBlock *FBB,
+                         SmallVectorImpl<MachineOperand> &Cond,
+                         MachineBasicBlock *&FoldedBB) const;
 };
 
 namespace RISCV {
diff --git a/llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll b/llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll
@@ -357,11 +357,6 @@ define i64 @ctpop_i64(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_two:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    beqz zero, .LBB6_2
-; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    sltiu a0, zero, 0
-; RV32I-NEXT:    ret
-; RV32I-NEXT:  .LBB6_2:
 ; RV32I-NEXT:    srli a2, a0, 1
 ; RV32I-NEXT:    lui a3, 349525
 ; RV32I-NEXT:    lui a4, 209715
@@ -404,11 +399,6 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 ;
 ; RV32ZBB-LABEL: ctpop_i64_ugt_two:
 ; RV32ZBB:       # %bb.0:
-; RV32ZBB-NEXT:    beqz zero, .LBB6_2
-; RV32ZBB-NEXT:  # %bb.1:
-; RV32ZBB-NEXT:    sltiu a0, zero, 0
-; RV32ZBB-NEXT:    ret
-; RV32ZBB-NEXT:  .LBB6_2:
 ; RV32ZBB-NEXT:    cpop a0, a0
 ; RV32ZBB-NEXT:    cpop a1, a1
 ; RV32ZBB-NEXT:    add a0, a1, a0
@@ -422,11 +412,6 @@ define i1 @ctpop_i64_ugt_two(i64 %a) nounwind {
 define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
 ; RV32I-LABEL: ctpop_i64_ugt_one:
 ; RV32I:       # %bb.0:
-; RV32I-NEXT:    beqz zero, .LBB7_2
-; RV32I-NEXT:  # %bb.1:
-; RV32I-NEXT:    snez a0, zero
-; RV32I-NEXT:    ret
-; RV32I-NEXT:  .LBB7_2:
 ; RV32I-NEXT:    srli a2, a0, 1
 ; RV32I-NEXT:    lui a3, 349525
 ; RV32I-NEXT:    lui a4, 209715
@@ -470,11 +455,6 @@ define i1 @ctpop_i64_ugt_one(i64 %a) nounwind {
 ;
 ; RV32ZBB-LABEL: ctpop_i64_ugt_one:
 ; RV32ZBB:       # %bb.0:
-; RV32ZBB-NEXT:    beqz zero, .LBB7_2
-; RV32ZBB-NEXT:  # %bb.1:
-; RV32ZBB-NEXT:    snez a0, zero
-; RV32ZBB-NEXT:    ret
-; RV32ZBB-NEXT:  .LBB7_2:
 ; RV32ZBB-NEXT:    cpop a0, a0
 ; RV32ZBB-NEXT:    cpop a1, a1
 ; RV32ZBB-NEXT:    add a0, a1, a0
diff --git a/llvm/test/CodeGen/RISCV/branch_zero.ll b/llvm/test/CodeGen/RISCV/branch_zero.ll
@@ -5,16 +5,10 @@
 define void @foo(i16 %finder_idx) {
 ; CHECK-LABEL: foo:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:  .LBB0_1: # %for.body
-; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    slli a0, a0, 48
-; CHECK-NEXT:    bltz a0, .LBB0_4
-; CHECK-NEXT:  # %bb.2: # %while.cond.preheader.i
-; CHECK-NEXT:    # in Loop: Header=BB0_1 Depth=1
-; CHECK-NEXT:    li a0, 0
-; CHECK-NEXT:    bnez zero, .LBB0_1
-; CHECK-NEXT:  # %bb.3: # %while.body
-; CHECK-NEXT:  .LBB0_4: # %while.cond1.preheader.i
+; CHECK-NEXT:    bltz a0, .LBB0_2
+; CHECK-NEXT:  # %bb.1: # %while.cond.preheader.i
+; CHECK-NEXT:  .LBB0_2: # %while.cond1.preheader.i
 entry:
   br label %for.body
 
@@ -46,16 +40,10 @@ if.then:
 define void @bar(i16 %finder_idx) {
 ; CHECK-LABEL: bar:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:  .LBB1_1: # %for.body
-; CHECK-NEXT:    # =>This Inner Loop Header: Depth=1
 ; CHECK-NEXT:    slli a0, a0, 48
-; CHECK-NEXT:    bgez a0, .LBB1_4
-; CHECK-NEXT:  # %bb.2: # %while.cond.preheader.i
-; CHECK-NEXT:    # in Loop: Header=BB1_1 Depth=1
-; CHECK-NEXT:    li a0, 0
-; CHECK-NEXT:    bnez zero, .LBB1_1
-; CHECK-NEXT:  # %bb.3: # %while.body
-; CHECK-NEXT:  .LBB1_4: # %while.cond1.preheader.i
+; CHECK-NEXT:    bgez a0, .LBB1_2
+; CHECK-NEXT:  # %bb.1: # %while.cond.preheader.i
+; CHECK-NEXT:  .LBB1_2: # %while.cond1.preheader.i
 entry:
   br label %for.body
 
@@ -87,11 +75,9 @@ if.then:
 define ptr @baz() {
 ; CHECK-LABEL: baz:
 ; CHECK:       # %bb.0: # %entry
-; CHECK-NEXT:    beqz zero, .LBB2_2
 ; CHECK-NEXT:  # %bb.1: # %if.end12
 ; CHECK-NEXT:    li a0, 0
 ; CHECK-NEXT:    ret
-; CHECK-NEXT:  .LBB2_2: # %if.then10
 entry:
   %or.cond = or i1 false, false
   br i1 %or.cond, label %if.end12, label %if.then10
diff --git a/llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll b/llvm/test/CodeGen/RISCV/push-pop-opt-crash.ll
@@ -1,3 +1,4 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
 ; RUN: llc -mattr=+zcmp -verify-machineinstrs  \
 ; RUN: -mtriple=riscv32 -target-abi=ilp32 < %s \
 ; RUN: | FileCheck %s -check-prefixes=RV32IZCMP
@@ -13,36 +14,27 @@ declare dso_local void @f1() local_unnamed_addr
 declare dso_local void @f2() local_unnamed_addr
 define  dso_local void @f0() local_unnamed_addr {
 ; RV32IZCMP-LABEL: f0:
-; RV32IZCMP: 	.cfi_startproc
-; RV32IZCMP-NEXT: # %bb.0:                                # %entry
-; RV32IZCMP-NEXT: 	bnez	zero, .LBB0_2
-; RV32IZCMP-NEXT: # %bb.1:                                # %if.T
-; RV32IZCMP-NEXT: 	cm.push	{ra}, -16
-; RV32IZCMP-NEXT: 	.cfi_def_cfa_offset 16
-; RV32IZCMP-NEXT: 	.cfi_offset ra, -4
-; RV32IZCMP-NEXT: 	call	f1
-; RV32IZCMP-NEXT: 	cm.pop	{ra}, 16
-; RV32IZCMP-NEXT:     .cfi_restore ra
-; RV32IZCMP-NEXT:     .cfi_def_cfa_offset 0
-; RV32IZCMP-NEXT: .LBB0_2:                                # %if.F
-; RV32IZCMP-NEXT: 	tail	f2
-; RV32IZCMP-NEXT: .Lfunc_end0:
-
+; RV32IZCMP:       # %bb.0: # %entry
+; RV32IZCMP-NEXT:    cm.push {ra}, -16
+; RV32IZCMP-NEXT:    .cfi_def_cfa_offset 16
+; RV32IZCMP-NEXT:    .cfi_offset ra, -4
+; RV32IZCMP-NEXT:    call f1
+; RV32IZCMP-NEXT:    cm.pop {ra}, 16
+; RV32IZCMP-NEXT:    .cfi_restore ra
+; RV32IZCMP-NEXT:    .cfi_def_cfa_offset 0
+; RV32IZCMP-NEXT:    tail f2
+;
 ; RV64IZCMP-LABEL: f0:
-; RV64IZCMP: 	.cfi_startproc
-; RV64IZCMP-NEXT: # %bb.0:                                # %entry
-; RV64IZCMP-NEXT: 	bnez	zero, .LBB0_2
-; RV64IZCMP-NEXT: # %bb.1:                                # %if.T
-; RV64IZCMP-NEXT: 	cm.push	{ra}, -16
-; RV64IZCMP-NEXT: 	.cfi_def_cfa_offset 16
-; RV64IZCMP-NEXT: 	.cfi_offset ra, -8
-; RV64IZCMP-NEXT: 	call	f1
-; RV64IZCMP-NEXT: 	cm.pop	{ra}, 16
-; RV64IZCMP-NEXT:     .cfi_restore ra
-; RV64IZCMP-NEXT:     .cfi_def_cfa_offset 0
-; RV64IZCMP-NEXT: .LBB0_2:                                # %if.F
-; RV64IZCMP-NEXT: 	tail	f2
-; RV64IZCMP-NEXT: .Lfunc_end0:
+; RV64IZCMP:       # %bb.0: # %entry
+; RV64IZCMP-NEXT:    cm.push {ra}, -16
+; RV64IZCMP-NEXT:    .cfi_def_cfa_offset 16
+; RV64IZCMP-NEXT:    .cfi_offset ra, -8
+; RV64IZCMP-NEXT:    call f1
+; RV64IZCMP-NEXT:    cm.pop {ra}, 16
+; RV64IZCMP-NEXT:    .cfi_restore ra
+; RV64IZCMP-NEXT:    .cfi_def_cfa_offset 0
+; RV64IZCMP-NEXT:    tail f2
+
 entry:
   br i1 poison, label %if.T, label %if.F
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/vxrm-insert-out-of-loop.ll b/llvm/test/CodeGen/RISCV/rvv/vxrm-insert-out-of-loop.ll
