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