[RISCV] Enable the machine outliner for RISC-V

This patch enables the machine outliner for RISC-V and adds the
necessary logic for checking whether sequences can be safely outlined,
and describing how they should be outlined. Outlined functions are
called using the register t0 (x5) as the return address register, which
must be available for an occurrence of a sequence to be safely outlined.

Differential Revision: https://reviews.llvm.org/D66210

lib/Target/RISCV/RISCVInstrInfo.cpp

@@ -637,3 +637,161 @@ RISCVInstrInfo::getSerializableDirectMachineOperandTargetFlags() const {
       {MO_TLS_GD_HI, "riscv-tls-gd-hi"}};
   return makeArrayRef(TargetFlags);
 }
+bool RISCVInstrInfo::isFunctionSafeToOutlineFrom(
+    MachineFunction &MF, bool OutlineFromLinkOnceODRs) const {
+  const Function &F = MF.getFunction();
+
+  // Can F be deduplicated by the linker? If it can, don't outline from it.
+  if (!OutlineFromLinkOnceODRs && F.hasLinkOnceODRLinkage())
+    return false;
+
+  // Don't outline from functions with section markings; the program could
+  // expect that all the code is in the named section.
+  if (F.hasSection())
+    return false;
+
+  // It's safe to outline from MF.
+  return true;
+}
+
+bool RISCVInstrInfo::isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
+                                            unsigned &Flags) const {
+  // More accurate safety checking is done in getOutliningCandidateInfo.
+  return true;
+}
+
+// Enum values indicating how an outlined call should be constructed.
+enum MachineOutlinerConstructionID {
+  MachineOutlinerDefault
+};
+
+outliner::OutlinedFunction RISCVInstrInfo::getOutliningCandidateInfo(
+    std::vector<outliner::Candidate> &RepeatedSequenceLocs) const {
+
+  // First we need to filter out candidates where the X5 register (IE t0) can't
+  // be used to setup the function call.
+  auto CannotInsertCall = [](outliner::Candidate &C) {
+    const TargetRegisterInfo *TRI = C.getMF()->getSubtarget().getRegisterInfo();
+
+    C.initLRU(*TRI);
+    LiveRegUnits LRU = C.LRU;
+    return !LRU.available(RISCV::X5);
+  };
+
+  RepeatedSequenceLocs.erase(std::remove_if(RepeatedSequenceLocs.begin(),
+                                            RepeatedSequenceLocs.end(),
+                                            CannotInsertCall),
+                             RepeatedSequenceLocs.end());
+
+  // If the sequence doesn't have enough candidates left, then we're done.
+  if (RepeatedSequenceLocs.size() < 2)
+    return outliner::OutlinedFunction();
+
+  unsigned SequenceSize = 0;
+
+  auto I = RepeatedSequenceLocs[0].front();
+  auto E = std::next(RepeatedSequenceLocs[0].back());
+  for (; I != E; ++I)
+    SequenceSize += getInstSizeInBytes(*I);
+
+  // call t0, function = 8 bytes.
+  unsigned CallOverhead = 8;
+  for (auto &C : RepeatedSequenceLocs)
+    C.setCallInfo(MachineOutlinerDefault, CallOverhead);
+
+  // jr t0 = 4 bytes, 2 bytes if compressed instructions are enabled.
+  unsigned FrameOverhead = 4;
+  if (RepeatedSequenceLocs[0].getMF()->getSubtarget()
+          .getFeatureBits()[RISCV::FeatureStdExtC])
+    FrameOverhead = 2;
+
+  return outliner::OutlinedFunction(RepeatedSequenceLocs, SequenceSize,
+                                    FrameOverhead, MachineOutlinerDefault);
+}
+
+outliner::InstrType
+RISCVInstrInfo::getOutliningType(MachineBasicBlock::iterator &MBBI,
+                                 unsigned Flags) const {
+  MachineInstr &MI = *MBBI;
+  MachineBasicBlock *MBB = MI.getParent();
+  const TargetRegisterInfo *TRI =
+      MBB->getParent()->getSubtarget().getRegisterInfo();
+
+  // Positions generally can't safely be outlined.
+  if (MI.isPosition()) {
+    // We can manually strip out CFI instructions later.
+    if (MI.isCFIInstruction())
+      return outliner::InstrType::Invisible;
+
+    return outliner::InstrType::Illegal;
+  }
+
+  // Don't trust the user to write safe inline assembly.
+  if (MI.isInlineAsm())
+    return outliner::InstrType::Illegal;
+
+  // We can't outline branches to other basic blocks.
+  if (MI.isTerminator() && !MBB->succ_empty())
+    return outliner::InstrType::Illegal;
+
+  // We need support for tail calls to outlined functions before return
+  // statements can be allowed.
+  if (MI.isReturn())
+    return outliner::InstrType::Illegal;
+
+  // Don't allow modifying the X5 register which we use for return addresses for
+  // these outlined functions.
+  if (MI.modifiesRegister(RISCV::X5, TRI) ||
+      MI.getDesc().hasImplicitDefOfPhysReg(RISCV::X5))
+    return outliner::InstrType::Illegal;
+
+  // Make sure the operands don't reference something unsafe.
+  for (const auto &MO : MI.operands())
+    if (MO.isMBB() || MO.isBlockAddress() || MO.isCPI())
+      return outliner::InstrType::Illegal;
+
+  // Don't allow instructions which won't be materialized to impact outlining
+  // analysis.
+  if (MI.isMetaInstruction())
+    return outliner::InstrType::Invisible;
+
+  return outliner::InstrType::Legal;
+}
+
+void RISCVInstrInfo::buildOutlinedFrame(
+    MachineBasicBlock &MBB, MachineFunction &MF,
+    const outliner::OutlinedFunction &OF) const {
+
+  // Strip out any CFI instructions
+  bool Changed = true;
+  while (Changed) {
+    Changed = false;
+    auto I = MBB.begin();
+    auto E = MBB.end();
+    for (; I != E; ++I) {
+      if (I->isCFIInstruction()) {
+        I->removeFromParent();
+        Changed = true;
+        break;
+      }
+    }
+  }
+
+  // Add in a return instruction to the end of the outlined frame.
+  MBB.insert(MBB.end(), BuildMI(MF, DebugLoc(), get(RISCV::JALR))
+      .addReg(RISCV::X0, RegState::Define)
+      .addReg(RISCV::X5)
+      .addImm(0));
+}
+
+MachineBasicBlock::iterator RISCVInstrInfo::insertOutlinedCall(
+    Module &M, MachineBasicBlock &MBB, MachineBasicBlock::iterator &It,
+    MachineFunction &MF, const outliner::Candidate &C) const {
+
+  // Add in a call instruction to the outlined function at the given location.
+  It = MBB.insert(It,
+                  BuildMI(MF, DebugLoc(), get(RISCV::PseudoCALLReg), RISCV::X5)
+                      .addGlobalAddress(M.getNamedValue(MF.getName()), 0,
+                                        RISCVII::MO_CALL));
+  return It;
+}

