[RISCV] Add late optimization pass for riscv

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -997,7 +997,7 @@ static RISCVCC::CondCode getCondFromBranchOpc(unsigned Opc) {
   }
 }
 
-static bool evaluateCondBranch(unsigned CC, int64_t C0, int64_t C1) {
+bool RISCVInstrInfo::evaluateCondBranch(unsigned CC, int64_t C0, int64_t C1) {
   switch (CC) {
   default:
     llvm_unreachable("Unexpected CC");
@@ -1301,6 +1301,30 @@ bool RISCVInstrInfo::reverseBranchCondition(
   return false;
 }
 
+bool RISCVInstrInfo::isLoadImm(const MachineInstr *MI, int64_t &Imm) {
+  // For now, Only checks for LI (i.e. ADDI x0, imm)
+  if (MI->getOpcode() == RISCV::ADDI && MI->getOperand(1).isReg() &&
+      MI->getOperand(1).getReg() == RISCV::X0) {
+    Imm = MI->getOperand(2).getImm();
+    return true;
+  }
+  return false;
+}
+
+bool RISCVInstrInfo::isFromLoadImm(const MachineRegisterInfo &MRI,
+                                   const MachineOperand &Op, int64_t &Imm) {
+  // Either a load from immediate instruction or X0.
+  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);
+}
+
 bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
   MachineBasicBlock *MBB = MI.getParent();
   MachineRegisterInfo &MRI = MBB->getParent()->getRegInfo();
@@ -1323,31 +1347,10 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
     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)) {
+  if (isFromLoadImm(MRI, Cond[1], C0) && isFromLoadImm(MRI, Cond[2], C1)) {
     unsigned NewCC =
         evaluateCondBranch(CC, C0, C1) ? RISCVCC::COND_EQ : RISCVCC::COND_NE;
     Cond[0] = MachineOperand::CreateImm(NewCC);
@@ -1393,7 +1396,7 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
     return Register();
   };
 
-  if (isFromLoadImm(LHS, C0) && MRI.hasOneUse(LHS.getReg())) {
+  if (isFromLoadImm(MRI, LHS, C0) && MRI.hasOneUse(LHS.getReg())) {
     // Might be case 1.
     // Signed integer overflow is UB. (UINT64_MAX is bigger so we don't need
     // to worry about unsigned overflow here)
@@ -1408,7 +1411,7 @@ bool RISCVInstrInfo::optimizeCondBranch(MachineInstr &MI) const {
         modifyBranch();
         return true;
       }
-  } else if (isFromLoadImm(RHS, C0) && MRI.hasOneUse(RHS.getReg())) {
+  } else if (isFromLoadImm(MRI, RHS, C0) && MRI.hasOneUse(RHS.getReg())) {
     // Might be case 2.
     // For unsigned cases, we don't want C1 to wrap back to UINT64_MAX
     // when C0 is zero.

llvm/lib/Target/RISCV/RISCVInstrInfo.h

@@ -307,6 +307,13 @@ class RISCVInstrInfo : public RISCVGenInstrInfo {
   static bool isLdStSafeToPair(const MachineInstr &LdSt,
                                const TargetRegisterInfo *TRI);
 
+  static bool evaluateCondBranch(unsigned CC, int64_t C0, int64_t C1);
+
+  static bool isLoadImm(const MachineInstr *MI, int64_t &Imm);
+
+  static bool isFromLoadImm(const MachineRegisterInfo &MRI,
+                            const MachineOperand &Op, int64_t &Imm);
+
 protected:
   const RISCVSubtarget &STI;
 

llvm/lib/Target/RISCV/RISCVLateOpt.cpp

@@ -6,22 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 ///
-/// This file provides RISC-V specific target descriptions.
+/// This file provides RISC-V specific target optimizations, currently it's
+/// limited to convert conditional branches into unconditional branches when
+/// the condition can be statically evaluated.
 ///
 //===----------------------------------------------------------------------===//
 
-#include "MCTargetDesc/RISCVMCTargetDesc.h"
-#include "RISCV.h"
 #include "RISCVInstrInfo.h"
 #include "RISCVSubtarget.h"
-#include "llvm/CodeGen/MachineBasicBlock.h"
-#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
-#include "llvm/CodeGen/MachineDominators.h"
-#include "llvm/CodeGen/MachineInstrBuilder.h"
-#include "llvm/CodeGen/Passes.h"
-#include "llvm/CodeGen/RegisterScavenging.h"
-#include "llvm/MC/TargetRegistry.h"
-#include "llvm/Support/Debug.h"
 
 using namespace llvm;
 
@@ -44,7 +36,7 @@ struct RISCVLateOpt : public MachineFunctionPass {
   bool runOnMachineFunction(MachineFunction &Fn) override;
 
 private:
-  bool trySimplifyCondBr(MachineInstr *MI, MachineBasicBlock *TBB,
+  bool trySimplifyCondBr(MachineBasicBlock &MBB, MachineBasicBlock *TBB,
                          MachineBasicBlock *FBB,
                          SmallVectorImpl<MachineOperand> &Cond) const;
 
@@ -57,104 +49,52 @@ INITIALIZE_PASS(RISCVLateOpt, "riscv-late-opt", RISCV_LATE_OPT_NAME, false,
                 false)
 
 bool RISCVLateOpt::trySimplifyCondBr(
-    MachineInstr *MI, MachineBasicBlock *TBB, MachineBasicBlock *FBB,
+    MachineBasicBlock &MBB, MachineBasicBlock *TBB, MachineBasicBlock *FBB,
     SmallVectorImpl<MachineOperand> &Cond) const {
 
+  if (!TBB || Cond.size() != 3)
+    return false;
+
   RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(Cond[0].getImm());
   assert(CC != RISCVCC::COND_INVALID);
 
-  // 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;
-  };
-
-  MachineBasicBlock *MBB = MI->getParent();
-  MachineRegisterInfo &MRI = MBB->getParent()->getRegInfo();
-  // 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);
-  };
-
   // Try and convert a conditional branch that can be evaluated statically
   // into an unconditional branch.
   MachineBasicBlock *Folded = nullptr;
   int64_t C0, C1;
-  if (isFromLoadImm(Cond[1], C0) && isFromLoadImm(Cond[2], C1)) {
-    switch (CC) {
-    case RISCVCC::COND_INVALID:
-      llvm_unreachable("Unexpected CC");
-    case RISCVCC::COND_EQ: {
-      Folded = (C0 == C1) ? TBB : FBB;
-      break;
-    }
-    case RISCVCC::COND_NE: {
-      Folded = (C0 != C1) ? TBB : FBB;
-      break;
-    }
-    case RISCVCC::COND_LT: {
-      Folded = (C0 < C1) ? TBB : FBB;
-      break;
-    }
-    case RISCVCC::COND_GE: {
-      Folded = (C0 >= C1) ? TBB : FBB;
-      break;
-    }
-    case RISCVCC::COND_LTU: {
-      Folded = ((uint64_t)C0 < (uint64_t)C1) ? TBB : FBB;
-      break;
-    }
-    case RISCVCC::COND_GEU: {
-      Folded = ((uint64_t)C0 >= (uint64_t)C1) ? TBB : FBB;
-      break;
-    }
-    }
 
-    // Do the conversion
-    // Build the new unconditional branch
-    DebugLoc DL = MBB->findBranchDebugLoc();
+  MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
+  if (RISCVInstrInfo::isFromLoadImm(MRI, Cond[1], C0) &&
+      RISCVInstrInfo::isFromLoadImm(MRI, Cond[2], C1)) {
+    Folded = RISCVInstrInfo::evaluateCondBranch(CC, C0, C1) ? TBB : FBB;
+
+    // At this point, its legal to optimize.
+    RII->removeBranch(MBB);
+    Cond.clear();
+
+    // Only need to insert a branch if we're not falling through.
     if (Folded) {
-      BuildMI(*MBB, MI, DL, RII->get(RISCV::PseudoBR)).addMBB(Folded);
-    } else {
-      MachineFunction::iterator Fallthrough = ++MBB->getIterator();
-      if (Fallthrough == MBB->getParent()->end())
-        return false;
-      BuildMI(*MBB, MI, DL, RII->get(RISCV::PseudoBR)).addMBB(&*Fallthrough);
+      DebugLoc DL = MBB.findBranchDebugLoc();
+      RII->insertBranch(MBB, Folded, nullptr, {}, DL);
     }
 
-    // Update successors of MBB->
-    if (Folded == TBB) {
-      // If we're taking TBB, then the succ to delete is the fallthrough (if
-      // it was a succ in the first place), or its the MBB from the
-      // unconditional branch.
-      if (!FBB) {
-        MachineFunction::iterator Fallthrough = ++MBB->getIterator();
-        if (Fallthrough != MBB->getParent()->end() &&
-            MBB->isSuccessor(&*Fallthrough))
-          MBB->removeSuccessor(&*Fallthrough, true);
-      } else {
-        MBB->removeSuccessor(FBB, true);
-      }
-    } else if (Folded == FBB) {
-      // If we're taking the fallthrough or unconditional branch, then the
-      // succ to remove is the one from the conditional branch.
-      MBB->removeSuccessor(TBB, true);
-    }
+    // Update the successors. Remove them all and add back the correct one.
+    while (!MBB.succ_empty())
+      MBB.removeSuccessor(MBB.succ_end() - 1);
+
+    // If it's a fallthrough, we need to figure out where MBB is going.
+    if (!Folded) {
+      MachineFunction::iterator Fallthrough = ++MBB.getIterator();
+      if (Fallthrough != MBB.getParent()->end())
+        MBB.addSuccessor(&*Fallthrough);
+    } else
+      MBB.addSuccessor(Folded);
 
-    MI->eraseFromParent();
+    TBB = Folded;
+    FBB = nullptr;
     return true;
   }
+
   return false;
 }
 
@@ -168,19 +108,10 @@ bool RISCVLateOpt::runOnMachineFunction(MachineFunction &Fn) {
   bool Changed = false;
 
   for (MachineBasicBlock &MBB : Fn) {
-    for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
-         MII != MIE;) {
-      MachineInstr *MI = &*MII;
-      // We may be erasing MI below, increment MII now.
-      ++MII;
-      if (!MI->isConditionalBranch())
-        continue;
-
-      MachineBasicBlock *TBB, *FBB;
-      SmallVector<MachineOperand, 4> Cond;
-      if (!RII->analyzeBranch(MBB, TBB, FBB, Cond, /*AllowModify=*/false))
-        Changed |= trySimplifyCondBr(MI, TBB, FBB, Cond);
-    }
+    MachineBasicBlock *TBB, *FBB;
+    SmallVector<MachineOperand, 4> Cond;
+    if (!RII->analyzeBranch(MBB, TBB, FBB, Cond, /*AllowModify=*/false))
+      Changed |= trySimplifyCondBr(MBB, TBB, FBB, Cond);
   }
 
   return Changed;

llvm/test/CodeGen/RISCV/branch_zero.ll

@@ -10,8 +10,7 @@ define void @foo(i16 %finder_idx) {
 ; CHECK-NEXT:    bltz a0, .LBB0_4
 ; CHECK-NEXT:  # %bb.2: # %while.cond.preheader.i
 ; CHECK-NEXT:    li a0, 0
-; CHECK-NEXT:    j .LBB0_3
-; CHECK-NEXT:  .LBB0_3: # %while.body
+; CHECK-NEXT:  # %bb.3: # %while.body
 ; CHECK-NEXT:  .LBB0_4: # %while.cond1.preheader.i
 entry:
   br label %for.body
@@ -49,8 +48,7 @@ define void @bar(i16 %finder_idx) {
 ; CHECK-NEXT:    bgez a0, .LBB1_4
 ; CHECK-NEXT:  # %bb.2: # %while.cond.preheader.i
 ; CHECK-NEXT:    li a0, 0
-; CHECK-NEXT:    j .LBB1_3
-; CHECK-NEXT:  .LBB1_3: # %while.body
+; CHECK-NEXT:  # %bb.3: # %while.body
 ; CHECK-NEXT:  .LBB1_4: # %while.cond1.preheader.i
 entry:
   br label %for.body

llvm/test/CodeGen/RISCV/machine-sink-load-immediate.ll

@@ -43,28 +43,22 @@ define i1 @sink_li(ptr %text, ptr %text.addr.0) nounwind {
 ; CHECK-NEXT:    # in Loop: Header=BB0_1 Depth=1
 ; CHECK-NEXT:    j .LBB0_1
 ; CHECK-NEXT:  # %bb.12: # %while.body.6
-; CHECK-NEXT:    j .LBB0_13
-; CHECK-NEXT:  .LBB0_13: # %while.body.6
-; CHECK-NEXT:    j .LBB0_14
-; CHECK-NEXT:  .LBB0_14: # %strdup.exit.split.loop.exit126
+; CHECK-NEXT:  # %bb.13: # %while.body.6
+; CHECK-NEXT:  # %bb.14: # %strdup.exit.split.loop.exit126
 ; CHECK-NEXT:    addi s0, s1, 7
 ; CHECK-NEXT:    j .LBB0_24
 ; CHECK-NEXT:  .LBB0_15: # %while.body
 ; CHECK-NEXT:    j .LBB0_17
-; CHECK-NEXT:    j .LBB0_17
 ; CHECK-NEXT:  .LBB0_16: # %while.body.1
-; CHECK-NEXT:    j .LBB0_17
 ; CHECK-NEXT:  .LBB0_17: # %strdup.exit.loopexit
 ; CHECK-NEXT:    li s0, 0
 ; CHECK-NEXT:    j .LBB0_24
 ; CHECK-NEXT:  .LBB0_18: # %while.body.3
-; CHECK-NEXT:    j .LBB0_19
-; CHECK-NEXT:  .LBB0_19: # %strdup.exit.split.loop.exit120
+; CHECK-NEXT:  # %bb.19: # %strdup.exit.split.loop.exit120
 ; CHECK-NEXT:    addi s0, s1, 4
 ; CHECK-NEXT:    j .LBB0_24
 ; CHECK-NEXT:  .LBB0_20: # %while.body.4
-; CHECK-NEXT:    j .LBB0_21
-; CHECK-NEXT:  .LBB0_21: # %strdup.exit.split.loop.exit122
+; CHECK-NEXT:  # %bb.21: # %strdup.exit.split.loop.exit122
 ; CHECK-NEXT:    addi s0, s1, 5
 ; CHECK-NEXT:    j .LBB0_24
 ; CHECK-NEXT:  .LBB0_22: # %while.body.5

llvm/test/CodeGen/RISCV/rvv/vxrm-insert-out-of-loop.ll

@@ -26,8 +26,7 @@ define void @test1(ptr nocapture noundef writeonly %dst, i32 noundef signext %i_
 ; RV32-NEXT:    slli t1, t2, 1
 ; RV32-NEXT:    li t6, 32
 ; RV32-NEXT:    mv t0, t1
-; RV32-NEXT:    j .LBB0_3
-; RV32-NEXT:  .LBB0_3: # %for.cond1.preheader.us.preheader
+; RV32-NEXT:  # %bb.3: # %for.cond1.preheader.us.preheader
 ; RV32-NEXT:    li t0, 32
 ; RV32-NEXT:  # %bb.4: # %for.cond1.preheader.us.preheader
 ; RV32-NEXT:    addi sp, sp, -16