[RISCV] Canonicalize foldable branch conditions in optimizeCondBranch

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -998,6 +998,25 @@ static RISCVCC::CondCode getCondFromBranchOpc(unsigned Opc) {
   }
 }
 
+static bool evaluateCondBranch(unsigned CC, int64_t C0, int64_t C1) {
+  switch (CC) {
+  default:
+    llvm_unreachable("Unexpected CC");
+  case RISCVCC::COND_EQ:
+    return C0 == C1;
+  case RISCVCC::COND_NE:
+    return C0 != C1;
+  case RISCVCC::COND_LT:
+    return C0 < C1;
+  case RISCVCC::COND_GE:
+    return C0 >= C1;
+  case RISCVCC::COND_LTU:
+    return (uint64_t)C0 < (uint64_t)C1;
+  case RISCVCC::COND_GEU:
+    return (uint64_t)C0 >= (uint64_t)C1;
+  }
+}
+
 // The contents of values added to Cond are not examined outside of
 // RISCVInstrInfo, giving us flexibility in what to push to it. For RISCV, we
 // push BranchOpcode, Reg1, Reg2.
@@ -1295,6 +1314,49 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
   RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(Cond[0].getImm());
   assert(CC != RISCVCC::COND_INVALID);
 
+  auto modifyBranch = [&]() {
+    // Build the new branch and remove the old one.
+    BuildMI(*MBB, MI, MI.getDebugLoc(),
+            getBrCond(static_cast<RISCVCC::CondCode>(Cond[0].getImm())))
+        .add(Cond[1])
+        .add(Cond[2])
+        .addMBB(TBB);
+    MI.eraseFromParent();
+  };
+
+  // Right now we only care about LI (i.e. ADDI x0, imm)
+  auto isLoadImm = [](const MachineInstr *MI, int64_t &Imm) -> bool {
+    if (MI->getOpcode() == RISCV::ADDI && MI->getOperand(1).isReg() &&
+        MI->getOperand(1).getReg() == RISCV::X0) {
+      Imm = MI->getOperand(2).getImm();
+      return true;
+    }
+    return false;
+  };
+  // Either a load from immediate instruction or X0.
+  auto isFromLoadImm = [&](const MachineOperand &Op, int64_t &Imm) -> bool {
+    if (!Op.isReg())
+      return false;
+    Register Reg = Op.getReg();
+    if (Reg == RISCV::X0) {
+      Imm = 0;
+      return true;
+    }
+    return Reg.isVirtual() && isLoadImm(MRI.getVRegDef(Reg), Imm);
+  };
+
+  // Canonicalize conditional branches which can be constant folded into
+  // beqz or bnez.  We can't modify the CFG here.
+  int64_t C0, C1;
+  if (isFromLoadImm(Cond[1], C0) && isFromLoadImm(Cond[2], C1)) {
+    unsigned NewCC =
+        evaluateCondBranch(CC, C0, C1) ? RISCVCC::COND_EQ : RISCVCC::COND_NE;
+    Cond[0] = MachineOperand::CreateImm(NewCC);
+    Cond[1] = Cond[2] = MachineOperand::CreateReg(RISCV::X0, /*isDef=*/false);
+    modifyBranch();
+    return true;
+  }
+
   if (CC == RISCVCC::COND_EQ || CC == RISCVCC::COND_NE)
     return false;
 
@@ -1315,24 +1377,6 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
   //
   // To make sure this optimization is really beneficial, we only
   // optimize for cases where Y had only one use (i.e. only used by the branch).
-
-  // Right now we only care about LI (i.e. ADDI x0, imm)
-  auto isLoadImm = [](const MachineInstr *MI, int64_t &Imm) -> bool {
-    if (MI->getOpcode() == RISCV::ADDI && MI->getOperand(1).isReg() &&
-        MI->getOperand(1).getReg() == RISCV::X0) {
-      Imm = MI->getOperand(2).getImm();
-      return true;
-    }
-    return false;
-  };
-  // Either a load from immediate instruction or X0.
-  auto isFromLoadImm = [&](const MachineOperand &Op, int64_t &Imm) -> bool {
-    if (!Op.isReg())
-      return false;
-    Register Reg = Op.getReg();
-    return Reg.isVirtual() && isLoadImm(MRI.getVRegDef(Reg), Imm);
-  };
-
   MachineOperand &LHS = MI.getOperand(0);
   MachineOperand &RHS = MI.getOperand(1);
   // Try to find the register for constant Z; return
@@ -1350,8 +1394,6 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
     return Register();
   };
 
-  bool Modify = false;
-  int64_t C0;
   if (isFromLoadImm(LHS, C0) && MRI.hasOneUse(LHS.getReg())) {
     // Might be case 1.
     // Signed integer overflow is UB. (UINT64_MAX is bigger so we don't need
@@ -1364,7 +1406,8 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
         // We might extend the live range of Z, clear its kill flag to
         // account for this.
         MRI.clearKillFlags(RegZ);
-        Modify = true;
+        modifyBranch();
+        return true;
       }
   } else if (isFromLoadImm(RHS, C0) && MRI.hasOneUse(RHS.getReg())) {
     // Might be case 2.
@@ -1378,22 +1421,12 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
         // We might extend the live range of Z, clear its kill flag to
         // account for this.
         MRI.clearKillFlags(RegZ);
-        Modify = true;
+        modifyBranch();
+        return true;
       }
   }
 
-  if (!Modify)
-    return false;
-
-  // Build the new branch and remove the old one.
-  BuildMI(*MBB, MI, MI.getDebugLoc(),
-          getBrCond(static_cast<RISCVCC::CondCode>(Cond[0].getImm())))
-      .add(Cond[1])
-      .add(Cond[2])
-      .addMBB(TBB);
-  MI.eraseFromParent();
-
-  return true;
+  return false;
 }
 
 MachineBasicBlock *

llvm/test/CodeGen/RISCV/bfloat-br-fcmp.ll

