1//===-- AMDGPURewriteAGPRCopyMFMA.cpp -------------------------------------===//
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 \brief Try to replace MFMA instructions using VGPRs with MFMA
10/// instructions using AGPRs. We expect MFMAs to be selected using VGPRs, and
11/// only use AGPRs if it helps avoid spilling. In this case, the MFMA will have
12/// copies between AGPRs and VGPRs and the AGPR variant of an MFMA pseudo. This
13/// pass will attempt to delete the cross register bank copy and replace the
14/// MFMA opcode.
15///
16/// TODO:
17/// - Handle rewrites of phis. This must be more careful than normal about the
18/// reassignment. We do not want to introduce an AGPR-to-AGPR copy inside of a
19/// loop, so it depends on the exact assignment of the copy.
20///
21/// - Update LiveIntervals incrementally instead of recomputing from scratch
22///
23//===----------------------------------------------------------------------===//
24
25#include "AMDGPU.h"
26#include "GCNSubtarget.h"
27#include "SIMachineFunctionInfo.h"
28#include "SIRegisterInfo.h"
29#include "llvm/ADT/Statistic.h"
30#include "llvm/CodeGen/LiveIntervals.h"
31#include "llvm/CodeGen/LiveRegMatrix.h"
32#include "llvm/CodeGen/LiveStacks.h"
33#include "llvm/CodeGen/MachineDominators.h"
34#include "llvm/CodeGen/MachineFrameInfo.h"
35#include "llvm/CodeGen/MachineFunctionPass.h"
36#include "llvm/CodeGen/SlotIndexes.h"
37#include "llvm/CodeGen/VirtRegMap.h"
38#include "llvm/InitializePasses.h"
39#include "llvm/Support/DebugCounter.h"
40
41using namespace llvm;
42
43#define DEBUG_TYPE "amdgpu-rewrite-agpr-copy-mfma"
44
45DEBUG_COUNTER(RewriteAGPRCopyMFMACounter, DEBUG_TYPE,
46 "Controls which MFMA chains are rewritten to AGPR form");
47
48namespace {
49
50STATISTIC(NumMFMAsRewrittenToAGPR,
51 "Number of MFMA instructions rewritten to use AGPR form");
52
53/// Map from spill slot frame index to list of instructions which reference it.
54using SpillReferenceMap = DenseMap<int, SmallVector<MachineInstr *, 4>>;
55
56class AMDGPURewriteAGPRCopyMFMAImpl {
57 MachineFunction &MF;
58 const GCNSubtarget &ST;
59 const SIInstrInfo &TII;
60 const SIRegisterInfo &TRI;
61 MachineRegisterInfo &MRI;
62 VirtRegMap &VRM;
63 LiveRegMatrix &LRM;
64 LiveIntervals &LIS;
65 LiveStacks &LSS;
66 const RegisterClassInfo &RegClassInfo;
67 MachineDominatorTree &MDT;
68
69 bool attemptReassignmentsToAGPR(SmallSetVector<Register, 4> &InterferingRegs,
70 MCPhysReg PrefPhysReg) const;
71
72public:
73 AMDGPURewriteAGPRCopyMFMAImpl(MachineFunction &MF, VirtRegMap &VRM,
74 LiveRegMatrix &LRM, LiveIntervals &LIS,
75 LiveStacks &LSS,
76 const RegisterClassInfo &RegClassInfo,
77 MachineDominatorTree &MDT)
78 : MF(MF), ST(MF.getSubtarget<GCNSubtarget>()), TII(*ST.getInstrInfo()),
79 TRI(*ST.getRegisterInfo()), MRI(MF.getRegInfo()), VRM(VRM), LRM(LRM),
80 LIS(LIS), LSS(LSS), RegClassInfo(RegClassInfo), MDT(MDT) {}
81
82 bool isRewriteCandidate(const MachineInstr &MI) const {
83 return TII.isMAI(MI) && AMDGPU::getAGPRFormOp(Opcode: MI.getOpcode()) != -1;
84 }
85
86 /// Find AV_* registers assigned to AGPRs (or virtual registers which were
87 /// already required to be AGPR).
88 ///
89 /// \return the assigned physical register that \p VReg is assigned to if it
90 /// is an AGPR, otherwise MCRegister().
91 MCRegister getAssignedAGPR(Register VReg) const {
92 MCRegister PhysReg = VRM.getPhys(virtReg: VReg);
93 if (!PhysReg)
94 return MCRegister();
95
96 // If this is an AV register, we have to check if the actual assignment is
97 // to an AGPR
98 const TargetRegisterClass *AssignedRC = TRI.getPhysRegBaseClass(Reg: PhysReg);
99 return TRI.isAGPRClass(RC: AssignedRC) ? PhysReg : MCRegister();
100 }
101
102 bool tryReassigningMFMAChain(MachineInstr &MFMA, Register MFMAHintReg,
103 MCPhysReg PhysRegHint) const;
104
105 /// Compute the register class constraints based on the uses of \p Reg,
106 /// excluding MFMA uses from which can be rewritten to change the register
107 /// class constraint. MFMA scale operands need to be constraint checked.
108 /// This should be nearly identical to MachineRegisterInfo::recomputeRegClass.
109
110 /// \p RewriteCandidates will collect the set of MFMA instructions that need
111 /// to have the opcode mutated to perform the replacement.
112 ///
113 /// \p RewriteRegs will accumulate the set of register used by those MFMAs
114 /// that need to have the register classes adjusted.
115 bool recomputeRegClassExceptRewritable(
116 Register Reg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
117 SmallSetVector<Register, 4> &RewriteRegs) const;
118
119 bool tryFoldCopiesToAGPR(Register VReg, MCRegister AssignedAGPR) const;
120 bool tryFoldCopiesFromAGPR(Register VReg, MCRegister AssignedAGPR) const;
121
122 /// Replace spill instruction \p SpillMI which loads/stores from/to \p SpillFI
123 /// with a COPY to the replacement register value \p VReg.
124 void replaceSpillWithCopyToVReg(MachineInstr &SpillMI, int SpillFI,
125 Register VReg) const;
126
127 /// Create a map from frame index to use instructions for spills. If a use of
128 /// the frame index does not consist only of spill instructions, it will not
129 /// be included in the map.
130 void collectSpillIndexUses(ArrayRef<LiveInterval *> StackIntervals,
131 SpillReferenceMap &Map) const;
132
133 /// Return true if the reload \p LoadMI of the stack slot with live interval
134 /// \p SlotLI is jointly dominated by the slot's spill stores, i.e. every path
135 /// from the entry block to the load passes through a store to the slot before
136 /// the load. \p StoreFreeReachable is the set of blocks reachable from the
137 /// entry block without passing through any store block for the slot.
138 bool isLoadJointlyDominatedByStores(
139 const MachineInstr &LoadMI, const LiveInterval &SlotLI,
140 const SmallPtrSetImpl<MachineBasicBlock *> &StoreFreeReachable) const;
141
142 /// Attempt to unspill VGPRs by finding a free register and replacing the
143 /// spill instructions with copies.
144 void eliminateSpillsOfReassignedVGPRs() const;
145
146 bool run(MachineFunction &MF) const;
147};
148
149bool AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
150 Register StartReg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
151 SmallSetVector<Register, 4> &RewriteRegs) const {
152 SmallVector<Register, 8> Worklist = {StartReg};
153
154 // Recursively visit all transitive MFMA users
155 while (!Worklist.empty()) {
156 Register Reg = Worklist.pop_back_val();
157 const TargetRegisterClass *OldRC = MRI.getRegClass(Reg);
158
159 // Inflate to the equivalent AV_* class.
160 const TargetRegisterClass *NewRC = TRI.getLargestLegalSuperClass(RC: OldRC, MF);
161 if (OldRC == NewRC)
162 return false;
163
164 // Accumulate constraints from all uses.
165 for (MachineOperand &MO : MRI.reg_nodbg_operands(Reg)) {
166 // Apply the effect of the given operand to NewRC.
167 MachineInstr *MI = MO.getParent();
168
169 // We can swap the classes of dst + src2 as a pair to AGPR, so ignore the
170 // effects of rewrite candidates. It just so happens that we can use
171 // either AGPR or VGPR in src0/src1. We still need to check constraint
172 // effects for scale variant, which does not allow AGPR.
173 if (isRewriteCandidate(MI: *MI)) {
174 int AGPROp = AMDGPU::getAGPRFormOp(Opcode: MI->getOpcode());
175 const MCInstrDesc &AGPRDesc = TII.get(Opcode: AGPROp);
176 const TargetRegisterClass *NewRC =
177 TII.getRegClass(MCID: AGPRDesc, OpNum: MO.getOperandNo());
178 if (!TRI.hasAGPRs(RC: NewRC))
179 return false;
180
181 const MachineOperand *VDst =
182 TII.getNamedOperand(MI&: *MI, OperandName: AMDGPU::OpName::vdst);
183 const MachineOperand *Src2 =
184 TII.getNamedOperand(MI&: *MI, OperandName: AMDGPU::OpName::src2);
185 for (const MachineOperand *Op : {VDst, Src2}) {
186 if (!Op->isReg())
187 continue;
188
189 Register OtherReg = Op->getReg();
190 if (OtherReg.isPhysical())
191 return false;
192
193 if (OtherReg != Reg && RewriteRegs.insert(X: OtherReg))
194 Worklist.push_back(Elt: OtherReg);
195 }
196
197 if (!is_contained(Range&: RewriteCandidates, Element: MI)) {
198 LLVM_DEBUG({
199 Register VDstPhysReg = VRM.getPhys(VDst->getReg());
200 dbgs() << "Attempting to replace VGPR MFMA with AGPR version:"
201 << " Dst=[" << printReg(VDst->getReg()) << " => "
202 << printReg(VDstPhysReg, &TRI);
203
204 if (Src2->isReg()) {
205 Register Src2PhysReg = VRM.getPhys(Src2->getReg());
206 dbgs() << "], Src2=[" << printReg(Src2->getReg(), &TRI) << " => "
207 << printReg(Src2PhysReg, &TRI);
208 }
209
210 dbgs() << "]: " << MI;
211 });
212
213 RewriteCandidates.push_back(Elt: MI);
214 }
215
216 continue;
217 }
218
219 unsigned OpNo = &MO - &MI->getOperand(i: 0);
220 NewRC = MI->getRegClassConstraintEffect(OpIdx: OpNo, CurRC: NewRC, TII: &TII, TRI: &TRI);
221 if (!NewRC || NewRC == OldRC) {
222 LLVM_DEBUG(dbgs() << "User of " << printReg(Reg, &TRI)
223 << " cannot be reassigned to "
224 << (NewRC ? TRI.getRegClassName(NewRC) : "NULL")
225 << ": " << *MI);
226 return false;
227 }
228 }
229 }
230
231 return true;
232}
233
234bool AMDGPURewriteAGPRCopyMFMAImpl::tryReassigningMFMAChain(
235 MachineInstr &MFMA, Register MFMAHintReg, MCPhysReg PhysRegHint) const {
236 // src2 and dst have the same physical class constraint; try to preserve
237 // the original src2 subclass if one were to exist.
238 SmallVector<MachineInstr *, 4> RewriteCandidates = {&MFMA};
239 SmallSetVector<Register, 4> RewriteRegs;
240
241 // Make sure we reassign the MFMA we found the copy from first. We want
242 // to ensure dst ends up in the physreg we were originally copying to.
243 RewriteRegs.insert(X: MFMAHintReg);
244
245 // We've found av = COPY (MFMA) (or MFMA (v = COPY av)) and need to verify
246 // that we can trivially rewrite src2 to use the new AGPR. If we can't
247 // trivially replace it, we're going to induce as many copies as we would have
248 // emitted in the first place, as well as need to assign another register, and
249 // need to figure out where to put them. The live range splitting is smarter
250 // than anything we're doing here, so trust it did something reasonable.
251 //
252 // Note recomputeRegClassExceptRewritable will consider the constraints of
253 // this MFMA's src2 as well as the src2/dst of any transitive MFMA users.
254 if (!recomputeRegClassExceptRewritable(StartReg: MFMAHintReg, RewriteCandidates,
255 RewriteRegs)) {
256 LLVM_DEBUG(dbgs() << "Could not recompute the regclass of dst reg "
257 << printReg(MFMAHintReg, &TRI) << '\n');
258 return false;
259 }
260
261 // If src2 and dst are different registers, we need to also reassign the
262 // input to an available AGPR if it is compatible with all other uses.
263 //
264 // If we can't reassign it, we'd need to introduce a different copy
265 // which is likely worse than the copy we'd be saving.
266 //
267 // It's likely that the MFMA is used in sequence with other MFMAs; if we
268 // cannot migrate the full use/def chain of MFMAs, we would need to
269 // introduce intermediate copies somewhere. So we only make the
270 // transform if all the interfering MFMAs can also be migrated. Collect
271 // the set of rewritable MFMAs and check if we can assign an AGPR at
272 // that point.
273 //
274 // If any of the MFMAs aren't reassignable, we give up and rollback to
275 // the original register assignments.
276
277 using RecoloringStack =
278 SmallVector<std::pair<const LiveInterval *, MCRegister>, 8>;
279 RecoloringStack TentativeReassignments;
280
281 for (Register RewriteReg : RewriteRegs) {
282 LiveInterval &LI = LIS.getInterval(Reg: RewriteReg);
283 TentativeReassignments.push_back(Elt: {&LI, VRM.getPhys(virtReg: RewriteReg)});
284 LRM.unassign(VirtReg: LI);
285 }
286
287 if (!DebugCounter::shouldExecute(Counter&: RewriteAGPRCopyMFMACounter) ||
288 !attemptReassignmentsToAGPR(InterferingRegs&: RewriteRegs, PrefPhysReg: PhysRegHint)) {
289 // Roll back the register assignments to the original state.
290 for (auto [LI, OldAssign] : TentativeReassignments) {
291 if (VRM.hasPhys(virtReg: LI->reg()))
292 LRM.unassign(VirtReg: *LI);
293 LRM.assign(VirtReg: *LI, PhysReg: OldAssign);
294 }
295
296 return false;
297 }
298
299 // Fixup the register classes of the virtual registers now that we've
300 // committed to the reassignments.
301 for (Register InterferingReg : RewriteRegs) {
302 const TargetRegisterClass *EquivalentAGPRRegClass =
303 TRI.getEquivalentAGPRClass(SRC: MRI.getRegClass(Reg: InterferingReg));
304 MRI.setRegClass(Reg: InterferingReg, RC: EquivalentAGPRRegClass);
305 }
306
307 for (MachineInstr *RewriteCandidate : RewriteCandidates) {
308 int NewMFMAOp = AMDGPU::getAGPRFormOp(Opcode: RewriteCandidate->getOpcode());
309 RewriteCandidate->setDesc(TII.get(Opcode: NewMFMAOp));
310 ++NumMFMAsRewrittenToAGPR;
311 }
312
313 return true;
314}
315
316/// Attempt to reassign the registers in \p InterferingRegs to be AGPRs, with a
317/// preference to use \p PhysReg first. Returns false if the reassignments
318/// cannot be trivially performed.
319bool AMDGPURewriteAGPRCopyMFMAImpl::attemptReassignmentsToAGPR(
320 SmallSetVector<Register, 4> &InterferingRegs, MCPhysReg PrefPhysReg) const {
321 // FIXME: The ordering may matter here, but we're just taking uselistorder
322 // with the special case of ensuring to process the starting instruction
323 // first. We probably should extract the priority advisor out of greedy and
324 // use that ordering.
325 for (Register InterferingReg : InterferingRegs) {
326 LiveInterval &ReassignLI = LIS.getInterval(Reg: InterferingReg);
327 const TargetRegisterClass *EquivalentAGPRRegClass =
328 TRI.getEquivalentAGPRClass(SRC: MRI.getRegClass(Reg: InterferingReg));
329
330 MCPhysReg Assignable = AMDGPU::NoRegister;
331 if (EquivalentAGPRRegClass->contains(Reg: PrefPhysReg) &&
332 LRM.checkInterference(VirtReg: ReassignLI, PhysReg: PrefPhysReg) ==
333 LiveRegMatrix::IK_Free) {
334 // First try to assign to the AGPR we were already copying to. This
335 // should be the first assignment we attempt. We have to guard
336 // against the use being a subregister (which doesn't have an exact
337 // class match).
338
339 // TODO: If this does happen to be a subregister use, we should
340 // still try to assign to a subregister of the original copy result.
341 Assignable = PrefPhysReg;
342 } else {
343 ArrayRef<MCPhysReg> AllocOrder =
344 RegClassInfo.getOrder(RC: EquivalentAGPRRegClass);
345 for (MCPhysReg Reg : AllocOrder) {
346 if (LRM.checkInterference(VirtReg: ReassignLI, PhysReg: Reg) == LiveRegMatrix::IK_Free) {
347 Assignable = Reg;
348 break;
349 }
350 }
351 }
352
353 if (!Assignable) {
354 LLVM_DEBUG(dbgs() << "Unable to reassign VGPR "
355 << printReg(InterferingReg, &TRI)
356 << " to a free AGPR\n");
357 return false;
358 }
359
360 LLVM_DEBUG(dbgs() << "Reassigning VGPR " << printReg(InterferingReg, &TRI)
361 << " to " << printReg(Assignable, &TRI) << '\n');
362 LRM.assign(VirtReg: ReassignLI, PhysReg: Assignable);
363 }
364
365 return true;
366}
367
368/// Identify copies that look like:
369/// %vdst:vgpr = V_MFMA_.. %src0:av, %src1:av, %src2:vgpr
370/// %agpr = COPY %vgpr
371///
372/// Then try to replace the transitive uses of %src2 and %vdst with the AGPR
373/// versions of the MFMA. This should cover the common case.
374bool AMDGPURewriteAGPRCopyMFMAImpl::tryFoldCopiesToAGPR(
375 Register VReg, MCRegister AssignedAGPR) const {
376 bool MadeChange = false;
377 for (MachineInstr &UseMI : MRI.def_instructions(Reg: VReg)) {
378 if (!UseMI.isCopy())
379 continue;
380
381 Register CopySrcReg = UseMI.getOperand(i: 1).getReg();
382 if (!CopySrcReg.isVirtual())
383 continue;
384
385 // TODO: Handle loop phis copied to AGPR. e.g.
386 //
387 // loop:
388 // %phi:vgpr = COPY %mfma:vgpr
389 // %mfma:vgpr = V_MFMA_xxx_vgprcd_e64 %a, %b, %phi
390 // s_cbranch_vccnz loop
391 //
392 // endloop:
393 // %agpr = mfma
394 //
395 // We need to be sure that %phi is assigned to the same physical register as
396 // %mfma, or else we will just be moving copies into the loop.
397
398 for (MachineInstr &CopySrcDefMI : MRI.def_instructions(Reg: CopySrcReg)) {
399 if (isRewriteCandidate(MI: CopySrcDefMI) &&
400 tryReassigningMFMAChain(
401 MFMA&: CopySrcDefMI, MFMAHintReg: CopySrcDefMI.getOperand(i: 0).getReg(), PhysRegHint: AssignedAGPR))
402 MadeChange = true;
403 }
404 }
405
406 return MadeChange;
407}
408
409/// Identify copies that look like:
410/// %src:vgpr = COPY %src:agpr
411/// %vdst:vgpr = V_MFMA_... %src0:av, %src1:av, %src:vgpr
412///
413/// Then try to replace the transitive uses of %src2 and %vdst with the AGPR
414/// versions of the MFMA. This should cover rarer cases, and will generally be
415/// redundant with tryFoldCopiesToAGPR.
416bool AMDGPURewriteAGPRCopyMFMAImpl::tryFoldCopiesFromAGPR(
417 Register VReg, MCRegister AssignedAGPR) const {
418 bool MadeChange = false;
419 for (MachineInstr &UseMI : MRI.use_instructions(Reg: VReg)) {
420 if (!UseMI.isCopy())
421 continue;
422
423 Register CopyDstReg = UseMI.getOperand(i: 0).getReg();
424 if (!CopyDstReg.isVirtual())
425 continue;
426 for (MachineOperand &CopyUseMO : MRI.reg_nodbg_operands(Reg: CopyDstReg)) {
427 if (!CopyUseMO.readsReg())
428 continue;
429
430 MachineInstr &CopyUseMI = *CopyUseMO.getParent();
431 if (isRewriteCandidate(MI: CopyUseMI)) {
432 if (tryReassigningMFMAChain(MFMA&: CopyUseMI, MFMAHintReg: CopyDstReg,
433 PhysRegHint: VRM.getPhys(virtReg: CopyDstReg)))
434 MadeChange = true;
435 }
436 }
437 }
438
439 return MadeChange;
440}
441
442void AMDGPURewriteAGPRCopyMFMAImpl::replaceSpillWithCopyToVReg(
443 MachineInstr &SpillMI, int SpillFI, Register VReg) const {
444 const DebugLoc &DL = SpillMI.getDebugLoc();
445 MachineBasicBlock &MBB = *SpillMI.getParent();
446 MachineInstr *NewCopy;
447 if (SpillMI.mayStore()) {
448 NewCopy = BuildMI(BB&: MBB, I&: SpillMI, MIMD: DL, MCID: TII.get(Opcode: TargetOpcode::COPY), DestReg: VReg)
449 .add(MO: SpillMI.getOperand(i: 0));
450 } else {
451 NewCopy = BuildMI(BB&: MBB, I&: SpillMI, MIMD: DL, MCID: TII.get(Opcode: TargetOpcode::COPY))
452 .add(MO: SpillMI.getOperand(i: 0))
453 .addReg(RegNo: VReg);
454 }
455
456 LIS.ReplaceMachineInstrInMaps(MI&: SpillMI, NewMI&: *NewCopy);
457 SpillMI.eraseFromParent();
458}
459
460void AMDGPURewriteAGPRCopyMFMAImpl::collectSpillIndexUses(
461 ArrayRef<LiveInterval *> StackIntervals, SpillReferenceMap &Map) const {
462
463 SmallSet<int, 4> NeededFrameIndexes;
464 for (const LiveInterval *LI : StackIntervals)
465 NeededFrameIndexes.insert(V: LI->reg().stackSlotIndex());
466
467 for (MachineBasicBlock &MBB : MF) {
468 for (MachineInstr &MI : MBB) {
469 for (MachineOperand &MO : MI.operands()) {
470 if (!MO.isFI() || !NeededFrameIndexes.count(V: MO.getIndex()))
471 continue;
472
473 if (TII.isVGPRSpill(MI)) {
474 SmallVector<MachineInstr *, 4> &References = Map[MO.getIndex()];
475 References.push_back(Elt: &MI);
476 break;
477 }
478
479 // Verify this was really a spill instruction, if it's not just ignore
480 // all uses.
481
482 // TODO: This should probably be verifier enforced.
483 NeededFrameIndexes.erase(V: MO.getIndex());
484 Map.erase(Val: MO.getIndex());
485 }
486 }
487 }
488}
489
490bool AMDGPURewriteAGPRCopyMFMAImpl::isLoadJointlyDominatedByStores(
491 const MachineInstr &LoadMI, const LiveInterval &SlotLI,
492 const SmallPtrSetImpl<MachineBasicBlock *> &StoreFreeReachable) const {
493 const MachineBasicBlock *LoadMBB = LoadMI.getParent();
494 if (!MDT.isReachableFromEntry(A: LoadMBB))
495 return true;
496
497 // Check if every path passed through a store block.
498 if (!StoreFreeReachable.contains(Ptr: LoadMBB))
499 return true;
500
501 // Otherwise, there exists a path to this block that has not seen any store
502 // yet. We must ensure that within this block there is a store to this slot
503 // before the load. Consult the slot's LiveStacks interval: a store to the
504 // slot before the load means the slot is not live into this block but is
505 // live at the load. If the load reads an undef value, the slot is not live
506 // at the load, failing the joint-dominance check.
507 SlotIndex LoadIdx = LIS.getInstructionIndex(Instr: LoadMI);
508 return SlotLI.liveAt(index: LoadIdx) && !LIS.isLiveInToMBB(LR: SlotLI, mbb: LoadMBB);
509}
510
511void AMDGPURewriteAGPRCopyMFMAImpl::eliminateSpillsOfReassignedVGPRs() const {
512 unsigned NumSlots = LSS.getNumIntervals();
513 if (NumSlots == 0)
514 return;
515
516 MachineFrameInfo &MFI = MF.getFrameInfo();
517
518 SmallVector<LiveInterval *, 32> StackIntervals;
519 StackIntervals.reserve(N: NumSlots);
520
521 for (auto &[Slot, LI] : LSS) {
522 if (!MFI.isSpillSlotObjectIndex(ObjectIdx: Slot) || MFI.isDeadObjectIndex(ObjectIdx: Slot))
523 continue;
524
525 const TargetRegisterClass *RC = LSS.getIntervalRegClass(Slot);
526 if (TRI.hasVGPRs(RC))
527 StackIntervals.push_back(Elt: &LI);
528 }
529
530 sort(C&: StackIntervals, Comp: [](const LiveInterval *A, const LiveInterval *B) {
531 // The ordering has to be strictly weak.
532 /// Sort heaviest intervals first to prioritize their unspilling
533 if (A->weight() != B->weight())
534 return A->weight() > B->weight();
535
536 if (A->getSize() != B->getSize())
537 return A->getSize() > B->getSize();
538
539 // Tie breaker by number to avoid need for stable sort
540 return A->reg().stackSlotIndex() < B->reg().stackSlotIndex();
541 });
542
543 // FIXME: The APIs for dealing with the LiveInterval of a frame index are
544 // cumbersome. LiveStacks owns its LiveIntervals which refer to stack
545 // slots. We cannot use the usual LiveRegMatrix::assign and unassign on these,
546 // and must create a substitute virtual register to do so. This makes
547 // incremental updating here difficult; we need to actually perform the IR
548 // mutation to get the new vreg references in place to compute the register
549 // LiveInterval to perform an assignment to track the new interference
550 // correctly, and we can't simply migrate the LiveInterval we already have.
551 //
552 // To avoid walking through the entire function for each index, pre-collect
553 // all the instructions slot referencess.
554
555 DenseMap<int, SmallVector<MachineInstr *, 4>> SpillSlotReferences;
556 collectSpillIndexUses(StackIntervals, Map&: SpillSlotReferences);
557
558 for (LiveInterval *LI : StackIntervals) {
559 int Slot = LI->reg().stackSlotIndex();
560 auto SpillReferences = SpillSlotReferences.find(Val: Slot);
561 if (SpillReferences == SpillSlotReferences.end())
562 continue;
563
564 // For each spill reload, every path from entry to the reload must pass
565 // through at least one spill store to the same stack slot.
566 SmallPtrSet<MachineBasicBlock *, 4> StoreBlocks;
567 for (MachineInstr *MI : SpillReferences->second) {
568 if (MI->mayStore() && MDT.isReachableFromEntry(A: MI->getParent()))
569 StoreBlocks.insert(Ptr: MI->getParent());
570 }
571
572 if (StoreBlocks.empty()) {
573 LLVM_DEBUG(dbgs() << "Skipping " << printReg(Slot, &TRI)
574 << ": no reachable stores\n");
575 continue;
576 }
577
578 // Compute blocks reachable from entry without passing through a store
579 // block.
580 MachineBasicBlock &EntryMBB = MF.front();
581 SmallPtrSet<MachineBasicBlock *, 16> StoreFreeReachable = {&EntryMBB};
582 SmallVector<MachineBasicBlock *, 16> Worklist = {&EntryMBB};
583
584 while (!Worklist.empty()) {
585 MachineBasicBlock *MBB = Worklist.pop_back_val();
586 if (StoreBlocks.contains(Ptr: MBB))
587 continue;
588
589 for (MachineBasicBlock *Succ : MBB->successors()) {
590 if (StoreFreeReachable.insert(Ptr: Succ).second)
591 Worklist.push_back(Elt: Succ);
592 }
593 }
594
595 // Every reachable reload must be jointly dominated by the slot's stores.
596 if (!llvm::all_of(Range&: SpillReferences->second, P: [&](const MachineInstr *MI) {
597 return !MI->mayLoad() ||
598 isLoadJointlyDominatedByStores(LoadMI: *MI, SlotLI: *LI, StoreFreeReachable);
599 })) {
600 LLVM_DEBUG(
601 dbgs() << "Skipping " << printReg(Slot, &TRI)
602 << ": some reachable load not jointly dominated by stores\n");
603 continue;
604 }
605
606 const TargetRegisterClass *RC = LSS.getIntervalRegClass(Slot);
607
608 LLVM_DEBUG(dbgs() << "Trying to eliminate " << printReg(Slot, &TRI)
609 << " by reassigning\n");
610
611 ArrayRef<MCPhysReg> AllocOrder = RegClassInfo.getOrder(RC);
612
613 for (MCPhysReg PhysReg : AllocOrder) {
614 if (LRM.checkInterference(VirtReg: *LI, PhysReg) != LiveRegMatrix::IK_Free)
615 continue;
616
617 LLVM_DEBUG(dbgs() << "Reassigning " << *LI << " to "
618 << printReg(PhysReg, &TRI) << '\n');
619
620 const TargetRegisterClass *RC = LSS.getIntervalRegClass(Slot);
621 Register NewVReg = MRI.createVirtualRegister(RegClass: RC);
622
623 for (MachineInstr *SpillMI : SpillReferences->second)
624 replaceSpillWithCopyToVReg(SpillMI&: *SpillMI, SpillFI: Slot, VReg: NewVReg);
625
626 // TODO: We should be able to transfer the information from the stack
627 // slot's LiveInterval without recomputing from scratch with the
628 // replacement vreg uses.
629 LiveInterval &NewLI = LIS.createAndComputeVirtRegInterval(Reg: NewVReg);
630 VRM.grow();
631
632 // A spill slot can be stored to multiple times, so the replacement
633 // vreg may have multiple disconnected live range components. Split
634 // them into separate vregs to maintain the single-component invariant.
635 SmallVector<LiveInterval *, 4> SplitLIs;
636 LIS.splitSeparateComponents(LI&: NewLI, SplitLIs);
637
638 LLVM_DEBUG({
639 if (!SplitLIs.empty()) {
640 dbgs() << "Split unspilled interval into " << (SplitLIs.size() + 1)
641 << " components\n";
642 }
643 });
644
645 LRM.assign(VirtReg: NewLI, PhysReg);
646 for (LiveInterval *SplitLI : SplitLIs) {
647 VRM.grow();
648 LRM.assign(VirtReg: *SplitLI, PhysReg);
649 }
650
651 MFI.RemoveStackObject(ObjectIdx: Slot);
652 break;
653 }
654 }
655}
656
657bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
658 // This only applies on subtargets that have a configurable AGPR vs. VGPR
659 // allocation.
660 if (!ST.hasGFX90AInsts())
661 return false;
662
663 // Early exit if no AGPRs were assigned.
664 if (!LRM.isPhysRegUsed(PhysReg: AMDGPU::AGPR0)) {
665 LLVM_DEBUG(dbgs() << "skipping function that did not allocate AGPRs\n");
666 return false;
667 }
668
669 bool MadeChange = false;
670
671 for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
672 Register VReg = Register::index2VirtReg(Index: I);
673 MCRegister AssignedAGPR = getAssignedAGPR(VReg);
674 if (!AssignedAGPR)
675 continue;
676
677 if (tryFoldCopiesToAGPR(VReg, AssignedAGPR))
678 MadeChange = true;
679 if (tryFoldCopiesFromAGPR(VReg, AssignedAGPR))
680 MadeChange = true;
681 }
682
683 // If we've successfully rewritten some MFMAs, we've alleviated some VGPR
684 // pressure. See if we can eliminate some spills now that those registers are
685 // more available.
686 if (MadeChange)
687 eliminateSpillsOfReassignedVGPRs();
688
689 return MadeChange;
690}
691
692class AMDGPURewriteAGPRCopyMFMALegacy : public MachineFunctionPass {
693public:
694 static char ID;
695 RegisterClassInfo RegClassInfo;
696
697 AMDGPURewriteAGPRCopyMFMALegacy() : MachineFunctionPass(ID) {}
698
699 bool runOnMachineFunction(MachineFunction &MF) override;
700
701 StringRef getPassName() const override {
702 return "AMDGPU Rewrite AGPR-Copy-MFMA";
703 }
704
705 void getAnalysisUsage(AnalysisUsage &AU) const override {
706 AU.addRequired<LiveIntervalsWrapperPass>();
707 AU.addRequired<VirtRegMapWrapperLegacy>();
708 AU.addRequired<LiveRegMatrixWrapperLegacy>();
709 AU.addRequired<LiveStacksWrapperLegacy>();
710 AU.addRequired<MachineDominatorTreeWrapperPass>();
711
712 AU.addPreserved<LiveIntervalsWrapperPass>();
713 AU.addPreserved<VirtRegMapWrapperLegacy>();
714 AU.addPreserved<LiveRegMatrixWrapperLegacy>();
715 AU.addPreserved<LiveStacksWrapperLegacy>();
716 AU.addPreserved<MachineDominatorTreeWrapperPass>();
717
718 AU.setPreservesAll();
719 MachineFunctionPass::getAnalysisUsage(AU);
720 }
721};
722
723} // End anonymous namespace.
724
725INITIALIZE_PASS_BEGIN(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE,
726 "AMDGPU Rewrite AGPR-Copy-MFMA", false, false)
727INITIALIZE_PASS_DEPENDENCY(LiveIntervalsWrapperPass)
728INITIALIZE_PASS_DEPENDENCY(VirtRegMapWrapperLegacy)
729INITIALIZE_PASS_DEPENDENCY(LiveRegMatrixWrapperLegacy)
730INITIALIZE_PASS_DEPENDENCY(LiveStacksWrapperLegacy)
731INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
732INITIALIZE_PASS_END(AMDGPURewriteAGPRCopyMFMALegacy, DEBUG_TYPE,
733 "AMDGPU Rewrite AGPR-Copy-MFMA", false, false)
734
735char AMDGPURewriteAGPRCopyMFMALegacy::ID = 0;
736
737char &llvm::AMDGPURewriteAGPRCopyMFMALegacyID =
738 AMDGPURewriteAGPRCopyMFMALegacy::ID;
739
740bool AMDGPURewriteAGPRCopyMFMALegacy::runOnMachineFunction(
741 MachineFunction &MF) {
742 if (skipFunction(F: MF.getFunction()))
743 return false;
744
745 RegClassInfo.runOnMachineFunction(MF);
746
747 auto &VRM = getAnalysis<VirtRegMapWrapperLegacy>().getVRM();
748 auto &LRM = getAnalysis<LiveRegMatrixWrapperLegacy>().getLRM();
749 auto &LIS = getAnalysis<LiveIntervalsWrapperPass>().getLIS();
750 auto &LSS = getAnalysis<LiveStacksWrapperLegacy>().getLS();
751 auto &MDT = getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
752 AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS, LSS, RegClassInfo, MDT);
753 return Impl.run(MF);
754}
755
756PreservedAnalyses
757AMDGPURewriteAGPRCopyMFMAPass::run(MachineFunction &MF,
758 MachineFunctionAnalysisManager &MFAM) {
759 VirtRegMap &VRM = MFAM.getResult<VirtRegMapAnalysis>(IR&: MF);
760 LiveRegMatrix &LRM = MFAM.getResult<LiveRegMatrixAnalysis>(IR&: MF);
761 LiveIntervals &LIS = MFAM.getResult<LiveIntervalsAnalysis>(IR&: MF);
762 LiveStacks &LSS = MFAM.getResult<LiveStacksAnalysis>(IR&: MF);
763 MachineDominatorTree &MDT = MFAM.getResult<MachineDominatorTreeAnalysis>(IR&: MF);
764 RegisterClassInfo RegClassInfo;
765 RegClassInfo.runOnMachineFunction(MF);
766
767 AMDGPURewriteAGPRCopyMFMAImpl Impl(MF, VRM, LRM, LIS, LSS, RegClassInfo, MDT);
768 if (!Impl.run(MF))
769 return PreservedAnalyses::all();
770 auto PA = getMachineFunctionPassPreservedAnalyses();
771 PA.preserveSet<CFGAnalyses>()
772 .preserve<LiveStacksAnalysis>()
773 .preserve<VirtRegMapAnalysis>()
774 .preserve<SlotIndexesAnalysis>()
775 .preserve<LiveIntervalsAnalysis>()
776 .preserve<LiveRegMatrixAnalysis>();
777 return PA;
778}
779