lib/Target/RISCV/RISCVInstrInfo.h

@@ -101,6 +101,35 @@ public:
   ArrayRef<std::pair<unsigned, const char *>>
   getSerializableDirectMachineOperandTargetFlags() const override;
 
+  // Return true if the function can safely be outlined from.
+  virtual bool
+  isFunctionSafeToOutlineFrom(MachineFunction &MF,
+                              bool OutlineFromLinkOnceODRs) const override;
+
+  // Return true if MBB is safe to outline from, and return any target-specific
+  // information in Flags.
+  virtual bool isMBBSafeToOutlineFrom(MachineBasicBlock &MBB,
+                                      unsigned &Flags) const override;
+
+  // Calculate target-specific information for a set of outlining candidates.
+  outliner::OutlinedFunction getOutliningCandidateInfo(
+      std::vector<outliner::Candidate> &RepeatedSequenceLocs) const override;
+
+  // Return if/how a given MachineInstr should be outlined.
+  virtual outliner::InstrType
+  getOutliningType(MachineBasicBlock::iterator &MBBI,
+                   unsigned Flags) const override;
+
+  // Insert a custom frame for outlined functions.
+  virtual void
+  buildOutlinedFrame(MachineBasicBlock &MBB, MachineFunction &MF,
+                     const outliner::OutlinedFunction &OF) const override;
+
+  // Insert a call to an outlined function into a given basic block.
+  virtual MachineBasicBlock::iterator
+  insertOutlinedCall(Module &M, MachineBasicBlock &MBB,
+                     MachineBasicBlock::iterator &It, MachineFunction &MF,
+                     const outliner::Candidate &C) const override;
 protected:
   const RISCVSubtarget &STI;
 };

