1//===-- AMDGPUPromoteKernelArguments.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 This pass recursively promotes generic pointer arguments of a kernel
10/// into the global address space.
11///
12/// The pass walks kernel's pointer arguments, then loads from them. If a loaded
13/// value is a pointer and loaded pointer is unmodified in the kernel before the
14/// load, then promote loaded pointer to global. Then recursively continue.
15//
16//===----------------------------------------------------------------------===//
17
18#include "AMDGPU.h"
19#include "AMDGPUMemoryUtils.h"
20#include "llvm/ADT/SmallVector.h"
21#include "llvm/Analysis/AliasAnalysis.h"
22#include "llvm/Analysis/MemorySSA.h"
23#include "llvm/IR/IRBuilder.h"
24#include "llvm/InitializePasses.h"
25
26#define DEBUG_TYPE "amdgpu-promote-kernel-arguments"
27
28using namespace llvm;
29
30namespace {
31
32class AMDGPUPromoteKernelArguments : public FunctionPass {
33 MemorySSA *MSSA;
34
35 AliasAnalysis *AA;
36
37 Instruction *ArgCastInsertPt;
38
39 SmallVector<Value *> Ptrs;
40
41 void enqueueUsers(Value *Ptr);
42
43 bool promotePointer(Value *Ptr);
44
45 bool promoteLoad(LoadInst *LI);
46
47public:
48 static char ID;
49
50 AMDGPUPromoteKernelArguments() : FunctionPass(ID) {}
51
52 bool run(Function &F, MemorySSA &MSSA, AliasAnalysis &AA);
53
54 bool runOnFunction(Function &F) override;
55
56 void getAnalysisUsage(AnalysisUsage &AU) const override {
57 AU.addRequired<AAResultsWrapperPass>();
58 AU.addRequired<MemorySSAWrapperPass>();
59 AU.setPreservesAll();
60 }
61};
62
63} // end anonymous namespace
64
65void AMDGPUPromoteKernelArguments::enqueueUsers(Value *Ptr) {
66 SmallVector<User *> PtrUsers(Ptr->users());
67
68 while (!PtrUsers.empty()) {
69 Instruction *U = dyn_cast<Instruction>(Val: PtrUsers.pop_back_val());
70 if (!U)
71 continue;
72
73 switch (U->getOpcode()) {
74 default:
75 break;
76 case Instruction::Load: {
77 LoadInst *LD = cast<LoadInst>(Val: U);
78 if (LD->getPointerOperand()->stripInBoundsOffsets() == Ptr &&
79 !AMDGPU::isClobberedInFunction(Load: LD, MSSA, AA))
80 Ptrs.push_back(Elt: LD);
81
82 break;
83 }
84 case Instruction::GetElementPtr:
85 case Instruction::AddrSpaceCast:
86 case Instruction::BitCast:
87 if (U->getOperand(i: 0)->stripInBoundsOffsets() == Ptr)
88 PtrUsers.append(in_start: U->user_begin(), in_end: U->user_end());
89 break;
90 }
91 }
92}
93
94bool AMDGPUPromoteKernelArguments::promotePointer(Value *Ptr) {
95 bool Changed = false;
96
97 LoadInst *LI = dyn_cast<LoadInst>(Val: Ptr);
98 if (LI)
99 Changed |= promoteLoad(LI);
100
101 PointerType *PT = dyn_cast<PointerType>(Val: Ptr->getType());
102 if (!PT)
103 return Changed;
104
105 if (PT->getAddressSpace() == AMDGPUAS::FLAT_ADDRESS ||
106 PT->getAddressSpace() == AMDGPUAS::GLOBAL_ADDRESS ||
107 PT->getAddressSpace() == AMDGPUAS::CONSTANT_ADDRESS)
108 enqueueUsers(Ptr);
109
110 if (PT->getAddressSpace() != AMDGPUAS::FLAT_ADDRESS)
111 return Changed;
112
113 // The promotion is valid because in the HSA offload model the host populates
114 // kernarg slots at dispatch time and can only supply global-aperture
115 // addresses. Any flat pointer reachable from a kernel argument can therefore
116 // be assumed to be in the global aperture. This also mirrors the assumption
117 // in getAssumedAddrSpace (AMDGPUTargetMachine.cpp) which independently
118 // promotes flat kernel arguments to global via InferAddressSpaces.
119 IRBuilder<> B(LI ? &*std::next(x: cast<Instruction>(Val: Ptr)->getIterator())
120 : ArgCastInsertPt);
121
122 // Cast pointer to global address space and back to flat and let
123 // Infer Address Spaces pass to do all necessary rewriting.
124 PointerType *NewPT =
125 PointerType::get(C&: PT->getContext(), AddressSpace: AMDGPUAS::GLOBAL_ADDRESS);
126 Value *Cast =
127 B.CreateAddrSpaceCast(V: Ptr, DestTy: NewPT, Name: Twine(Ptr->getName(), ".global"));
128 Value *CastBack =
129 B.CreateAddrSpaceCast(V: Cast, DestTy: PT, Name: Twine(Ptr->getName(), ".flat"));
130 Ptr->replaceUsesWithIf(New: CastBack,
131 ShouldReplace: [Cast](Use &U) { return U.getUser() != Cast; });
132
133 return true;
134}
135
136bool AMDGPUPromoteKernelArguments::promoteLoad(LoadInst *LI) {
137 if (!LI->isSimple())
138 return false;
139
140 LI->setMetadata(Kind: "amdgpu.noclobber", Node: MDNode::get(Context&: LI->getContext(), MDs: {}));
141 return true;
142}
143
144// skip allocas
145static BasicBlock::iterator getInsertPt(BasicBlock &BB) {
146 BasicBlock::iterator InsPt = BB.getFirstInsertionPt();
147 for (BasicBlock::iterator E = BB.end(); InsPt != E; ++InsPt) {
148 AllocaInst *AI = dyn_cast<AllocaInst>(Val: &*InsPt);
149
150 // If this is a dynamic alloca, the value may depend on the loaded kernargs,
151 // so loads will need to be inserted before it.
152 if (!AI || !AI->isStaticAlloca())
153 break;
154 }
155
156 return InsPt;
157}
158
159bool AMDGPUPromoteKernelArguments::run(Function &F, MemorySSA &MSSA,
160 AliasAnalysis &AA) {
161 if (skipFunction(F))
162 return false;
163
164 CallingConv::ID CC = F.getCallingConv();
165 if (CC != CallingConv::AMDGPU_KERNEL || F.arg_empty())
166 return false;
167
168 ArgCastInsertPt = &*getInsertPt(BB&: *F.begin());
169 this->MSSA = &MSSA;
170 this->AA = &AA;
171
172 for (Argument &Arg : F.args()) {
173 if (Arg.use_empty())
174 continue;
175
176 PointerType *PT = dyn_cast<PointerType>(Val: Arg.getType());
177 if (!PT || (PT->getAddressSpace() != AMDGPUAS::FLAT_ADDRESS &&
178 PT->getAddressSpace() != AMDGPUAS::GLOBAL_ADDRESS &&
179 PT->getAddressSpace() != AMDGPUAS::CONSTANT_ADDRESS))
180 continue;
181
182 Ptrs.push_back(Elt: &Arg);
183 }
184
185 bool Changed = false;
186 while (!Ptrs.empty()) {
187 Value *Ptr = Ptrs.pop_back_val();
188 Changed |= promotePointer(Ptr);
189 }
190
191 return Changed;
192}
193
194bool AMDGPUPromoteKernelArguments::runOnFunction(Function &F) {
195 MemorySSA &MSSA = getAnalysis<MemorySSAWrapperPass>().getMSSA();
196 AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
197 return run(F, MSSA, AA);
198}
199
200INITIALIZE_PASS_BEGIN(AMDGPUPromoteKernelArguments, DEBUG_TYPE,
201 "AMDGPU Promote Kernel Arguments", false, false)
202INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
203INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
204INITIALIZE_PASS_END(AMDGPUPromoteKernelArguments, DEBUG_TYPE,
205 "AMDGPU Promote Kernel Arguments", false, false)
206
207char AMDGPUPromoteKernelArguments::ID = 0;
208
209FunctionPass *llvm::createAMDGPUPromoteKernelArgumentsPass() {
210 return new AMDGPUPromoteKernelArguments();
211}
212
213PreservedAnalyses
214AMDGPUPromoteKernelArgumentsPass::run(Function &F,
215 FunctionAnalysisManager &AM) {
216 MemorySSA &MSSA = AM.getResult<MemorySSAAnalysis>(IR&: F).getMSSA();
217 AliasAnalysis &AA = AM.getResult<AAManager>(IR&: F);
218 if (AMDGPUPromoteKernelArguments().run(F, MSSA, AA)) {
219 PreservedAnalyses PA;
220 PA.preserveSet<CFGAnalyses>();
221 PA.preserve<MemorySSAAnalysis>();
222 return PA;
223 }
224 return PreservedAnalyses::all();
225}
226