RISCVPushPopOptimizer.cpp

Go to the documentation of this file.

//===------- RISCVPushPopOptimizer.cpp - RISC-V Push/Pop opt. pass --------===//
//
// 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 contains a pass that replaces Zcmp POP instructions with
// POPRET[Z] where possible.
//
//===----------------------------------------------------------------------===//
 
#include "RISCV.h"
#include "RISCVInstrInfo.h"
#include "RISCVMachineFunctionInfo.h"
#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstr.h"
#include "llvm/CodeGen/MachinePassManager.h"
#include "llvm/IR/Analysis.h"
#include "llvm/Pass.h"
 
using namespace llvm;
 
#define RISCV_PUSH_POP_OPT_NAME "RISC-V Zcmp Push/Pop optimization pass"
 
namespace {
class RISCVPushPopOptLegacy : public MachineFunctionPass {
public:
  static char ID;
 
  RISCVPushPopOptLegacy() : MachineFunctionPass(ID) {}
  bool runOnMachineFunction(MachineFunction &MF) override;
 
  StringRef getPassName() const override { return RISCV_PUSH_POP_OPT_NAME; }
 
  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
};
 
class RISCVPushPopOptImpl {
  const RISCVInstrInfo *TII;
  const TargetRegisterInfo *TRI;
 
  // Track which register units have been modified and used.
  LiveRegUnits ModifiedRegUnits, UsedRegUnits;
 
  bool usePopRet(MachineBasicBlock::iterator &MBBI,
                 MachineBasicBlock::iterator &NextI, bool IsReturnZero);
  bool adjustRetVal(MachineBasicBlock::iterator &MBBI);
 
public:
  bool runOnMachineFunction(MachineFunction &MF);
};
 
char RISCVPushPopOptLegacy::ID = 0;
 
} // end of anonymous namespace
 
INITIALIZE_PASS(RISCVPushPopOptLegacy, "riscv-push-pop-opt",
                RISCV_PUSH_POP_OPT_NAME, false, false)
 
static bool isPop(unsigned Opcode) {
  switch (Opcode) {
  case RISCV::CM_POP:
  case RISCV::QC_CM_POP:
    return true;
  default:
    return false;
  }
}
 
static unsigned getPopRetOpcode(unsigned PopOpcode, bool IsReturnZero) {
  assert(isPop(PopOpcode) && "Unexpected Pop Opcode");
 
  switch (PopOpcode) {
  case RISCV::CM_POP:
    return IsReturnZero ? RISCV::CM_POPRETZ : RISCV::CM_POPRET;
  case RISCV::QC_CM_POP:
    return IsReturnZero ? RISCV::QC_CM_POPRETZ : RISCV::QC_CM_POPRET;
  default:
    llvm_unreachable("Unhandled Pop Opcode");
  }
}
 
bool RISCVPushPopOptImpl::usePopRet(MachineBasicBlock::iterator &MBBI,
                                    MachineBasicBlock::iterator &NextI,
                                    bool IsReturnZero) {
  // Since Pseudo instruction lowering happen later in the pipeline,
  // this will detect all ret instruction.
  DebugLoc DL = NextI->getDebugLoc();
  unsigned Opc = getPopRetOpcode(MBBI->getOpcode(), IsReturnZero);
  MachineInstrBuilder PopRetBuilder =
      BuildMI(*NextI->getParent(), NextI, DL, TII->get(Opc))
          .add(MBBI->getOperand(0))
          .add(MBBI->getOperand(1))
          .setMIFlag(MachineInstr::FrameDestroy);
 
  // Copy over the variable implicit uses and defs from the CM_POP. They depend
  // on what register list has been picked during frame lowering.
  const MCInstrDesc &PopDesc = MBBI->getDesc();
  unsigned FirstNonDeclaredOp = PopDesc.getNumOperands() +
                                PopDesc.NumImplicitUses +
                                PopDesc.NumImplicitDefs;
  for (unsigned i = FirstNonDeclaredOp; i < MBBI->getNumOperands(); ++i)
    PopRetBuilder.add(MBBI->getOperand(i));
 
  MBBI->eraseFromParent();
  NextI->eraseFromParent();
  return true;
}
 