lib/Target/RISCV/RISCVTargetMachine.cpp

@@ -66,6 +66,9 @@ RISCVTargetMachine::RISCVTargetMachine(const Target &T, const Triple &TT,
                         getEffectiveCodeModel(CM, CodeModel::Small), OL),
       TLOF(std::make_unique<RISCVELFTargetObjectFile>()) {
   initAsmInfo();
+
+  // RISC-V supports the MachineOutliner.
+  setMachineOutliner(true);
 }
 
 const RISCVSubtarget *

test/CodeGen/RISCV/machineoutliner.mir

@@ -0,0 +1,132 @@
+# RUN: llc -march=riscv32 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs < %s \
+# RUN: | FileCheck -check-prefix=RV32I-MO %s
+# RUN: llc -march=riscv64 -x mir -run-pass=machine-outliner -simplify-mir -verify-machineinstrs < %s \
+# RUN: | FileCheck -check-prefix=RV64I-MO %s
+
+--- |
+  define i32 @outline_0(i32 %a, i32 %b) { ret i32 0 }
+
+  define i32 @outline_1(i32 %a, i32 %b) { ret i32 0 }
+
+  define i32 @outline_2(i32 %a, i32 %b) { ret i32 0 }
+
+  ; Should not outline linkonce_odr functions which could be deduplicated by the
+  ; linker.
+  define linkonce_odr i32 @dont_outline_0(i32 %a, i32 %b) { ret i32 0 }
+
+  ; Should not outline functions with named linker sections
+  define i32 @dont_outline_1(i32 %a, i32 %b) section "named" { ret i32 0 }
+
+  ; Cannot outline if the X5 (t0) register is not free
+  define i32 @dont_outline_2(i32 %a, i32 %b) { ret i32 0 }
+
+...
+---
+name:            outline_0
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+    ; RV32I-MO-LABEL: name: outline_0
+    ; RV32I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: outline_0
+    ; RV64I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    PseudoRET implicit $x10
+
+...
+---
+name:            outline_1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+    ; RV32I-MO-LABEL: name: outline_1
+    ; RV32I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: outline_1
+    ; RV64I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    PseudoRET implicit $x10
+
+...
+---
+name:            outline_2
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+    ; RV32I-MO-LABEL: name: outline_2
+    ; RV32I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: outline_2
+    ; RV64I-MO: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    PseudoRET implicit $x10
+
+...
+---
+name:            dont_outline_0
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+    ; RV32I-MO-LABEL: name: dont_outline_0
+    ; RV32I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: dont_outline_0
+    ; RV64I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    PseudoRET implicit $x10
+
+...
+---
+name:            dont_outline_1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11
+    ; RV32I-MO-LABEL: name: dont_outline_1
+    ; RV32I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: dont_outline_1
+    ; RV64I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    PseudoRET implicit $x10
+
+...
+---
+name:            dont_outline_2
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x10, $x11, $x5
+    ; RV32I-MO-LABEL: name: dont_outline_2
+    ; RV32I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    ;
+    ; RV64I-MO-LABEL: name: dont_outline_2
+    ; RV64I-MO-NOT: $x5 = PseudoCALLReg {{.*}} @OUTLINED_FUNCTION_0
+    $x11 = ORI $x11, 1023
+    $x12 = ADDI $x10, 17
+    $x11 = AND $x12, $x11
+    $x10 = SUB $x10, $x11
+    $x10 = ADD $x10, $x5
+    PseudoRET implicit $x10
+
+...