1//===- SIPreAllocateWWMRegs.cpp - WWM Register Pre-allocation -------------===//
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/// \file
10/// Pass to pre-allocated WWM registers
11//
12//===----------------------------------------------------------------------===//
13
14#include "SIPreAllocateWWMRegs.h"
15#include "AMDGPU.h"
16#include "GCNSubtarget.h"
17#include "MCTargetDesc/AMDGPUMCTargetDesc.h"
18#include "SIMachineFunctionInfo.h"
19#include "llvm/ADT/PostOrderIterator.h"
20#include "llvm/CodeGen/LiveIntervals.h"
21#include "llvm/CodeGen/LiveRegMatrix.h"
22#include "llvm/CodeGen/MachineFrameInfo.h"
23#include "llvm/CodeGen/MachineFunctionPass.h"
24#include "llvm/CodeGen/RegisterClassInfo.h"
25#include "llvm/CodeGen/VirtRegMap.h"
26#include "llvm/InitializePasses.h"
27
28using namespace llvm;
29
30#define DEBUG_TYPE "si-pre-allocate-wwm-regs"
31
32static cl::opt<bool>
33 EnablePreallocateSGPRSpillVGPRs("amdgpu-prealloc-sgpr-spill-vgprs",
34 cl::init(Val: false), cl::Hidden);
35
36bool llvm::isPreallocateSGPRSpillVGPRsEnabled(const MachineFunction &MF) {
37 return EnablePreallocateSGPRSpillVGPRs ||
38 MF.getFunction().hasFnAttribute(Kind: "amdgpu-prealloc-sgpr-spill-vgprs");
39}
40
41namespace {
42
43class SIPreAllocateWWMRegs {
44private:
45 const SIInstrInfo *TII;
46 const SIRegisterInfo *TRI;
47 MachineRegisterInfo *MRI;
48 LiveIntervals *LIS;
49 LiveRegMatrix *Matrix;
50 VirtRegMap *VRM;
51 RegisterClassInfo &RCI;
52
53 std::vector<unsigned> RegsToRewrite;
54#ifndef NDEBUG
55 void printWWMInfo(const MachineInstr &MI);
56#endif
57 bool processDef(MachineOperand &MO);
58 void rewriteRegs(MachineFunction &MF);
59
60public:
61 SIPreAllocateWWMRegs(LiveIntervals *LIS, LiveRegMatrix *Matrix,
62 VirtRegMap *VRM, RegisterClassInfo &RCI)
63 : LIS(LIS), Matrix(Matrix), VRM(VRM), RCI(RCI) {}
64 bool run(MachineFunction &MF);
65};
66
67class SIPreAllocateWWMRegsLegacy : public MachineFunctionPass {
68public:
69 static char ID;
70
71 SIPreAllocateWWMRegsLegacy() : MachineFunctionPass(ID) {}
72
73 bool runOnMachineFunction(MachineFunction &MF) override;
74
75 void getAnalysisUsage(AnalysisUsage &AU) const override {
76 AU.addRequired<LiveIntervalsWrapperPass>();
77 AU.addRequired<VirtRegMapWrapperLegacy>();
78 AU.addRequired<LiveRegMatrixWrapperLegacy>();
79 AU.addRequired<MachineRegisterClassInfoWrapperPass>();
80 AU.setPreservesAll();
81 MachineFunctionPass::getAnalysisUsage(AU);
82 }
83};
84
85} // End anonymous namespace.
86
87INITIALIZE_PASS_BEGIN(SIPreAllocateWWMRegsLegacy, DEBUG_TYPE,
88 "SI Pre-allocate WWM Registers", false, false)
89INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
90INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
91INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
92INITIALIZE_PASS_DEPENDENCY(MachineRegisterClassInfoWrapperPass)
93INITIALIZE_PASS_END(SIPreAllocateWWMRegsLegacy, DEBUG_TYPE,
94 "SI Pre-allocate WWM Registers", false, false)
95
96char SIPreAllocateWWMRegsLegacy::ID = 0;
97
98char &llvm::SIPreAllocateWWMRegsLegacyID = SIPreAllocateWWMRegsLegacy::ID;
99
100FunctionPass *llvm::createSIPreAllocateWWMRegsLegacyPass() {
101 return new SIPreAllocateWWMRegsLegacy();
102}
103
104bool SIPreAllocateWWMRegs::processDef(MachineOperand &MO) {
105 Register Reg = MO.getReg();
106 if (Reg.isPhysical())
107 return false;
108
109 if (!SIRegisterInfo::hasVGPRs(RC: MRI->getRegClass(Reg)))
110 return false;
111
112 if (VRM->hasPhys(virtReg: Reg))
113 return false;
114
115 LiveInterval &LI = LIS->getInterval(Reg);
116
117 for (MCRegister PhysReg : RCI.getOrder(RC: MRI->getRegClass(Reg))) {
118 if (!MRI->isPhysRegUsed(PhysReg, /*SkipRegMaskTest=*/true) &&
119 Matrix->checkInterference(VirtReg: LI, PhysReg) == LiveRegMatrix::IK_Free) {
120 Matrix->assign(VirtReg: LI, PhysReg);
121 assert(PhysReg != 0);
122 RegsToRewrite.push_back(x: Reg);
123 return true;
124 }
125 }
126
127 llvm_unreachable("physreg not found for WWM expression");
128}
129
130void SIPreAllocateWWMRegs::rewriteRegs(MachineFunction &MF) {
131 for (MachineBasicBlock &MBB : MF) {
132 for (MachineInstr &MI : MBB) {
133 for (MachineOperand &MO : MI.operands()) {
134 if (!MO.isReg())
135 continue;
136
137 const Register VirtReg = MO.getReg();
138 if (VirtReg.isPhysical())
139 continue;
140
141 if (!VirtReg.isValid())
142 continue;
143
144 if (!VRM->hasPhys(virtReg: VirtReg))
145 continue;
146
147 Register PhysReg = VRM->getPhys(virtReg: VirtReg);
148 const unsigned SubReg = MO.getSubReg();
149 if (SubReg != 0) {
150 PhysReg = TRI->getSubReg(Reg: PhysReg, Idx: SubReg);
151 MO.setSubReg(0);
152 }
153
154 MO.setReg(PhysReg);
155 MO.setIsRenamable(false);
156 }
157 }
158 }
159
160 SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
161
162 for (unsigned Reg : RegsToRewrite) {
163 const Register PhysReg = VRM->getPhys(virtReg: Reg);
164 assert(PhysReg != 0);
165
166 LiveInterval &LI = LIS->getInterval(Reg);
167 Matrix->unassign(VirtReg: LI, /*ClearAllReferencingSegments=*/true);
168 LIS->removeInterval(Reg);
169
170 MFI->reserveWWMRegister(Reg: PhysReg);
171 }
172
173 RegsToRewrite.clear();
174
175 // Update the set of reserved registers to include WWM ones
176 // without unnecessarily invalidating RegClassInfo.
177 MRI->freezeReservedRegs();
178 RCI.updateReservedRegs(ReservedInput: MRI->getReservedRegs());
179}
180
181#ifndef NDEBUG
182LLVM_DUMP_METHOD void
183SIPreAllocateWWMRegs::printWWMInfo(const MachineInstr &MI) {
184
185 unsigned Opc = MI.getOpcode();
186
187 if (Opc == AMDGPU::ENTER_STRICT_WWM || Opc == AMDGPU::ENTER_STRICT_WQM) {
188 dbgs() << "Entering ";
189 } else {
190 assert(Opc == AMDGPU::EXIT_STRICT_WWM || Opc == AMDGPU::EXIT_STRICT_WQM);
191 dbgs() << "Exiting ";
192 }
193
194 if (Opc == AMDGPU::ENTER_STRICT_WWM || Opc == AMDGPU::EXIT_STRICT_WWM) {
195 dbgs() << "Strict WWM ";
196 } else {
197 assert(Opc == AMDGPU::ENTER_STRICT_WQM || Opc == AMDGPU::EXIT_STRICT_WQM);
198 dbgs() << "Strict WQM ";
199 }
200
201 dbgs() << "region: " << MI;
202}
203
204#endif
205
206bool SIPreAllocateWWMRegsLegacy::runOnMachineFunction(MachineFunction &MF) {
207 auto *LIS = &getAnalysis<LiveIntervalsWrapperPass>().getLIS();
208 auto *Matrix = &getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
209 auto *VRM = &getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
210 auto &RCI = getAnalysis<MachineRegisterClassInfoWrapperPass>().getRCI();
211 return SIPreAllocateWWMRegs(LIS, Matrix, VRM, RCI).run(MF);
212}
213
214bool SIPreAllocateWWMRegs::run(MachineFunction &MF) {
215 LLVM_DEBUG(dbgs() << "SIPreAllocateWWMRegs: function " << MF.getName() << "\n");
216
217 const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
218
219 TII = ST.getInstrInfo();
220 TRI = &TII->getRegisterInfo();
221 MRI = &MF.getRegInfo();
222
223 bool PreallocateSGPRSpillVGPRs = isPreallocateSGPRSpillVGPRsEnabled(MF);
224
225 bool RegsAssigned = false;
226
227 // We use a reverse post-order traversal of the control-flow graph to
228 // guarantee that we visit definitions in dominance order. Since WWM
229 // expressions are guaranteed to never involve phi nodes, and we can only
230 // escape WWM through the special WWM instruction, this means that this is a
231 // perfect elimination order, so we can never do any better.
232 ReversePostOrderTraversal<MachineFunction*> RPOT(&MF);
233
234 for (MachineBasicBlock *MBB : RPOT) {
235 bool InWWM = false;
236 for (MachineInstr &MI : *MBB) {
237 if (MI.getOpcode() == AMDGPU::SI_SPILL_S32_TO_VGPR) {
238 if (PreallocateSGPRSpillVGPRs)
239 RegsAssigned |= processDef(MO&: MI.getOperand(i: 0));
240 continue;
241 }
242
243 if (MI.getOpcode() == AMDGPU::ENTER_STRICT_WWM ||
244 MI.getOpcode() == AMDGPU::ENTER_STRICT_WQM) {
245 LLVM_DEBUG(printWWMInfo(MI));
246 InWWM = true;
247 continue;
248 }
249
250 if (MI.getOpcode() == AMDGPU::EXIT_STRICT_WWM ||
251 MI.getOpcode() == AMDGPU::EXIT_STRICT_WQM) {
252 LLVM_DEBUG(printWWMInfo(MI));
253 InWWM = false;
254 }
255
256 if (!InWWM)
257 continue;
258
259 LLVM_DEBUG(dbgs() << "Processing " << MI);
260
261 for (MachineOperand &DefOpnd : MI.defs()) {
262 RegsAssigned |= processDef(MO&: DefOpnd);
263 }
264 }
265 }
266
267 if (!RegsAssigned)
268 return false;
269
270 rewriteRegs(MF);
271 return true;
272}
273
274PreservedAnalyses
275SIPreAllocateWWMRegsPass::run(MachineFunction &MF,
276 MachineFunctionAnalysisManager &MFAM) {
277 auto *LIS = &MFAM.getResult<LiveIntervalsAnalysis>(IR&: MF);
278 auto *Matrix = &MFAM.getResult<LiveRegMatrixAnalysis>(IR&: MF);
279 auto *VRM = &MFAM.getResult<VirtRegMapAnalysis>(IR&: MF);
280 auto &RCI = MFAM.getResult<MachineRegisterClassAnalysis>(IR&: MF);
281 SIPreAllocateWWMRegs(LIS, Matrix, VRM, RCI).run(MF);
282 return PreservedAnalyses::all();
283}
284