[RISCV] Add late optimization pass for riscv (#133256)

This patch is an alternative to PRs #117060, #131684, #131728.

The patch adds a late optimization pass that replaces conditional
branches that can be statically evaluated with an unconditinal branch.

Adding Michael as a co-author as most of the code that evaluates the
condition comes from #131684.

Co-authored-by: Michael Maitland [email protected]

Author: mikhailramalho

llvm/lib/Target/RISCV/CMakeLists.txt

@@ -47,6 +47,7 @@ add_llvm_target(RISCVCodeGen
   RISCVISelDAGToDAG.cpp
   RISCVISelLowering.cpp
   RISCVLandingPadSetup.cpp
+  RISCVLateBranchOpt.cpp
   RISCVLoadStoreOptimizer.cpp
   RISCVMachineFunctionInfo.cpp
   RISCVMakeCompressible.cpp

llvm/lib/Target/RISCV/RISCV.h

@@ -40,6 +40,9 @@ void initializeRISCVLandingPadSetupPass(PassRegistry &);
 FunctionPass *createRISCVISelDag(RISCVTargetMachine &TM,
                                  CodeGenOptLevel OptLevel);
 
+FunctionPass *createRISCVLateBranchOptPass();
+void initializeRISCVLateBranchOptPass(PassRegistry &);
+
 FunctionPass *createRISCVMakeCompressibleOptPass();
 void initializeRISCVMakeCompressibleOptPass(PassRegistry &);
 

llvm/lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -993,7 +993,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");
@@ -1297,6 +1297,31 @@ bool RISCVInstrInfo::reverseBranchCondition(
   return false;
 }
 
+// Return true if the instruction is a load immediate instruction (i.e.
+// ADDI x0, imm).
+static bool isLoadImm(const MachineInstr *MI, int64_t &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();
@@ -1319,31 +1344,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);
@@ -1389,7 +1393,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)
@@ -1404,7 +1408,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,15 @@ class RISCVInstrInfo : public RISCVGenInstrInfo {
   static bool isLdStSafeToPair(const MachineInstr &LdSt,
                                const TargetRegisterInfo *TRI);
 
+  /// Return the result of the evaluation of C0 CC C1, where CC is a
+  /// RISCVCC::CondCode.
+  static bool evaluateCondBranch(unsigned CC, int64_t C0, int64_t C1);
+
+  /// Return true if the operand is a load immediate instruction and
+  /// sets Imm to the immediate value.
+  static bool isFromLoadImm(const MachineRegisterInfo &MRI,
+                            const MachineOperand &Op, int64_t &Imm);
+
 protected:
   const RISCVSubtarget &STI;
 

llvm/lib/Target/RISCV/RISCVLateBranchOpt.cpp

@@ -0,0 +1,111 @@
+//===-- RISCVLateBranchOpt.cpp - Late Stage Branch Optimization -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// 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 "RISCVInstrInfo.h"
+#include "RISCVSubtarget.h"
+
+using namespace llvm;
+
+#define RISCV_LATE_BRANCH_OPT_NAME "RISC-V Late Branch Optimisation Pass"
+
+namespace {
+
+struct RISCVLateBranchOpt : public MachineFunctionPass {
+  static char ID;
+
+  RISCVLateBranchOpt() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override { return RISCV_LATE_BRANCH_OPT_NAME; }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &Fn) override;
+
+private:
+  bool runOnBasicBlock(MachineBasicBlock &MBB) const;
+
+  const RISCVInstrInfo *RII = nullptr;
+};
+} // namespace
+
+char RISCVLateBranchOpt::ID = 0;
+INITIALIZE_PASS(RISCVLateBranchOpt, "riscv-late-branch-opt",
+                RISCV_LATE_BRANCH_OPT_NAME, false, false)
+
+bool RISCVLateBranchOpt::runOnBasicBlock(MachineBasicBlock &MBB) const {
+  MachineBasicBlock *TBB, *FBB;
+  SmallVector<MachineOperand, 4> Cond;
+  if (RII->analyzeBranch(MBB, TBB, FBB, Cond, /*AllowModify=*/false))
+    return false;
+
+  if (!TBB || Cond.size() != 3)
+    return false;
+
+  RISCVCC::CondCode CC = static_cast<RISCVCC::CondCode>(Cond[0].getImm());
+  assert(CC != RISCVCC::COND_INVALID);
+
+  MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
+
+  // Try and convert a conditional branch that can be evaluated statically
+  // into an unconditional branch.
+  int64_t C0, C1;
+  if (!RISCVInstrInfo::isFromLoadImm(MRI, Cond[1], C0) ||
+      !RISCVInstrInfo::isFromLoadImm(MRI, Cond[2], C1))
+    return false;
+
+  MachineBasicBlock *Folded =
+      RISCVInstrInfo::evaluateCondBranch(CC, C0, C1) ? TBB : FBB;
+
+  // At this point, its legal to optimize.
+  RII->removeBranch(MBB);
+
+  // Only need to insert a branch if we're not falling through.
+  if (Folded) {
+    DebugLoc DL = MBB.findBranchDebugLoc();
+    RII->insertBranch(MBB, Folded, nullptr, {}, DL);
+  }
+
+  // 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);
+
+  return true;
+}
+
+bool RISCVLateBranchOpt::runOnMachineFunction(MachineFunction &Fn) {
+  if (skipFunction(Fn.getFunction()))
+    return false;
+
+  auto &ST = Fn.getSubtarget<RISCVSubtarget>();
+  RII = ST.getInstrInfo();
+
+  bool Changed = false;
+  for (MachineBasicBlock &MBB : Fn)
+    Changed |= runOnBasicBlock(MBB);
+  return Changed;
+}
+
+FunctionPass *llvm::createRISCVLateBranchOptPass() {
+  return new RISCVLateBranchOpt();
+}

llvm/lib/Target/RISCV/RISCVTargetMachine.cpp

@@ -127,6 +127,7 @@ extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeRISCVTarget() {
   initializeRISCVPostLegalizerCombinerPass(*PR);
   initializeKCFIPass(*PR);
   initializeRISCVDeadRegisterDefinitionsPass(*PR);
+  initializeRISCVLateBranchOptPass(*PR);
   initializeRISCVMakeCompressibleOptPass(*PR);
   initializeRISCVGatherScatterLoweringPass(*PR);
   initializeRISCVCodeGenPreparePass(*PR);
@@ -565,6 +566,8 @@ void RISCVPassConfig::addPreEmitPass() {
   if (TM->getOptLevel() >= CodeGenOptLevel::Default &&
       EnableRISCVCopyPropagation)
     addPass(createMachineCopyPropagationPass(true));
+  if (TM->getOptLevel() >= CodeGenOptLevel::Default)
+    addPass(createRISCVLateBranchOptPass());
   addPass(&BranchRelaxationPassID);
   addPass(createRISCVMakeCompressibleOptPass());
 }

llvm/test/CodeGen/RISCV/GlobalISel/rv32zbb.ll

@@ -357,7 +357,7 @@ 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:    j .LBB6_2
 ; RV32I-NEXT:  # %bb.1:
 ; RV32I-NEXT:    sltiu a0, zero, 0
 ; RV32I-NEXT:    ret
@@ -404,7 +404,7 @@ 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:    j .LBB6_2
 ; RV32ZBB-NEXT:  # %bb.1:
 ; RV32ZBB-NEXT:    sltiu a0, zero, 0
 ; RV32ZBB-NEXT:    ret
@@ -422,7 +422,7 @@ 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:    j .LBB7_2
 ; RV32I-NEXT:  # %bb.1:
 ; RV32I-NEXT:    snez a0, zero
 ; RV32I-NEXT:    ret
@@ -470,7 +470,7 @@ 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:    j .LBB7_2
 ; RV32ZBB-NEXT:  # %bb.1:
 ; RV32ZBB-NEXT:    snez a0, zero
 ; RV32ZBB-NEXT:    ret

llvm/test/CodeGen/RISCV/O3-pipeline.ll

@@ -194,6 +194,7 @@
 ; CHECK-NEXT:       Insert XRay ops
 ; CHECK-NEXT:       Implement the 'patchable-function' attribute
 ; CHECK-NEXT:       Machine Copy Propagation Pass
+; CHECK-NEXT:       RISC-V Late Branch Optimisation Pass
 ; CHECK-NEXT:       Branch relaxation pass
 ; CHECK-NEXT:       RISC-V Make Compressible
 ; CHECK-NEXT:       Contiguously Lay Out Funclets

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

@@ -11,7 +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:    beqz zero, .LBB0_2
+; RV32IZFBFMIN-NEXT:    j .LBB0_2
 ; RV32IZFBFMIN-NEXT:  # %bb.1: # %if.then
 ; RV32IZFBFMIN-NEXT:    ret
 ; RV32IZFBFMIN-NEXT:  .LBB0_2: # %if.else
@@ -21,7 +21,7 @@ define void @br_fcmp_false(bfloat %a, bfloat %b) nounwind {
 ;
 ; RV64IZFBFMIN-LABEL: br_fcmp_false:
 ; RV64IZFBFMIN:       # %bb.0:
-; RV64IZFBFMIN-NEXT:    beqz zero, .LBB0_2
+; RV64IZFBFMIN-NEXT:    j .LBB0_2
 ; RV64IZFBFMIN-NEXT:  # %bb.1: # %if.then
 ; RV64IZFBFMIN-NEXT:    ret
 ; RV64IZFBFMIN-NEXT:  .LBB0_2: # %if.else
@@ -581,7 +581,7 @@ if.then:
 define void @br_fcmp_true(bfloat %a, bfloat %b) nounwind {
 ; RV32IZFBFMIN-LABEL: br_fcmp_true:
 ; RV32IZFBFMIN:       # %bb.0:
-; RV32IZFBFMIN-NEXT:    beqz zero, .LBB16_2
+; RV32IZFBFMIN-NEXT:    j .LBB16_2
 ; RV32IZFBFMIN-NEXT:  # %bb.1: # %if.else
 ; RV32IZFBFMIN-NEXT:    ret
 ; RV32IZFBFMIN-NEXT:  .LBB16_2: # %if.then
@@ -591,7 +591,7 @@ define void @br_fcmp_true(bfloat %a, bfloat %b) nounwind {
 ;
 ; RV64IZFBFMIN-LABEL: br_fcmp_true:
 ; RV64IZFBFMIN:       # %bb.0:
-; RV64IZFBFMIN-NEXT:    beqz zero, .LBB16_2
+; RV64IZFBFMIN-NEXT:    j .LBB16_2
 ; RV64IZFBFMIN-NEXT:  # %bb.1: # %if.else
 ; RV64IZFBFMIN-NEXT:    ret
 ; RV64IZFBFMIN-NEXT:  .LBB16_2: # %if.then

llvm/test/CodeGen/RISCV/branch_zero.ll

@@ -5,14 +5,11 @@
 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:  # %bb.1: # %for.body
 ; 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
 entry:
@@ -46,14 +43,11 @@ 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:  # %bb.1: # %for.body
 ; 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
 entry: