xref: /llvm-project/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp (revision c4e517f59c086eafe2eb61d23197820f05be799c) 
1 //===- SIMachineFunctionInfo.cpp - SI Machine Function Info ---------------===//
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 #include "SIMachineFunctionInfo.h"
10 #include "AMDGPUSubtarget.h"
11 #include "AMDGPUTargetMachine.h"
12 #include "GCNSubtarget.h"
13 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
14 #include "SIRegisterInfo.h"
15 #include "Utils/AMDGPUBaseInfo.h"
16 #include "llvm/CodeGen/LiveIntervals.h"
17 #include "llvm/CodeGen/MIRParser/MIParser.h"
18 #include "llvm/CodeGen/MachineBasicBlock.h"
19 #include "llvm/CodeGen/MachineFrameInfo.h"
20 #include "llvm/CodeGen/MachineFunction.h"
21 #include "llvm/CodeGen/MachineRegisterInfo.h"
22 #include "llvm/IR/CallingConv.h"
23 #include "llvm/IR/DiagnosticInfo.h"
24 #include "llvm/IR/Function.h"
25 #include <cassert>
26 #include <optional>
27 #include <vector>
28 
29 #define MAX_LANES 64
30 
31 using namespace llvm;
32 
33 const GCNTargetMachine &getTM(const GCNSubtarget *STI) {
34   const SITargetLowering *TLI = STI->getTargetLowering();
35   return static_cast<const GCNTargetMachine &>(TLI->getTargetMachine());
36 }
37 
38 SIMachineFunctionInfo::SIMachineFunctionInfo(const Function &F,
39                                              const GCNSubtarget *STI)
40     : AMDGPUMachineFunction(F, *STI), Mode(F, *STI), GWSResourcePSV(getTM(STI)),
41       UserSGPRInfo(F, *STI), WorkGroupIDX(false), WorkGroupIDY(false),
42       WorkGroupIDZ(false), WorkGroupInfo(false), LDSKernelId(false),
43       PrivateSegmentWaveByteOffset(false), WorkItemIDX(false),
44       WorkItemIDY(false), WorkItemIDZ(false), ImplicitArgPtr(false),
45       GITPtrHigh(0xffffffff), HighBitsOf32BitAddress(0) {
46   const GCNSubtarget &ST = *static_cast<const GCNSubtarget *>(STI);
47   FlatWorkGroupSizes = ST.getFlatWorkGroupSizes(F);
48   WavesPerEU = ST.getWavesPerEU(F);
49   MaxNumWorkGroups = ST.getMaxNumWorkGroups(F);
50   assert(MaxNumWorkGroups.size() == 3);
51 
52   Occupancy = ST.computeOccupancy(F, getLDSSize());
53   CallingConv::ID CC = F.getCallingConv();
54 
55   VRegFlags.reserve(1024);
56 
57   const bool IsKernel = CC == CallingConv::AMDGPU_KERNEL ||
58                         CC == CallingConv::SPIR_KERNEL;
59 
60   if (IsKernel) {
61     WorkGroupIDX = true;
62     WorkItemIDX = true;
63   } else if (CC == CallingConv::AMDGPU_PS) {
64     PSInputAddr = AMDGPU::getInitialPSInputAddr(F);
65   }
66 
67   MayNeedAGPRs = ST.hasMAIInsts();
68 
69   if (AMDGPU::isChainCC(CC)) {
70     // Chain functions don't receive an SP from their caller, but are free to
71     // set one up. For now, we can use s32 to match what amdgpu_gfx functions
72     // would use if called, but this can be revisited.
73     // FIXME: Only reserve this if we actually need it.
74     StackPtrOffsetReg = AMDGPU::SGPR32;
75 
76     ScratchRSrcReg = AMDGPU::SGPR48_SGPR49_SGPR50_SGPR51;
77 
78     ArgInfo.PrivateSegmentBuffer =
79         ArgDescriptor::createRegister(ScratchRSrcReg);
80 
81     ImplicitArgPtr = false;
82   } else if (!isEntryFunction()) {
83     if (CC != CallingConv::AMDGPU_Gfx)
84       ArgInfo = AMDGPUArgumentUsageInfo::FixedABIFunctionInfo;
85 
86     // TODO: Pick a high register, and shift down, similar to a kernel.
87     FrameOffsetReg = AMDGPU::SGPR33;
88     StackPtrOffsetReg = AMDGPU::SGPR32;
89 
90     if (!ST.enableFlatScratch()) {
91       // Non-entry functions have no special inputs for now, other registers
92       // required for scratch access.
93       ScratchRSrcReg = AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3;
94 
95       ArgInfo.PrivateSegmentBuffer =
96         ArgDescriptor::createRegister(ScratchRSrcReg);
97     }
98 
99     if (!F.hasFnAttribute("amdgpu-no-implicitarg-ptr"))
100       ImplicitArgPtr = true;
101   } else {
102     ImplicitArgPtr = false;
103     MaxKernArgAlign = std::max(ST.getAlignmentForImplicitArgPtr(),
104                                MaxKernArgAlign);
105 
106     if (ST.hasGFX90AInsts() &&
107         ST.getMaxNumVGPRs(F) <= AMDGPU::VGPR_32RegClass.getNumRegs() &&
108         !mayUseAGPRs(F))
109       MayNeedAGPRs = false; // We will select all MAI with VGPR operands.
110   }
111 
112   if (!AMDGPU::isGraphics(CC) ||
113       (CC == CallingConv::AMDGPU_CS && ST.hasArchitectedSGPRs())) {
114     if (IsKernel || !F.hasFnAttribute("amdgpu-no-workgroup-id-x"))
115       WorkGroupIDX = true;
116 
117     if (!F.hasFnAttribute("amdgpu-no-workgroup-id-y"))
118       WorkGroupIDY = true;
119 
120     if (!F.hasFnAttribute("amdgpu-no-workgroup-id-z"))
121       WorkGroupIDZ = true;
122   }
123 
124   if (!AMDGPU::isGraphics(CC)) {
125     if (IsKernel || !F.hasFnAttribute("amdgpu-no-workitem-id-x"))
126       WorkItemIDX = true;
127 
128     if (!F.hasFnAttribute("amdgpu-no-workitem-id-y") &&
129         ST.getMaxWorkitemID(F, 1) != 0)
130       WorkItemIDY = true;
131 
132     if (!F.hasFnAttribute("amdgpu-no-workitem-id-z") &&
133         ST.getMaxWorkitemID(F, 2) != 0)
134       WorkItemIDZ = true;
135 
136     if (!IsKernel && !F.hasFnAttribute("amdgpu-no-lds-kernel-id"))
137       LDSKernelId = true;
138   }
139 
140   if (isEntryFunction()) {
141     // X, XY, and XYZ are the only supported combinations, so make sure Y is
142     // enabled if Z is.
143     if (WorkItemIDZ)
144       WorkItemIDY = true;
145 
146     if (!ST.flatScratchIsArchitected()) {
147       PrivateSegmentWaveByteOffset = true;
148 
149       // HS and GS always have the scratch wave offset in SGPR5 on GFX9.
150       if (ST.getGeneration() >= AMDGPUSubtarget::GFX9 &&
151           (CC == CallingConv::AMDGPU_HS || CC == CallingConv::AMDGPU_GS))
152         ArgInfo.PrivateSegmentWaveByteOffset =
153             ArgDescriptor::createRegister(AMDGPU::SGPR5);
154     }
155   }
156 
157   Attribute A = F.getFnAttribute("amdgpu-git-ptr-high");
158   StringRef S = A.getValueAsString();
159   if (!S.empty())
160     S.consumeInteger(0, GITPtrHigh);
161 
162   A = F.getFnAttribute("amdgpu-32bit-address-high-bits");
163   S = A.getValueAsString();
164   if (!S.empty())
165     S.consumeInteger(0, HighBitsOf32BitAddress);
166 
167   // On GFX908, in order to guarantee copying between AGPRs, we need a scratch
168   // VGPR available at all times. For now, reserve highest available VGPR. After
169   // RA, shift it to the lowest available unused VGPR if the one exist.
170   if (ST.hasMAIInsts() && !ST.hasGFX90AInsts()) {
171     VGPRForAGPRCopy =
172         AMDGPU::VGPR_32RegClass.getRegister(ST.getMaxNumVGPRs(F) - 1);
173   }
174 }
175 
176 MachineFunctionInfo *SIMachineFunctionInfo::clone(
177     BumpPtrAllocator &Allocator, MachineFunction &DestMF,
178     const DenseMap<MachineBasicBlock *, MachineBasicBlock *> &Src2DstMBB)
179     const {
180   return DestMF.cloneInfo<SIMachineFunctionInfo>(*this);
181 }
182 
183 void SIMachineFunctionInfo::limitOccupancy(const MachineFunction &MF) {
184   limitOccupancy(getMaxWavesPerEU());
185   const GCNSubtarget& ST = MF.getSubtarget<GCNSubtarget>();
186   limitOccupancy(ST.getOccupancyWithLocalMemSize(getLDSSize(),
187                  MF.getFunction()));
188 }
189 
190 Register SIMachineFunctionInfo::addPrivateSegmentBuffer(
191   const SIRegisterInfo &TRI) {
192   ArgInfo.PrivateSegmentBuffer =
193     ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
194     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SGPR_128RegClass));
195   NumUserSGPRs += 4;
196   return ArgInfo.PrivateSegmentBuffer.getRegister();
197 }
198 
199 Register SIMachineFunctionInfo::addDispatchPtr(const SIRegisterInfo &TRI) {
200   ArgInfo.DispatchPtr = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
201     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
202   NumUserSGPRs += 2;
203   return ArgInfo.DispatchPtr.getRegister();
204 }
205 
206 Register SIMachineFunctionInfo::addQueuePtr(const SIRegisterInfo &TRI) {
207   ArgInfo.QueuePtr = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
208     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
209   NumUserSGPRs += 2;
210   return ArgInfo.QueuePtr.getRegister();
211 }
212 
213 Register SIMachineFunctionInfo::addKernargSegmentPtr(const SIRegisterInfo &TRI) {
214   ArgInfo.KernargSegmentPtr
215     = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
216     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
217   NumUserSGPRs += 2;
218   return ArgInfo.KernargSegmentPtr.getRegister();
219 }
220 
221 Register SIMachineFunctionInfo::addDispatchID(const SIRegisterInfo &TRI) {
222   ArgInfo.DispatchID = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
223     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
224   NumUserSGPRs += 2;
225   return ArgInfo.DispatchID.getRegister();
226 }
227 
228 Register SIMachineFunctionInfo::addFlatScratchInit(const SIRegisterInfo &TRI) {
229   ArgInfo.FlatScratchInit = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
230     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
231   NumUserSGPRs += 2;
232   return ArgInfo.FlatScratchInit.getRegister();
233 }
234 
235 Register SIMachineFunctionInfo::addImplicitBufferPtr(const SIRegisterInfo &TRI) {
236   ArgInfo.ImplicitBufferPtr = ArgDescriptor::createRegister(TRI.getMatchingSuperReg(
237     getNextUserSGPR(), AMDGPU::sub0, &AMDGPU::SReg_64RegClass));
238   NumUserSGPRs += 2;
239   return ArgInfo.ImplicitBufferPtr.getRegister();
240 }
241 
242 Register SIMachineFunctionInfo::addLDSKernelId() {
243   ArgInfo.LDSKernelId = ArgDescriptor::createRegister(getNextUserSGPR());
244   NumUserSGPRs += 1;
245   return ArgInfo.LDSKernelId.getRegister();
246 }
247 
248 SmallVectorImpl<MCRegister> *SIMachineFunctionInfo::addPreloadedKernArg(
249     const SIRegisterInfo &TRI, const TargetRegisterClass *RC,
250     unsigned AllocSizeDWord, int KernArgIdx, int PaddingSGPRs) {
251   assert(!ArgInfo.PreloadKernArgs.count(KernArgIdx) &&
252          "Preload kernel argument allocated twice.");
253   NumUserSGPRs += PaddingSGPRs;
254   // If the available register tuples are aligned with the kernarg to be
255   // preloaded use that register, otherwise we need to use a set of SGPRs and
256   // merge them.
257   Register PreloadReg =
258       TRI.getMatchingSuperReg(getNextUserSGPR(), AMDGPU::sub0, RC);
259   if (PreloadReg &&
260       (RC == &AMDGPU::SReg_32RegClass || RC == &AMDGPU::SReg_64RegClass)) {
261     ArgInfo.PreloadKernArgs[KernArgIdx].Regs.push_back(PreloadReg);
262     NumUserSGPRs += AllocSizeDWord;
263   } else {
264     for (unsigned I = 0; I < AllocSizeDWord; ++I) {
265       ArgInfo.PreloadKernArgs[KernArgIdx].Regs.push_back(getNextUserSGPR());
266       NumUserSGPRs++;
267     }
268   }
269 
270   // Track the actual number of SGPRs that HW will preload to.
271   UserSGPRInfo.allocKernargPreloadSGPRs(AllocSizeDWord + PaddingSGPRs);
272   return &ArgInfo.PreloadKernArgs[KernArgIdx].Regs;
273 }
274 
275 void SIMachineFunctionInfo::allocateWWMSpill(MachineFunction &MF, Register VGPR,
276                                              uint64_t Size, Align Alignment) {
277   // Skip if it is an entry function or the register is already added.
278   if (isEntryFunction() || WWMSpills.count(VGPR))
279     return;
280 
281   // Skip if this is a function with the amdgpu_cs_chain or
282   // amdgpu_cs_chain_preserve calling convention and this is a scratch register.
283   // We never need to allocate a spill for these because we don't even need to
284   // restore the inactive lanes for them (they're scratchier than the usual
285   // scratch registers).
286   if (isChainFunction() && SIRegisterInfo::isChainScratchRegister(VGPR))
287     return;
288 
289   WWMSpills.insert(std::make_pair(
290       VGPR, MF.getFrameInfo().CreateSpillStackObject(Size, Alignment)));
291 }
292 
293 // Separate out the callee-saved and scratch registers.
294 void SIMachineFunctionInfo::splitWWMSpillRegisters(
295     MachineFunction &MF,
296     SmallVectorImpl<std::pair<Register, int>> &CalleeSavedRegs,
297     SmallVectorImpl<std::pair<Register, int>> &ScratchRegs) const {
298   const MCPhysReg *CSRegs = MF.getRegInfo().getCalleeSavedRegs();
299   for (auto &Reg : WWMSpills) {
300     if (isCalleeSavedReg(CSRegs, Reg.first))
301       CalleeSavedRegs.push_back(Reg);
302     else
303       ScratchRegs.push_back(Reg);
304   }
305 }
306 
307 bool SIMachineFunctionInfo::isCalleeSavedReg(const MCPhysReg *CSRegs,
308                                              MCPhysReg Reg) const {
309   for (unsigned I = 0; CSRegs[I]; ++I) {
310     if (CSRegs[I] == Reg)
311       return true;
312   }
313 
314   return false;
315 }
316 
317 void SIMachineFunctionInfo::shiftSpillPhysVGPRsToLowestRange(
318     MachineFunction &MF) {
319   const SIRegisterInfo *TRI = MF.getSubtarget<GCNSubtarget>().getRegisterInfo();
320   MachineRegisterInfo &MRI = MF.getRegInfo();
321   for (unsigned I = 0, E = SpillPhysVGPRs.size(); I < E; ++I) {
322     Register Reg = SpillPhysVGPRs[I];
323     Register NewReg =
324         TRI->findUnusedRegister(MRI, &AMDGPU::VGPR_32RegClass, MF);
325     if (!NewReg || NewReg >= Reg)
326       break;
327 
328     MRI.replaceRegWith(Reg, NewReg);
329 
330     // Update various tables with the new VGPR.
331     SpillPhysVGPRs[I] = NewReg;
332     WWMReservedRegs.remove(Reg);
333     WWMReservedRegs.insert(NewReg);
334     WWMSpills.insert(std::make_pair(NewReg, WWMSpills[Reg]));
335     WWMSpills.erase(Reg);
336 
337     for (MachineBasicBlock &MBB : MF) {
338       MBB.removeLiveIn(Reg);
339       MBB.sortUniqueLiveIns();
340     }
341   }
342 }
343 
344 bool SIMachineFunctionInfo::allocateVirtualVGPRForSGPRSpills(
345     MachineFunction &MF, int FI, unsigned LaneIndex) {
346   MachineRegisterInfo &MRI = MF.getRegInfo();
347   Register LaneVGPR;
348   if (!LaneIndex) {
349     LaneVGPR = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
350     SpillVGPRs.push_back(LaneVGPR);
351   } else {
352     LaneVGPR = SpillVGPRs.back();
353   }
354 
355   SGPRSpillsToVirtualVGPRLanes[FI].push_back(
356       SIRegisterInfo::SpilledReg(LaneVGPR, LaneIndex));
357   return true;
358 }
359 
360 bool SIMachineFunctionInfo::allocatePhysicalVGPRForSGPRSpills(
361     MachineFunction &MF, int FI, unsigned LaneIndex, bool IsPrologEpilog) {
362   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
363   const SIRegisterInfo *TRI = ST.getRegisterInfo();
364   MachineRegisterInfo &MRI = MF.getRegInfo();
365   Register LaneVGPR;
366   if (!LaneIndex) {
367     // Find the highest available register if called before RA to ensure the
368     // lowest registers are available for allocation. The LaneVGPR, in that
369     // case, will be shifted back to the lowest range after VGPR allocation.
370     LaneVGPR = TRI->findUnusedRegister(MRI, &AMDGPU::VGPR_32RegClass, MF,
371                                        !IsPrologEpilog);
372     if (LaneVGPR == AMDGPU::NoRegister) {
373       // We have no VGPRs left for spilling SGPRs. Reset because we will not
374       // partially spill the SGPR to VGPRs.
375       SGPRSpillsToPhysicalVGPRLanes.erase(FI);
376       return false;
377     }
378 
379     allocateWWMSpill(MF, LaneVGPR);
380     reserveWWMRegister(LaneVGPR);
381     for (MachineBasicBlock &MBB : MF) {
382       MBB.addLiveIn(LaneVGPR);
383       MBB.sortUniqueLiveIns();
384     }
385     SpillPhysVGPRs.push_back(LaneVGPR);
386   } else {
387     LaneVGPR = SpillPhysVGPRs.back();
388   }
389 
390   SGPRSpillsToPhysicalVGPRLanes[FI].push_back(
391       SIRegisterInfo::SpilledReg(LaneVGPR, LaneIndex));
392   return true;
393 }
394 
395 bool SIMachineFunctionInfo::allocateSGPRSpillToVGPRLane(
396     MachineFunction &MF, int FI, bool SpillToPhysVGPRLane,
397     bool IsPrologEpilog) {
398   std::vector<SIRegisterInfo::SpilledReg> &SpillLanes =
399       SpillToPhysVGPRLane ? SGPRSpillsToPhysicalVGPRLanes[FI]
400                           : SGPRSpillsToVirtualVGPRLanes[FI];
401 
402   // This has already been allocated.
403   if (!SpillLanes.empty())
404     return true;
405 
406   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
407   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
408   unsigned WaveSize = ST.getWavefrontSize();
409 
410   unsigned Size = FrameInfo.getObjectSize(FI);
411   unsigned NumLanes = Size / 4;
412 
413   if (NumLanes > WaveSize)
414     return false;
415 
416   assert(Size >= 4 && "invalid sgpr spill size");
417   assert(ST.getRegisterInfo()->spillSGPRToVGPR() &&
418          "not spilling SGPRs to VGPRs");
419 
420   unsigned &NumSpillLanes = SpillToPhysVGPRLane ? NumPhysicalVGPRSpillLanes
421                                                 : NumVirtualVGPRSpillLanes;
422 
423   for (unsigned I = 0; I < NumLanes; ++I, ++NumSpillLanes) {
424     unsigned LaneIndex = (NumSpillLanes % WaveSize);
425 
426     bool Allocated = SpillToPhysVGPRLane
427                          ? allocatePhysicalVGPRForSGPRSpills(MF, FI, LaneIndex,
428                                                              IsPrologEpilog)
429                          : allocateVirtualVGPRForSGPRSpills(MF, FI, LaneIndex);
430     if (!Allocated) {
431       NumSpillLanes -= I;
432       return false;
433     }
434   }
435 
436   return true;
437 }
438 
439 /// Reserve AGPRs or VGPRs to support spilling for FrameIndex \p FI.
440 /// Either AGPR is spilled to VGPR to vice versa.
441 /// Returns true if a \p FI can be eliminated completely.
442 bool SIMachineFunctionInfo::allocateVGPRSpillToAGPR(MachineFunction &MF,
443                                                     int FI,
444                                                     bool isAGPRtoVGPR) {
445   MachineRegisterInfo &MRI = MF.getRegInfo();
446   MachineFrameInfo &FrameInfo = MF.getFrameInfo();
447   const GCNSubtarget &ST =  MF.getSubtarget<GCNSubtarget>();
448 
449   assert(ST.hasMAIInsts() && FrameInfo.isSpillSlotObjectIndex(FI));
450 
451   auto &Spill = VGPRToAGPRSpills[FI];
452 
453   // This has already been allocated.
454   if (!Spill.Lanes.empty())
455     return Spill.FullyAllocated;
456 
457   unsigned Size = FrameInfo.getObjectSize(FI);
458   unsigned NumLanes = Size / 4;
459   Spill.Lanes.resize(NumLanes, AMDGPU::NoRegister);
460 
461   const TargetRegisterClass &RC =
462       isAGPRtoVGPR ? AMDGPU::VGPR_32RegClass : AMDGPU::AGPR_32RegClass;
463   auto Regs = RC.getRegisters();
464 
465   auto &SpillRegs = isAGPRtoVGPR ? SpillAGPR : SpillVGPR;
466   const SIRegisterInfo *TRI = ST.getRegisterInfo();
467   Spill.FullyAllocated = true;
468 
469   // FIXME: Move allocation logic out of MachineFunctionInfo and initialize
470   // once.
471   BitVector OtherUsedRegs;
472   OtherUsedRegs.resize(TRI->getNumRegs());
473 
474   const uint32_t *CSRMask =
475       TRI->getCallPreservedMask(MF, MF.getFunction().getCallingConv());
476   if (CSRMask)
477     OtherUsedRegs.setBitsInMask(CSRMask);
478 
479   // TODO: Should include register tuples, but doesn't matter with current
480   // usage.
481   for (MCPhysReg Reg : SpillAGPR)
482     OtherUsedRegs.set(Reg);
483   for (MCPhysReg Reg : SpillVGPR)
484     OtherUsedRegs.set(Reg);
485 
486   SmallVectorImpl<MCPhysReg>::const_iterator NextSpillReg = Regs.begin();
487   for (int I = NumLanes - 1; I >= 0; --I) {
488     NextSpillReg = std::find_if(
489         NextSpillReg, Regs.end(), [&MRI, &OtherUsedRegs](MCPhysReg Reg) {
490           return MRI.isAllocatable(Reg) && !MRI.isPhysRegUsed(Reg) &&
491                  !OtherUsedRegs[Reg];
492         });
493 
494     if (NextSpillReg == Regs.end()) { // Registers exhausted
495       Spill.FullyAllocated = false;
496       break;
497     }
498 
499     OtherUsedRegs.set(*NextSpillReg);
500     SpillRegs.push_back(*NextSpillReg);
501     MRI.reserveReg(*NextSpillReg, TRI);
502     Spill.Lanes[I] = *NextSpillReg++;
503   }
504 
505   return Spill.FullyAllocated;
506 }
507 
508 bool SIMachineFunctionInfo::removeDeadFrameIndices(
509     MachineFrameInfo &MFI, bool ResetSGPRSpillStackIDs) {
510   // Remove dead frame indices from function frame, however keep FP & BP since
511   // spills for them haven't been inserted yet. And also make sure to remove the
512   // frame indices from `SGPRSpillsToVirtualVGPRLanes` data structure,
513   // otherwise, it could result in an unexpected side effect and bug, in case of
514   // any re-mapping of freed frame indices by later pass(es) like "stack slot
515   // coloring".
516   for (auto &R : make_early_inc_range(SGPRSpillsToVirtualVGPRLanes)) {
517     MFI.RemoveStackObject(R.first);
518     SGPRSpillsToVirtualVGPRLanes.erase(R.first);
519   }
520 
521   // Remove the dead frame indices of CSR SGPRs which are spilled to physical
522   // VGPR lanes during SILowerSGPRSpills pass.
523   if (!ResetSGPRSpillStackIDs) {
524     for (auto &R : make_early_inc_range(SGPRSpillsToPhysicalVGPRLanes)) {
525       MFI.RemoveStackObject(R.first);
526       SGPRSpillsToPhysicalVGPRLanes.erase(R.first);
527     }
528   }
529   bool HaveSGPRToMemory = false;
530 
531   if (ResetSGPRSpillStackIDs) {
532     // All other SGPRs must be allocated on the default stack, so reset the
533     // stack ID.
534     for (int I = MFI.getObjectIndexBegin(), E = MFI.getObjectIndexEnd(); I != E;
535          ++I) {
536       if (!checkIndexInPrologEpilogSGPRSpills(I)) {
537         if (MFI.getStackID(I) == TargetStackID::SGPRSpill) {
538           MFI.setStackID(I, TargetStackID::Default);
539           HaveSGPRToMemory = true;
540         }
541       }
542     }
543   }
544 
545   for (auto &R : VGPRToAGPRSpills) {
546     if (R.second.IsDead)
547       MFI.RemoveStackObject(R.first);
548   }
549 
550   return HaveSGPRToMemory;
551 }
552 
553 int SIMachineFunctionInfo::getScavengeFI(MachineFrameInfo &MFI,
554                                          const SIRegisterInfo &TRI) {
555   if (ScavengeFI)
556     return *ScavengeFI;
557 
558   ScavengeFI =
559       MFI.CreateStackObject(TRI.getSpillSize(AMDGPU::SGPR_32RegClass),
560                             TRI.getSpillAlign(AMDGPU::SGPR_32RegClass), false);
561   return *ScavengeFI;
562 }
563 
564 MCPhysReg SIMachineFunctionInfo::getNextUserSGPR() const {
565   assert(NumSystemSGPRs == 0 && "System SGPRs must be added after user SGPRs");
566   return AMDGPU::SGPR0 + NumUserSGPRs;
567 }
568 
569 MCPhysReg SIMachineFunctionInfo::getNextSystemSGPR() const {
570   return AMDGPU::SGPR0 + NumUserSGPRs + NumSystemSGPRs;
571 }
572 
573 void SIMachineFunctionInfo::MRI_NoteNewVirtualRegister(Register Reg) {
574   VRegFlags.grow(Reg);
575 }
576 
577 void SIMachineFunctionInfo::MRI_NoteCloneVirtualRegister(Register NewReg,
578                                                          Register SrcReg) {
579   VRegFlags.grow(NewReg);
580   VRegFlags[NewReg] = VRegFlags[SrcReg];
581 }
582 
583 Register
584 SIMachineFunctionInfo::getGITPtrLoReg(const MachineFunction &MF) const {
585   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
586   if (!ST.isAmdPalOS())
587     return Register();
588   Register GitPtrLo = AMDGPU::SGPR0; // Low GIT address passed in
589   if (ST.hasMergedShaders()) {
590     switch (MF.getFunction().getCallingConv()) {
591     case CallingConv::AMDGPU_HS:
592     case CallingConv::AMDGPU_GS:
593       // Low GIT address is passed in s8 rather than s0 for an LS+HS or
594       // ES+GS merged shader on gfx9+.
595       GitPtrLo = AMDGPU::SGPR8;
596       return GitPtrLo;
597     default:
598       return GitPtrLo;
599     }
600   }
601   return GitPtrLo;
602 }
603 
604 static yaml::StringValue regToString(Register Reg,
605                                      const TargetRegisterInfo &TRI) {
606   yaml::StringValue Dest;
607   {
608     raw_string_ostream OS(Dest.Value);
609     OS << printReg(Reg, &TRI);
610   }
611   return Dest;
612 }
613 
614 static std::optional<yaml::SIArgumentInfo>
615 convertArgumentInfo(const AMDGPUFunctionArgInfo &ArgInfo,
616                     const TargetRegisterInfo &TRI) {
617   yaml::SIArgumentInfo AI;
618 
619   auto convertArg = [&](std::optional<yaml::SIArgument> &A,
620                         const ArgDescriptor &Arg) {
621     if (!Arg)
622       return false;
623 
624     // Create a register or stack argument.
625     yaml::SIArgument SA = yaml::SIArgument::createArgument(Arg.isRegister());
626     if (Arg.isRegister()) {
627       raw_string_ostream OS(SA.RegisterName.Value);
628       OS << printReg(Arg.getRegister(), &TRI);
629     } else
630       SA.StackOffset = Arg.getStackOffset();
631     // Check and update the optional mask.
632     if (Arg.isMasked())
633       SA.Mask = Arg.getMask();
634 
635     A = SA;
636     return true;
637   };
638 
639   // TODO: Need to serialize kernarg preloads.
640   bool Any = false;
641   Any |= convertArg(AI.PrivateSegmentBuffer, ArgInfo.PrivateSegmentBuffer);
642   Any |= convertArg(AI.DispatchPtr, ArgInfo.DispatchPtr);
643   Any |= convertArg(AI.QueuePtr, ArgInfo.QueuePtr);
644   Any |= convertArg(AI.KernargSegmentPtr, ArgInfo.KernargSegmentPtr);
645   Any |= convertArg(AI.DispatchID, ArgInfo.DispatchID);
646   Any |= convertArg(AI.FlatScratchInit, ArgInfo.FlatScratchInit);
647   Any |= convertArg(AI.LDSKernelId, ArgInfo.LDSKernelId);
648   Any |= convertArg(AI.PrivateSegmentSize, ArgInfo.PrivateSegmentSize);
649   Any |= convertArg(AI.WorkGroupIDX, ArgInfo.WorkGroupIDX);
650   Any |= convertArg(AI.WorkGroupIDY, ArgInfo.WorkGroupIDY);
651   Any |= convertArg(AI.WorkGroupIDZ, ArgInfo.WorkGroupIDZ);
652   Any |= convertArg(AI.WorkGroupInfo, ArgInfo.WorkGroupInfo);
653   Any |= convertArg(AI.PrivateSegmentWaveByteOffset,
654                     ArgInfo.PrivateSegmentWaveByteOffset);
655   Any |= convertArg(AI.ImplicitArgPtr, ArgInfo.ImplicitArgPtr);
656   Any |= convertArg(AI.ImplicitBufferPtr, ArgInfo.ImplicitBufferPtr);
657   Any |= convertArg(AI.WorkItemIDX, ArgInfo.WorkItemIDX);
658   Any |= convertArg(AI.WorkItemIDY, ArgInfo.WorkItemIDY);
659   Any |= convertArg(AI.WorkItemIDZ, ArgInfo.WorkItemIDZ);
660 
661   if (Any)
662     return AI;
663 
664   return std::nullopt;
665 }
666 
667 yaml::SIMachineFunctionInfo::SIMachineFunctionInfo(
668     const llvm::SIMachineFunctionInfo &MFI, const TargetRegisterInfo &TRI,
669     const llvm::MachineFunction &MF)
670     : ExplicitKernArgSize(MFI.getExplicitKernArgSize()),
671       MaxKernArgAlign(MFI.getMaxKernArgAlign()), LDSSize(MFI.getLDSSize()),
672       GDSSize(MFI.getGDSSize()),
673       DynLDSAlign(MFI.getDynLDSAlign()), IsEntryFunction(MFI.isEntryFunction()),
674       NoSignedZerosFPMath(MFI.hasNoSignedZerosFPMath()),
675       MemoryBound(MFI.isMemoryBound()), WaveLimiter(MFI.needsWaveLimiter()),
676       HasSpilledSGPRs(MFI.hasSpilledSGPRs()),
677       HasSpilledVGPRs(MFI.hasSpilledVGPRs()),
678       HighBitsOf32BitAddress(MFI.get32BitAddressHighBits()),
679       Occupancy(MFI.getOccupancy()),
680       ScratchRSrcReg(regToString(MFI.getScratchRSrcReg(), TRI)),
681       FrameOffsetReg(regToString(MFI.getFrameOffsetReg(), TRI)),
682       StackPtrOffsetReg(regToString(MFI.getStackPtrOffsetReg(), TRI)),
683       BytesInStackArgArea(MFI.getBytesInStackArgArea()),
684       ReturnsVoid(MFI.returnsVoid()),
685       ArgInfo(convertArgumentInfo(MFI.getArgInfo(), TRI)),
686       PSInputAddr(MFI.getPSInputAddr()),
687       PSInputEnable(MFI.getPSInputEnable()),
688       Mode(MFI.getMode()) {
689   for (Register Reg : MFI.getWWMReservedRegs())
690     WWMReservedRegs.push_back(regToString(Reg, TRI));
691 
692   if (MFI.getLongBranchReservedReg())
693     LongBranchReservedReg = regToString(MFI.getLongBranchReservedReg(), TRI);
694   if (MFI.getVGPRForAGPRCopy())
695     VGPRForAGPRCopy = regToString(MFI.getVGPRForAGPRCopy(), TRI);
696 
697   if (MFI.getSGPRForEXECCopy())
698     SGPRForEXECCopy = regToString(MFI.getSGPRForEXECCopy(), TRI);
699 
700   auto SFI = MFI.getOptionalScavengeFI();
701   if (SFI)
702     ScavengeFI = yaml::FrameIndex(*SFI, MF.getFrameInfo());
703 }
704 
705 void yaml::SIMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
706   MappingTraits<SIMachineFunctionInfo>::mapping(YamlIO, *this);
707 }
708 
709 bool SIMachineFunctionInfo::initializeBaseYamlFields(
710     const yaml::SIMachineFunctionInfo &YamlMFI, const MachineFunction &MF,
711     PerFunctionMIParsingState &PFS, SMDiagnostic &Error, SMRange &SourceRange) {
712   ExplicitKernArgSize = YamlMFI.ExplicitKernArgSize;
713   MaxKernArgAlign = YamlMFI.MaxKernArgAlign;
714   LDSSize = YamlMFI.LDSSize;
715   GDSSize = YamlMFI.GDSSize;
716   DynLDSAlign = YamlMFI.DynLDSAlign;
717   PSInputAddr = YamlMFI.PSInputAddr;
718   PSInputEnable = YamlMFI.PSInputEnable;
719   HighBitsOf32BitAddress = YamlMFI.HighBitsOf32BitAddress;
720   Occupancy = YamlMFI.Occupancy;
721   IsEntryFunction = YamlMFI.IsEntryFunction;
722   NoSignedZerosFPMath = YamlMFI.NoSignedZerosFPMath;
723   MemoryBound = YamlMFI.MemoryBound;
724   WaveLimiter = YamlMFI.WaveLimiter;
725   HasSpilledSGPRs = YamlMFI.HasSpilledSGPRs;
726   HasSpilledVGPRs = YamlMFI.HasSpilledVGPRs;
727   BytesInStackArgArea = YamlMFI.BytesInStackArgArea;
728   ReturnsVoid = YamlMFI.ReturnsVoid;
729 
730   if (YamlMFI.ScavengeFI) {
731     auto FIOrErr = YamlMFI.ScavengeFI->getFI(MF.getFrameInfo());
732     if (!FIOrErr) {
733       // Create a diagnostic for a the frame index.
734       const MemoryBuffer &Buffer =
735           *PFS.SM->getMemoryBuffer(PFS.SM->getMainFileID());
736 
737       Error = SMDiagnostic(*PFS.SM, SMLoc(), Buffer.getBufferIdentifier(), 1, 1,
738                            SourceMgr::DK_Error, toString(FIOrErr.takeError()),
739                            "", std::nullopt, std::nullopt);
740       SourceRange = YamlMFI.ScavengeFI->SourceRange;
741       return true;
742     }
743     ScavengeFI = *FIOrErr;
744   } else {
745     ScavengeFI = std::nullopt;
746   }
747   return false;
748 }
749 
750 bool SIMachineFunctionInfo::mayUseAGPRs(const Function &F) const {
751   for (const BasicBlock &BB : F) {
752     for (const Instruction &I : BB) {
753       const auto *CB = dyn_cast<CallBase>(&I);
754       if (!CB)
755         continue;
756 
757       if (CB->isInlineAsm()) {
758         const InlineAsm *IA = dyn_cast<InlineAsm>(CB->getCalledOperand());
759         for (const auto &CI : IA->ParseConstraints()) {
760           for (StringRef Code : CI.Codes) {
761             Code.consume_front("{");
762             if (Code.starts_with("a"))
763               return true;
764           }
765         }
766         continue;
767       }
768 
769       const Function *Callee =
770           dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
771       if (!Callee)
772         return true;
773 
774       if (Callee->getIntrinsicID() == Intrinsic::not_intrinsic)
775         return true;
776     }
777   }
778 
779   return false;
780 }
781 
782 bool SIMachineFunctionInfo::usesAGPRs(const MachineFunction &MF) const {
783   if (UsesAGPRs)
784     return *UsesAGPRs;
785 
786   if (!mayNeedAGPRs()) {
787     UsesAGPRs = false;
788     return false;
789   }
790 
791   if (!AMDGPU::isEntryFunctionCC(MF.getFunction().getCallingConv()) ||
792       MF.getFrameInfo().hasCalls()) {
793     UsesAGPRs = true;
794     return true;
795   }
796 
797   const MachineRegisterInfo &MRI = MF.getRegInfo();
798 
799   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
800     const Register Reg = Register::index2VirtReg(I);
801     const TargetRegisterClass *RC = MRI.getRegClassOrNull(Reg);
802     if (RC && SIRegisterInfo::isAGPRClass(RC)) {
803       UsesAGPRs = true;
804       return true;
805     } else if (!RC && !MRI.use_empty(Reg) && MRI.getType(Reg).isValid()) {
806       // Defer caching UsesAGPRs, function might not yet been regbank selected.
807       return true;
808     }
809   }
810 
811   for (MCRegister Reg : AMDGPU::AGPR_32RegClass) {
812     if (MRI.isPhysRegUsed(Reg)) {
813       UsesAGPRs = true;
814       return true;
815     }
816   }
817 
818   UsesAGPRs = false;
819   return false;
820 }
821