// Search for last assignment to a0 and if possible use ret_val slot of POP to
// store return value.
bool RISCVPushPopOptImpl::adjustRetVal(MachineBasicBlock::iterator &MBBI) {
  MachineBasicBlock::reverse_iterator RE = MBBI->getParent()->rend();
  // Track which register units have been modified and used between the POP
  // insn and the last assignment to register a0.
  ModifiedRegUnits.clear();
  UsedRegUnits.clear();
  // Since POP instruction is in Epilogue no normal instructions will follow
  // after it. Therefore search only previous ones to find the return value.
  for (MachineBasicBlock::reverse_iterator I =
           next_nodbg(MBBI.getReverse(), RE);
       I != RE; I = next_nodbg(I, RE)) {
    MachineInstr &MI = *I;
    if (auto OperandPair = TII->isCopyInstrImpl(MI)) {
      Register DestReg = OperandPair->Destination->getReg();
      Register Source = OperandPair->Source->getReg();
      if (DestReg == RISCV::X10 && Source == RISCV::X0) {
        MI.removeFromParent();
        return true;
      }
    }
    // Update modified / used register units.
    LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
    // If a0 was modified or used, there is no possibility
    // of using ret_val slot of POP instruction.
    if (!ModifiedRegUnits.available(RISCV::X10) ||
        !UsedRegUnits.available(RISCV::X10))
      return false;
  }
  return false;
}
 
bool RISCVPushPopOptImpl::runOnMachineFunction(MachineFunction &MF) {
  // If Zcmp extension is not supported, abort.
  const RISCVSubtarget *Subtarget = &MF.getSubtarget<RISCVSubtarget>();
  if (!Subtarget->hasStdExtZcmp() && !Subtarget->hasVendorXqccmp())
    return false;
 
  // We don't want any popret[z] instructions when emitting code with shadow
  // stack protection. Note that this pass would actually fail to insert any
  // popret[z] instructions in this case since the cm.pop and ret will not be
  // adjacent. But there's no point in running a pass that won't do anything.
  if ((MF.getFunction().hasFnAttribute("hw-shadow-stack") &&
       Subtarget->hasStdExtZimop()) ||
      MF.getFunction().hasFnAttribute(Attribute::ShadowCallStack))
    return false;
 
  TII = Subtarget->getInstrInfo();
  TRI = Subtarget->getRegisterInfo();
 
  // Resize the modified and used register unit trackers.  We do this once
  // per function and then clear the register units each time we determine
  // correct return value for the POP.
  ModifiedRegUnits.init(*TRI);
  UsedRegUnits.init(*TRI);
 
  bool Modified = false;
  for (auto &MBB : MF) {
    // RET should be the only terminator.
    auto RetMBBI = MBB.getFirstTerminator();
    if (RetMBBI == MBB.end() || RetMBBI->getOpcode() != RISCV::PseudoRET ||
        RetMBBI == MBB.begin())
      continue;
 
    // The previous instruction should be a POP.
    auto PopMBBI = prev_nodbg(RetMBBI, MBB.begin());
    if (isPop(PopMBBI->getOpcode()) &&
        PopMBBI->getFlag(MachineInstr::FrameDestroy))
      Modified |= usePopRet(PopMBBI, RetMBBI, adjustRetVal(PopMBBI));
  }
 
  return Modified;
}
 
/// createRISCVPushPopOptimizationPass - returns an instance of the
/// Push/Pop optimization pass.
FunctionPass *llvm::createRISCVPushPopOptLegacyPass() {
  return new RISCVPushPopOptLegacy();
}
 
bool RISCVPushPopOptLegacy::runOnMachineFunction(MachineFunction &MF) {
  if (skipFunction(MF.getFunction()))
    return false;
  RISCVPushPopOptImpl Impl;
  return Impl.runOnMachineFunction(MF);
}
 
PreservedAnalyses
RISCVPushPopOptPass::run(MachineFunction &MF,
                         MachineFunctionAnalysisManager &MFAM) {
  RISCVPushPopOptImpl Impl;
  return Impl.runOnMachineFunction(MF)
             ? getMachineFunctionPassPreservedAnalyses()
                   .preserveSet<CFGAnalyses>()
             : PreservedAnalyses::all();
}