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