1 //===- AMDGPUSetWavePriority.cpp - Set wave priority ----------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 /// \file
10 /// Pass to temporarily raise the wave priority beginning the start of
11 /// the shader function until its last VMEM instructions to allow younger
12 /// waves to issue their VMEM instructions as well.
13 //
14 //===----------------------------------------------------------------------===//
15 
16 #include "AMDGPU.h"
17 #include "GCNSubtarget.h"
18 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
19 #include "SIInstrInfo.h"
20 #include "llvm/ADT/PostOrderIterator.h"
21 #include "llvm/CodeGen/MachineFunctionPass.h"
22 #include "llvm/InitializePasses.h"
23 #include "llvm/Support/Allocator.h"
24 
25 using namespace llvm;
26 
27 #define DEBUG_TYPE "amdgpu-set-wave-priority"
28 
29 namespace {
30 
31 struct MBBInfo {
32   MBBInfo() = default;
33   bool MayReachVMEMLoad = false;
34 };
35 
36 using MBBInfoSet = DenseMap<const MachineBasicBlock *, MBBInfo>;
37 
38 class AMDGPUSetWavePriority : public MachineFunctionPass {
39 public:
40   static char ID;
41 
42   AMDGPUSetWavePriority() : MachineFunctionPass(ID) {}
43 
44   StringRef getPassName() const override { return "Set wave priority"; }
45 
46   bool runOnMachineFunction(MachineFunction &MF) override;
47 
48 private:
49   MachineInstr *BuildSetprioMI(MachineFunction &MF, unsigned priority) const;
50 
51   const SIInstrInfo *TII;
52 };
53 
54 } // End anonymous namespace.
55 
56 INITIALIZE_PASS(AMDGPUSetWavePriority, DEBUG_TYPE, "Set wave priority", false,
57                 false)
58 
59 char AMDGPUSetWavePriority::ID = 0;
60 
61 FunctionPass *llvm::createAMDGPUSetWavePriorityPass() {
62   return new AMDGPUSetWavePriority();
63 }
64 
65 MachineInstr *AMDGPUSetWavePriority::BuildSetprioMI(MachineFunction &MF,
66                                                     unsigned priority) const {
67   return BuildMI(MF, DebugLoc(), TII->get(AMDGPU::S_SETPRIO)).addImm(priority);
68 }
69 
70 // Checks that for every predecessor Pred that can reach a VMEM load,
71 // none of Pred's successors can reach a VMEM load.
72 static bool CanLowerPriorityDirectlyInPredecessors(const MachineBasicBlock &MBB,
73                                                    MBBInfoSet &MBBInfos) {
74   for (const MachineBasicBlock *Pred : MBB.predecessors()) {
75     if (!MBBInfos[Pred].MayReachVMEMLoad)
76       continue;
77     for (const MachineBasicBlock *Succ : Pred->successors()) {
78       if (MBBInfos[Succ].MayReachVMEMLoad)
79         return false;
80     }
81   }
82   return true;
83 }
84 
85 static bool isVMEMLoad(const MachineInstr &MI) {
86   return SIInstrInfo::isVMEM(MI) && MI.mayLoad();
87 }
88 
89 bool AMDGPUSetWavePriority::runOnMachineFunction(MachineFunction &MF) {
90   const unsigned HighPriority = 3;
91   const unsigned LowPriority = 0;
92 
93   Function &F = MF.getFunction();
94   if (skipFunction(F) || !AMDGPU::isEntryFunctionCC(F.getCallingConv()))
95     return false;
96 
97   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
98   TII = ST.getInstrInfo();
99 
100   MBBInfoSet MBBInfos;
101   SmallVector<const MachineBasicBlock *, 16> Worklist;
102   for (MachineBasicBlock &MBB : MF) {
103     if (any_of(MBB, isVMEMLoad))
104       Worklist.push_back(&MBB);
105   }
106 
107   // Mark blocks from which control may reach VMEM loads.
108   while (!Worklist.empty()) {
109     const MachineBasicBlock *MBB = Worklist.pop_back_val();
110     MBBInfo &Info = MBBInfos[MBB];
111     if (!Info.MayReachVMEMLoad) {
112       Info.MayReachVMEMLoad = true;
113       Worklist.append(MBB->pred_begin(), MBB->pred_end());
114     }
115   }
116 
117   MachineBasicBlock &Entry = MF.front();
118   if (!MBBInfos[&Entry].MayReachVMEMLoad)
119     return false;
120 
121   // Raise the priority at the beginning of the shader.
122   MachineBasicBlock::iterator I = Entry.begin(), E = Entry.end();
123   while (I != E && !SIInstrInfo::isVALU(*I) && !I->isTerminator())
124     ++I;
125   Entry.insert(I, BuildSetprioMI(MF, HighPriority));
126 
127   // Lower the priority on edges where control leaves blocks from which
128   // VMEM loads are reachable.
129   SmallSet<MachineBasicBlock *, 16> PriorityLoweringBlocks;
130   for (MachineBasicBlock &MBB : MF) {
131     if (MBBInfos[&MBB].MayReachVMEMLoad) {
132       if (MBB.succ_empty())
133         PriorityLoweringBlocks.insert(&MBB);
134       continue;
135     }
136 
137     if (CanLowerPriorityDirectlyInPredecessors(MBB, MBBInfos)) {
138       for (MachineBasicBlock *Pred : MBB.predecessors()) {
139         if (MBBInfos[Pred].MayReachVMEMLoad)
140           PriorityLoweringBlocks.insert(Pred);
141       }
142       continue;
143     }
144 
145     // Where lowering the priority in predecessors is not possible, the
146     // block receiving control either was not part of a loop in the first
147     // place or the loop simplification/canonicalization pass should have
148     // already tried to split the edge and insert a preheader, and if for
149     // whatever reason it failed to do so, then this leaves us with the
150     // only option of lowering the priority within the loop.
151     PriorityLoweringBlocks.insert(&MBB);
152   }
153 
154   for (MachineBasicBlock *MBB : PriorityLoweringBlocks) {
155     MachineBasicBlock::iterator I = MBB->end(), B = MBB->begin();
156     while (I != B) {
157       if (isVMEMLoad(*--I)) {
158         ++I;
159         break;
160       }
161     }
162     MBB->insert(I, BuildSetprioMI(MF, LowPriority));
163   }
164 
165   return true;
166 }
167