1//===------- RISCVPushPopOptimizer.cpp - RISC-V Push/Pop opt. pass --------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This file contains a pass that replaces Zcmp POP instructions with
10// POPRET[Z] where possible.
11//
12//===----------------------------------------------------------------------===//
13
14#include "RISCVInstrInfo.h"
15#include "RISCVMachineFunctionInfo.h"
16#include "llvm/CodeGen/MachineInstr.h"
17
18using namespace llvm;
19
20#define RISCV_PUSH_POP_OPT_NAME "RISC-V Zcmp Push/Pop optimization pass"
21
22namespace {
23struct RISCVPushPopOpt : public MachineFunctionPass {
24 static char ID;
25
26 RISCVPushPopOpt() : MachineFunctionPass(ID) {}
27
28 const RISCVInstrInfo *TII;
29 const TargetRegisterInfo *TRI;
30
31 // Track which register units have been modified and used.
32 LiveRegUnits ModifiedRegUnits, UsedRegUnits;
33
34 bool usePopRet(MachineBasicBlock::iterator &MBBI,
35 MachineBasicBlock::iterator &NextI, bool IsReturnZero);
36 bool adjustRetVal(MachineBasicBlock::iterator &MBBI);
37 bool runOnMachineFunction(MachineFunction &Fn) override;
38
39 StringRef getPassName() const override { return RISCV_PUSH_POP_OPT_NAME; }
40};
41
42char RISCVPushPopOpt::ID = 0;
43
44} // end of anonymous namespace
45
46INITIALIZE_PASS(RISCVPushPopOpt, "riscv-push-pop-opt", RISCV_PUSH_POP_OPT_NAME,
47 false, false)
48
49static bool isPop(unsigned Opcode) {
50 switch (Opcode) {
51 case RISCV::CM_POP:
52 case RISCV::QC_CM_POP:
53 return true;
54 default:
55 return false;
56 }
57}
58
59static unsigned getPopRetOpcode(unsigned PopOpcode, bool IsReturnZero) {
60 assert(isPop(PopOpcode) && "Unexpected Pop Opcode");
61
62 switch (PopOpcode) {
63 case RISCV::CM_POP:
64 return IsReturnZero ? RISCV::CM_POPRETZ : RISCV::CM_POPRET;
65 case RISCV::QC_CM_POP:
66 return IsReturnZero ? RISCV::QC_CM_POPRETZ : RISCV::QC_CM_POPRET;
67 default:
68 llvm_unreachable("Unhandled Pop Opcode");
69 }
70}
71
72bool RISCVPushPopOpt::usePopRet(MachineBasicBlock::iterator &MBBI,
73 MachineBasicBlock::iterator &NextI,
74 bool IsReturnZero) {
75 // Since Pseudo instruction lowering happen later in the pipeline,
76 // this will detect all ret instruction.
77 DebugLoc DL = NextI->getDebugLoc();
78 unsigned Opc = getPopRetOpcode(PopOpcode: MBBI->getOpcode(), IsReturnZero);
79 MachineInstrBuilder PopRetBuilder =
80 BuildMI(BB&: *NextI->getParent(), I: NextI, MIMD: DL, MCID: TII->get(Opcode: Opc))
81 .add(MO: MBBI->getOperand(i: 0))
82 .add(MO: MBBI->getOperand(i: 1))
83 .setMIFlag(MachineInstr::FrameDestroy);
84
85 // Copy over the variable implicit uses and defs from the CM_POP. They depend
86 // on what register list has been picked during frame lowering.
87 const MCInstrDesc &PopDesc = MBBI->getDesc();
88 unsigned FirstNonDeclaredOp = PopDesc.getNumOperands() +
89 PopDesc.NumImplicitUses +
90 PopDesc.NumImplicitDefs;
91 for (unsigned i = FirstNonDeclaredOp; i < MBBI->getNumOperands(); ++i)
92 PopRetBuilder.add(MO: MBBI->getOperand(i));
93
94 MBBI->eraseFromParent();
95 NextI->eraseFromParent();
96 return true;
97}
98
99// Search for last assignment to a0 and if possible use ret_val slot of POP to
100// store return value.
101bool RISCVPushPopOpt::adjustRetVal(MachineBasicBlock::iterator &MBBI) {
102 MachineBasicBlock::reverse_iterator RE = MBBI->getParent()->rend();
103 // Track which register units have been modified and used between the POP
104 // insn and the last assignment to register a0.
105 ModifiedRegUnits.clear();
106 UsedRegUnits.clear();
107 // Since POP instruction is in Epilogue no normal instructions will follow
108 // after it. Therefore search only previous ones to find the return value.
109 for (MachineBasicBlock::reverse_iterator I =
110 next_nodbg(It: MBBI.getReverse(), End: RE);
111 I != RE; I = next_nodbg(It: I, End: RE)) {
112 MachineInstr &MI = *I;
113 if (auto OperandPair = TII->isCopyInstrImpl(MI)) {
114 Register DestReg = OperandPair->Destination->getReg();
115 Register Source = OperandPair->Source->getReg();
116 if (DestReg == RISCV::X10 && Source == RISCV::X0) {
117 MI.removeFromParent();
118 return true;
119 }
120 }
121 // Update modified / used register units.
122 LiveRegUnits::accumulateUsedDefed(MI, ModifiedRegUnits, UsedRegUnits, TRI);
123 // If a0 was modified or used, there is no possibility
124 // of using ret_val slot of POP instruction.
125 if (!ModifiedRegUnits.available(Reg: RISCV::X10) ||
126 !UsedRegUnits.available(Reg: RISCV::X10))
127 return false;
128 }
129 return false;
130}
131
132bool RISCVPushPopOpt::runOnMachineFunction(MachineFunction &Fn) {
133 if (skipFunction(F: Fn.getFunction()))
134 return false;
135
136 // If Zcmp extension is not supported, abort.
137 const RISCVSubtarget *Subtarget = &Fn.getSubtarget<RISCVSubtarget>();
138 if (!Subtarget->hasStdExtZcmp() && !Subtarget->hasVendorXqccmp())
139 return false;
140
141 // We don't want any popret[z] instructions when emitting code with shadow
142 // stack protection. Note that this pass would actually fail to insert any
143 // popret[z] instructions in this case since the cm.pop and ret will not be
144 // adjacent. But there's no point in running a pass that won't do anything.
145 if ((Fn.getFunction().hasFnAttribute(Kind: "hw-shadow-stack") &&
146 Subtarget->hasStdExtZimop()) ||
147 Fn.getFunction().hasFnAttribute(Kind: Attribute::ShadowCallStack))
148 return false;
149
150 TII = Subtarget->getInstrInfo();
151 TRI = Subtarget->getRegisterInfo();
152
153 // Resize the modified and used register unit trackers. We do this once
154 // per function and then clear the register units each time we determine
155 // correct return value for the POP.
156 ModifiedRegUnits.init(TRI: *TRI);
157 UsedRegUnits.init(TRI: *TRI);
158
159 bool Modified = false;
160 for (auto &MBB : Fn) {
161 // RET should be the only terminator.
162 auto RetMBBI = MBB.getFirstTerminator();
163 if (RetMBBI == MBB.end() || RetMBBI->getOpcode() != RISCV::PseudoRET ||
164 RetMBBI == MBB.begin())
165 continue;
166
167 // The previous instruction should be a POP.
168 auto PopMBBI = prev_nodbg(It: RetMBBI, Begin: MBB.begin());
169 if (isPop(Opcode: PopMBBI->getOpcode()) &&
170 PopMBBI->getFlag(Flag: MachineInstr::FrameDestroy))
171 Modified |= usePopRet(MBBI&: PopMBBI, NextI&: RetMBBI, IsReturnZero: adjustRetVal(MBBI&: PopMBBI));
172 }
173
174 return Modified;
175}
176
177/// createRISCVPushPopOptimizationPass - returns an instance of the
178/// Push/Pop optimization pass.
179FunctionPass *llvm::createRISCVPushPopOptimizationPass() {
180 return new RISCVPushPopOpt();
181}
182