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

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.

Author: michaelmaitland

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);

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -1005,6 +1005,96 @@ RISCVCC::CondCode RISCVCC::getOppositeBranchCondition(RISCVCC::CondCode CC) {
   }
 }
 
+/// Return true if the branch represented by the conditional branch with
+/// components TBB, FBB, and CurCond can folded into an unconditional branch
+/// that branches to FoldedBB.
+///
+/// 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.
+static bool canFoldBrOfCondBr(MachineBasicBlock *TBB, MachineBasicBlock *FBB,
+                              SmallVectorImpl<MachineOperand> &CurCond,
+                              MachineBasicBlock *&FoldedBB) {
+  if (!TBB || !FBB || CurCond.size() != 3)
+    return false;
+
+  RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(CurCond[0].getImm());
+  auto LHS = CurCond[1];
+  auto RHS = CurCond[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))) {
+      FoldedBB = TBB;
+      return true;
+    }
+    // We can statically evaluate that we take the second branch
+    if ((isZero(LHS) && isOne(RHS)) || (isOne(LHS) && isZero(RHS))) {
+      FoldedBB = FBB;
+      return true;
+    }
+    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))) {
+      FoldedBB = TBB;
+      return true;
+    }
+    // We can statically evaluate that we take the second branch
+    if ((isZero(LHS) && isZero(RHS)) || (isOne(LHS) && isOne(RHS))) {
+      FoldedBB = FBB;
+      return true;
+    }
+    return false;
+  }
+  }
+}
+
 bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
                                    MachineBasicBlock *&TBB,
                                    MachineBasicBlock *&FBB,
@@ -1068,8 +1158,24 @@ bool RISCVInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
   // 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);
+
+    // Try to fold the branch of the conditional branch into an unconditional
+    // branch.
+    MachineInstr &CondBr = *std::prev(I);
+    parseCondBranch(CondBr, TBB, Cond);
     FBB = getBranchDestBlock(*I);
+    MachineBasicBlock *Folded = nullptr;
+    if (AllowModify && canFoldBrOfCondBr(TBB, FBB, Cond, Folded)) {
+      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 false;
   }
 

llvm/test/CodeGen/RISCV/branch_zero.ll

@@ -87,11 +87,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