@@ -11,8 +11,7 @@ declare bfloat @dummy(bfloat)
 define void @br_fcmp_false(bfloat %a, bfloat %b) nounwind {
 ; RV32IZFBFMIN-LABEL: br_fcmp_false:
 ; RV32IZFBFMIN:       # %bb.0:
-; RV32IZFBFMIN-NEXT:    li a0, 1
-; RV32IZFBFMIN-NEXT:    bnez a0, .LBB0_2
+; RV32IZFBFMIN-NEXT:    beqz zero, .LBB0_2
 ; RV32IZFBFMIN-NEXT:  # %bb.1: # %if.then
 ; RV32IZFBFMIN-NEXT:    ret
 ; RV32IZFBFMIN-NEXT:  .LBB0_2: # %if.else
@@ -22,8 +21,7 @@ define void @br_fcmp_false(bfloat %a, bfloat %b) nounwind {
 ;
 ; RV64IZFBFMIN-LABEL: br_fcmp_false:
 ; RV64IZFBFMIN:       # %bb.0:
-; RV64IZFBFMIN-NEXT:    li a0, 1
-; RV64IZFBFMIN-NEXT:    bnez a0, .LBB0_2
+; RV64IZFBFMIN-NEXT:    beqz zero, .LBB0_2
 ; RV64IZFBFMIN-NEXT:  # %bb.1: # %if.then
 ; RV64IZFBFMIN-NEXT:    ret
 ; RV64IZFBFMIN-NEXT:  .LBB0_2: # %if.else
@@ -583,8 +581,7 @@ if.then:
 define void @br_fcmp_true(bfloat %a, bfloat %b) nounwind {
 ; RV32IZFBFMIN-LABEL: br_fcmp_true:
 ; RV32IZFBFMIN:       # %bb.0:
-; RV32IZFBFMIN-NEXT:    li a0, 1
-; RV32IZFBFMIN-NEXT:    bnez a0, .LBB16_2
+; RV32IZFBFMIN-NEXT:    beqz zero, .LBB16_2
 ; RV32IZFBFMIN-NEXT:  # %bb.1: # %if.else
 ; RV32IZFBFMIN-NEXT:    ret
 ; RV32IZFBFMIN-NEXT:  .LBB16_2: # %if.then
@@ -594,8 +591,7 @@ define void @br_fcmp_true(bfloat %a, bfloat %b) nounwind {
 ;
 ; RV64IZFBFMIN-LABEL: br_fcmp_true:
 ; RV64IZFBFMIN:       # %bb.0:
-; RV64IZFBFMIN-NEXT:    li a0, 1
-; RV64IZFBFMIN-NEXT:    bnez a0, .LBB16_2
+; RV64IZFBFMIN-NEXT:    beqz zero, .LBB16_2
 ; RV64IZFBFMIN-NEXT:  # %bb.1: # %if.else
 ; RV64IZFBFMIN-NEXT:    ret
 ; RV64IZFBFMIN-NEXT:  .LBB16_2: # %if.then

llvm/test/CodeGen/RISCV/double-br-fcmp.ll

@@ -14,8 +14,7 @@ declare void @exit(i32)
 define void @br_fcmp_false(double %a, double %b) nounwind {
 ; RV32IFD-LABEL: br_fcmp_false:
 ; RV32IFD:       # %bb.0:
-; RV32IFD-NEXT:    li a0, 1
-; RV32IFD-NEXT:    bnez a0, .LBB0_2
+; RV32IFD-NEXT:    beqz zero, .LBB0_2
 ; RV32IFD-NEXT:  # %bb.1: # %if.then
 ; RV32IFD-NEXT:    ret
 ; RV32IFD-NEXT:  .LBB0_2: # %if.else
@@ -25,8 +24,7 @@ define void @br_fcmp_false(double %a, double %b) nounwind {
 ;
 ; RV64IFD-LABEL: br_fcmp_false:
 ; RV64IFD:       # %bb.0:
-; RV64IFD-NEXT:    li a0, 1
-; RV64IFD-NEXT:    bnez a0, .LBB0_2
+; RV64IFD-NEXT:    beqz zero, .LBB0_2
 ; RV64IFD-NEXT:  # %bb.1: # %if.then
 ; RV64IFD-NEXT:    ret
 ; RV64IFD-NEXT:  .LBB0_2: # %if.else
@@ -36,8 +34,7 @@ define void @br_fcmp_false(double %a, double %b) nounwind {
 ;
 ; RV32IZFINXZDINX-LABEL: br_fcmp_false:
 ; RV32IZFINXZDINX:       # %bb.0:
-; RV32IZFINXZDINX-NEXT:    li a0, 1
-; RV32IZFINXZDINX-NEXT:    bnez a0, .LBB0_2
+; RV32IZFINXZDINX-NEXT:    beqz zero, .LBB0_2
 ; RV32IZFINXZDINX-NEXT:  # %bb.1: # %if.then
 ; RV32IZFINXZDINX-NEXT:    ret
 ; RV32IZFINXZDINX-NEXT:  .LBB0_2: # %if.else
@@ -47,8 +44,7 @@ define void @br_fcmp_false(double %a, double %b) nounwind {
 ;
 ; RV64IZFINXZDINX-LABEL: br_fcmp_false:
 ; RV64IZFINXZDINX:       # %bb.0:
-; RV64IZFINXZDINX-NEXT:    li a0, 1
-; RV64IZFINXZDINX-NEXT:    bnez a0, .LBB0_2
+; RV64IZFINXZDINX-NEXT:    beqz zero, .LBB0_2
 ; RV64IZFINXZDINX-NEXT:  # %bb.1: # %if.then
 ; RV64IZFINXZDINX-NEXT:    ret
 ; RV64IZFINXZDINX-NEXT:  .LBB0_2: # %if.else
@@ -897,8 +893,7 @@ if.then:
 define void @br_fcmp_true(double %a, double %b) nounwind {
 ; RV32IFD-LABEL: br_fcmp_true:
 ; RV32IFD:       # %bb.0:
-; RV32IFD-NEXT:    li a0, 1
-; RV32IFD-NEXT:    bnez a0, .LBB16_2
+; RV32IFD-NEXT:    beqz zero, .LBB16_2
 ; RV32IFD-NEXT:  # %bb.1: # %if.else
 ; RV32IFD-NEXT:    ret
 ; RV32IFD-NEXT:  .LBB16_2: # %if.then
@@ -908,8 +903,7 @@ define void @br_fcmp_true(double %a, double %b) nounwind {
 ;
 ; RV64IFD-LABEL: br_fcmp_true:
 ; RV64IFD:       # %bb.0:
-; RV64IFD-NEXT:    li a0, 1
-; RV64IFD-NEXT:    bnez a0, .LBB16_2
+; RV64IFD-NEXT:    beqz zero, .LBB16_2
 ; RV64IFD-NEXT:  # %bb.1: # %if.else
 ; RV64IFD-NEXT:    ret
 ; RV64IFD-NEXT:  .LBB16_2: # %if.then
@@ -919,8 +913,7 @@ define void @br_fcmp_true(double %a, double %b) nounwind {
 ;
 ; RV32IZFINXZDINX-LABEL: br_fcmp_true:
 ; RV32IZFINXZDINX:       # %bb.0:
-; RV32IZFINXZDINX-NEXT:    li a0, 1
-; RV32IZFINXZDINX-NEXT:    bnez a0, .LBB16_2
+; RV32IZFINXZDINX-NEXT:    beqz zero, .LBB16_2
 ; RV32IZFINXZDINX-NEXT:  # %bb.1: # %if.else
 ; RV32IZFINXZDINX-NEXT:    ret
 ; RV32IZFINXZDINX-NEXT:  .LBB16_2: # %if.then
@@ -930,8 +923,7 @@ define void @br_fcmp_true(double %a, double %b) nounwind {
 ;
 ; RV64IZFINXZDINX-LABEL: br_fcmp_true:
 ; RV64IZFINXZDINX:       # %bb.0:
-; RV64IZFINXZDINX-NEXT:    li a0, 1
-; RV64IZFINXZDINX-NEXT:    bnez a0, .LBB16_2
+; RV64IZFINXZDINX-NEXT:    beqz zero, .LBB16_2
 ; RV64IZFINXZDINX-NEXT:  # %bb.1: # %if.else
 ; RV64IZFINXZDINX-NEXT:    ret
 ; RV64IZFINXZDINX-NEXT:  .LBB16_2: # %if.then

llvm/test/CodeGen/RISCV/float-br-fcmp.ll

@@ -15,8 +15,7 @@ declare float @dummy(float)
 define void @br_fcmp_false(float %a, float %b) nounwind {
 ; RV32IF-LABEL: br_fcmp_false:
 ; RV32IF:       # %bb.0:
-; RV32IF-NEXT:    li a0, 1
-; RV32IF-NEXT:    bnez a0, .LBB0_2
+; RV32IF-NEXT:    beqz zero, .LBB0_2
 ; RV32IF-NEXT:  # %bb.1: # %if.then
 ; RV32IF-NEXT:    ret
 ; RV32IF-NEXT:  .LBB0_2: # %if.else
@@ -26,8 +25,7 @@ define void @br_fcmp_false(float %a, float %b) nounwind {
 ;
 ; RV64IF-LABEL: br_fcmp_false:
 ; RV64IF:       # %bb.0:
-; RV64IF-NEXT:    li a0, 1
-; RV64IF-NEXT:    bnez a0, .LBB0_2
+; RV64IF-NEXT:    beqz zero, .LBB0_2
 ; RV64IF-NEXT:  # %bb.1: # %if.then
 ; RV64IF-NEXT:    ret
 ; RV64IF-NEXT:  .LBB0_2: # %if.else
@@ -37,8 +35,7 @@ define void @br_fcmp_false(float %a, float %b) nounwind {
 ;
 ; RV32IZFINX-LABEL: br_fcmp_false:
 ; RV32IZFINX:       # %bb.0:
-; RV32IZFINX-NEXT:    li a0, 1
-; RV32IZFINX-NEXT:    bnez a0, .LBB0_2
+; RV32IZFINX-NEXT:    beqz zero, .LBB0_2
 ; RV32IZFINX-NEXT:  # %bb.1: # %if.then
 ; RV32IZFINX-NEXT:    ret
 ; RV32IZFINX-NEXT:  .LBB0_2: # %if.else
@@ -48,8 +45,7 @@ define void @br_fcmp_false(float %a, float %b) nounwind {
 ;
 ; RV64IZFINX-LABEL: br_fcmp_false:
 ; RV64IZFINX:       # %bb.0:
-; RV64IZFINX-NEXT:    li a0, 1
-; RV64IZFINX-NEXT:    bnez a0, .LBB0_2
+; RV64IZFINX-NEXT:    beqz zero, .LBB0_2
 ; RV64IZFINX-NEXT:  # %bb.1: # %if.then
 ; RV64IZFINX-NEXT:    ret
 ; RV64IZFINX-NEXT:  .LBB0_2: # %if.else
@@ -898,8 +894,7 @@ if.then:
 define void @br_fcmp_true(float %a, float %b) nounwind {
 ; RV32IF-LABEL: br_fcmp_true:
 ; RV32IF:       # %bb.0:
-; RV32IF-NEXT:    li a0, 1
-; RV32IF-NEXT:    bnez a0, .LBB16_2
+; RV32IF-NEXT:    beqz zero, .LBB16_2
 ; RV32IF-NEXT:  # %bb.1: # %if.else
 ; RV32IF-NEXT:    ret
 ; RV32IF-NEXT:  .LBB16_2: # %if.then
@@ -909,8 +904,7 @@ define void @br_fcmp_true(float %a, float %b) nounwind {
 ;
 ; RV64IF-LABEL: br_fcmp_true:
 ; RV64IF:       # %bb.0:
-; RV64IF-NEXT:    li a0, 1
-; RV64IF-NEXT:    bnez a0, .LBB16_2
+; RV64IF-NEXT:    beqz zero, .LBB16_2
 ; RV64IF-NEXT:  # %bb.1: # %if.else
 ; RV64IF-NEXT:    ret
 ; RV64IF-NEXT:  .LBB16_2: # %if.then
@@ -920,8 +914,7 @@ define void @br_fcmp_true(float %a, float %b) nounwind {
 ;
 ; RV32IZFINX-LABEL: br_fcmp_true:
 ; RV32IZFINX:       # %bb.0:
-; RV32IZFINX-NEXT:    li a0, 1
-; RV32IZFINX-NEXT:    bnez a0, .LBB16_2
+; RV32IZFINX-NEXT:    beqz zero, .LBB16_2
 ; RV32IZFINX-NEXT:  # %bb.1: # %if.else
 ; RV32IZFINX-NEXT:    ret
 ; RV32IZFINX-NEXT:  .LBB16_2: # %if.then
@@ -931,8 +924,7 @@ define void @br_fcmp_true(float %a, float %b) nounwind {
 ;
 ; RV64IZFINX-LABEL: br_fcmp_true:
 ; RV64IZFINX:       # %bb.0:
-; RV64IZFINX-NEXT:    li a0, 1
-; RV64IZFINX-NEXT:    bnez a0, .LBB16_2
+; RV64IZFINX-NEXT:    beqz zero, .LBB16_2
 ; RV64IZFINX-NEXT:  # %bb.1: # %if.else
 ; RV64IZFINX-NEXT:    ret
 ; RV64IZFINX-NEXT:  .LBB16_2: # %if.then