1//===- AMDGPURewriteOutArgumentsPass.cpp - Create struct returns ----------===//
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 attempts to replace out argument usage with a return of a
10/// struct.
11///
12/// We can support returning a lot of values directly in registers, but
13/// idiomatic C code frequently uses a pointer argument to return a second value
14/// rather than returning a struct by value. GPU stack access is also quite
15/// painful, so we want to avoid that if possible. Passing a stack object
16/// pointer to a function also requires an additional address expansion code
17/// sequence to convert the pointer to be relative to the kernel's scratch wave
18/// offset register since the callee doesn't know what stack frame the incoming
19/// pointer is relative to.
20///
21/// The goal is to try rewriting code that looks like this:
22///
23/// int foo(int a, int b, int* out) {
24/// *out = bar();
25/// return a + b;
26/// }
27///
28/// into something like this:
29///
30/// std::pair<int, int> foo(int a, int b) {
31/// return std::pair(a + b, bar());
32/// }
33///
34/// Typically the incoming pointer is a simple alloca for a temporary variable
35/// to use the API, which if replaced with a struct return will be easily SROA'd
36/// out when the stub function we create is inlined
37///
38/// This pass introduces the struct return, but leaves the unused pointer
39/// arguments and introduces a new stub function calling the struct returning
40/// body. DeadArgumentElimination should be run after this to clean these up.
41//
42//===----------------------------------------------------------------------===//
43
44#include "AMDGPU.h"
45#include "Utils/AMDGPUBaseInfo.h"
46#include "llvm/ADT/Statistic.h"
47#include "llvm/Analysis/MemorySSA.h"
48#include "llvm/Analysis/MemorySSAUpdater.h"
49#include "llvm/IR/AttributeMask.h"
50#include "llvm/IR/IRBuilder.h"
51#include "llvm/IR/Instructions.h"
52#include "llvm/InitializePasses.h"
53#include "llvm/Pass.h"
54#include "llvm/Support/CommandLine.h"
55#include "llvm/Support/Debug.h"
56#include "llvm/Support/raw_ostream.h"
57
58#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
59
60using namespace llvm;
61
62static cl::opt<bool> AnyAddressSpace(
63 "amdgpu-any-address-space-out-arguments",
64 cl::desc("Replace pointer out arguments with "
65 "struct returns for non-private address space"),
66 cl::Hidden,
67 cl::init(Val: false));
68
69static cl::opt<unsigned> MaxNumRetRegs(
70 "amdgpu-max-return-arg-num-regs",
71 cl::desc("Approximately limit number of return registers for replacing out arguments"),
72 cl::Hidden,
73 cl::init(Val: 16));
74
75STATISTIC(NumOutArgumentsReplaced,
76 "Number out arguments moved to struct return values");
77STATISTIC(NumOutArgumentFunctionsReplaced,
78 "Number of functions with out arguments moved to struct return values");
79
80namespace {
81
82class AMDGPURewriteOutArguments : public FunctionPass {
83private:
84 const DataLayout *DL = nullptr;
85 MemorySSA *MSSA = nullptr;
86 MemorySSAUpdater *MSSAU = nullptr;
87 AAResults *AA = nullptr;
88
89 Type *getStoredType(Value &Arg) const;
90 Type *getOutArgumentType(Argument &Arg) const;
91
92public:
93 static char ID;
94
95 AMDGPURewriteOutArguments() : FunctionPass(ID) {}
96
97 void getAnalysisUsage(AnalysisUsage &AU) const override {
98 AU.addRequired<MemorySSAWrapperPass>();
99 AU.addRequired<AAResultsWrapperPass>();
100 FunctionPass::getAnalysisUsage(AU);
101 }
102
103 bool doInitialization(Module &M) override;
104 bool runOnFunction(Function &F) override;
105};
106
107} // end anonymous namespace
108
109INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
110 "AMDGPU Rewrite Out Arguments", false, false)
111INITIALIZE_PASS_DEPENDENCY(MemorySSAWrapperPass)
112INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
113 "AMDGPU Rewrite Out Arguments", false, false)
114
115char AMDGPURewriteOutArguments::ID = 0;
116
117Type *AMDGPURewriteOutArguments::getStoredType(Value &Arg) const {
118 const int MaxUses = 10;
119 int UseCount = 0;
120
121 SmallVector<Use *> Worklist(llvm::make_pointer_range(Range: Arg.uses()));
122
123 Type *StoredType = nullptr;
124 while (!Worklist.empty()) {
125 Use *U = Worklist.pop_back_val();
126
127 if (auto *BCI = dyn_cast<BitCastInst>(Val: U->getUser())) {
128 for (Use &U : BCI->uses())
129 Worklist.push_back(Elt: &U);
130 continue;
131 }
132
133 if (auto *SI = dyn_cast<StoreInst>(Val: U->getUser())) {
134 if (UseCount++ > MaxUses)
135 return nullptr;
136
137 if (!SI->isSimple() ||
138 U->getOperandNo() != StoreInst::getPointerOperandIndex())
139 return nullptr;
140
141 if (StoredType && StoredType != SI->getValueOperand()->getType())
142 return nullptr; // More than one type.
143 StoredType = SI->getValueOperand()->getType();
144 continue;
145 }
146
147 // Unsupported user.
148 return nullptr;
149 }
150
151 return StoredType;
152}
153
154Type *AMDGPURewriteOutArguments::getOutArgumentType(Argument &Arg) const {
155 const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
156 PointerType *ArgTy = dyn_cast<PointerType>(Val: Arg.getType());
157
158 // TODO: It might be useful for any out arguments, not just privates.
159 if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
160 !AnyAddressSpace) ||
161 Arg.hasByValAttr() || Arg.hasStructRetAttr()) {
162 return nullptr;
163 }
164
165 Type *StoredType = getStoredType(Arg);
166 if (!StoredType || DL->getTypeStoreSize(Ty: StoredType) > MaxOutArgSizeBytes)
167 return nullptr;
168
169 return StoredType;
170}
171
172bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
173 DL = &M.getDataLayout();
174 return false;
175}
176
177static StoreInst *findStoreForOutArgument(BasicBlock *BB, Argument *OutArg,
178 MemorySSA &MSSA,
179 BatchAAResults &BAA) {
180 MemoryLocation ArgLoc = MemoryLocation::getBeforeOrAfter(Ptr: OutArg);
181 const auto *Accesses = MSSA.getBlockAccesses(BB);
182 if (!Accesses)
183 return nullptr;
184
185 for (const MemoryAccess &Access : reverse(C: *Accesses)) {
186 const auto *UseOrDef = dyn_cast<MemoryUseOrDef>(Val: &Access);
187 if (!UseOrDef)
188 continue;
189
190 Instruction *I = UseOrDef->getMemoryInst();
191
192 // Return the must-alias store to the out argument.
193 if (auto *Store = dyn_cast<StoreInst>(Val: I))
194 if (Store->getPointerOperand() == OutArg)
195 return Store;
196
197 if (auto *FI = dyn_cast<FenceInst>(Val: I))
198 if (FI->getOrdering() == AtomicOrdering::Release)
199 continue;
200
201 if (auto *LI = dyn_cast<LoadInst>(Val: I)) {
202 if (LI->isAtomic()) {
203 // May-alias reads with monotonic ordering are ignored.
204 if (isStrongerThan(AO: LI->getOrdering(), Other: AtomicOrdering::Monotonic))
205 return nullptr;
206 continue;
207 }
208 }
209
210 // Any other memory access that writes the location prevents the
211 // rewrite.
212 // FIXME: should handle aliasing reads too.
213 if (isModSet(MRI: BAA.getModRefInfo(I, OptLoc: ArgLoc)))
214 return nullptr;
215 }
216
217 return nullptr;
218}
219
220bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
221 if (skipFunction(F))
222 return false;
223
224 // TODO: Could probably handle variadic functions.
225 if (F.isVarArg() || F.hasStructRetAttr() ||
226 AMDGPU::isEntryFunctionCC(CC: F.getCallingConv()))
227 return false;
228
229 unsigned ReturnNumRegs = 0;
230 // Maps an out-argument number to its field index in the return struct.
231 // Fields are in processing order, which the retry loop below can reorder
232 // relative to argument order, so the index must be tracked, not recomputed.
233 SmallDenseMap<unsigned, unsigned, 4> OutArgIndexes;
234 SmallVector<Type *, 4> ReturnTypes;
235 Type *RetTy = F.getReturnType();
236 if (!RetTy->isVoidTy()) {
237 ReturnNumRegs = DL->getTypeStoreSize(Ty: RetTy) / 4;
238
239 if (ReturnNumRegs >= MaxNumRetRegs)
240 return false;
241
242 ReturnTypes.push_back(Elt: RetTy);
243 }
244
245 SmallVector<std::pair<Argument *, Type *>, 4> OutArgs;
246 for (Argument &Arg : F.args()) {
247 if (Type *Ty = getOutArgumentType(Arg)) {
248 LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg
249 << " in function " << F.getName() << '\n');
250 OutArgs.push_back(Elt: {&Arg, Ty});
251 }
252 }
253
254 if (OutArgs.empty())
255 return false;
256
257 using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
258
259 DenseMap<ReturnInst *, ReplacementVec> Replacements;
260
261 SmallVector<ReturnInst *, 4> Returns;
262 for (BasicBlock &BB : F) {
263 if (ReturnInst *RI = dyn_cast<ReturnInst>(Val: &BB.back()))
264 Returns.push_back(Elt: RI);
265 }
266
267 if (Returns.empty())
268 return false;
269
270 AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
271 MSSA = &getAnalysis<MemorySSAWrapperPass>().getMSSA();
272
273 BatchAAResults BatchAA(*AA);
274 MemorySSAUpdater MSSAUpdater(MSSA);
275 MSSAU = &MSSAUpdater;
276
277 bool Changing;
278
279 do {
280 Changing = false;
281
282 // Keep retrying if we are able to successfully eliminate an argument. This
283 // helps with cases with multiple arguments which may alias, such as in a
284 // sincos implementation. With 2 stores to may-aliasing arguments, MDA
285 // returns the second store for the first argument too; the identity guard
286 // below rejects it, and a later iteration folds the first argument once the
287 // second store has been removed.
288 for (const auto &Pair : OutArgs) {
289 bool ThisReplaceable = true;
290 SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
291
292 Argument *OutArg = Pair.first;
293 Type *ArgTy = Pair.second;
294
295 // Skip this argument if converting it will push us over the register
296 // count to return limit.
297
298 // TODO: This is an approximation. When legalized this could be more. We
299 // can ask TLI for exactly how many.
300 unsigned ArgNumRegs = DL->getTypeStoreSize(Ty: ArgTy) / 4;
301 if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
302 continue;
303
304 // An argument is convertible only if all exit blocks are able to replace
305 // it.
306 for (ReturnInst *RI : Returns) {
307 BasicBlock *BB = RI->getParent();
308
309 StoreInst *SI = findStoreForOutArgument(BB, OutArg, MSSA&: *MSSA, BAA&: BatchAA);
310 if (SI) {
311 LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
312 ReplaceableStores.emplace_back(Args&: RI, Args&: SI);
313 } else {
314 ThisReplaceable = false;
315 break;
316 }
317 }
318
319 if (!ThisReplaceable)
320 continue; // Try the next argument candidate.
321
322 for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
323 Value *ReplVal = Store.second->getValueOperand();
324
325 auto &ValVec = Replacements[Store.first];
326 if (llvm::is_contained(Range: llvm::make_first_range(c&: ValVec), Element: OutArg)) {
327 LLVM_DEBUG(dbgs()
328 << "Saw multiple out arg stores" << *OutArg << '\n');
329 // It is possible to see stores to the same argument multiple times,
330 // but we expect these would have been optimized out already.
331 ThisReplaceable = false;
332 break;
333 }
334
335 ValVec.emplace_back(Args&: OutArg, Args&: ReplVal);
336 MSSAU->removeMemoryAccess(I: Store.second);
337 Store.second->eraseFromParent();
338 }
339
340 if (ThisReplaceable) {
341 OutArgIndexes.insert(KV: {OutArg->getArgNo(), ReturnTypes.size()});
342 ReturnTypes.push_back(Elt: ArgTy);
343 ++NumOutArgumentsReplaced;
344 Changing = true;
345 }
346 }
347 } while (Changing);
348
349 if (Replacements.empty())
350 return false;
351
352 LLVMContext &Ctx = F.getContext();
353 StructType *NewRetTy = StructType::create(Context&: Ctx, Elements: ReturnTypes, Name: F.getName());
354
355 FunctionType *NewFuncTy = FunctionType::get(Result: NewRetTy,
356 Params: F.getFunctionType()->params(),
357 isVarArg: F.isVarArg());
358
359 LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
360
361 Function *NewFunc = Function::Create(Ty: NewFuncTy, Linkage: Function::PrivateLinkage,
362 N: F.getName() + ".body");
363 F.getParent()->getFunctionList().insert(where: F.getIterator(), New: NewFunc);
364 NewFunc->copyAttributesFrom(Src: &F);
365 NewFunc->setComdat(F.getComdat());
366
367 // We want to preserve the function and param attributes, but need to strip
368 // off any return attributes, e.g. zeroext doesn't make sense with a struct.
369 NewFunc->stealArgumentListFrom(Src&: F);
370
371 AttributeMask RetAttrs;
372 RetAttrs.addAttribute(Val: Attribute::SExt);
373 RetAttrs.addAttribute(Val: Attribute::ZExt);
374 RetAttrs.addAttribute(Val: Attribute::NoAlias);
375 NewFunc->removeRetAttrs(Attrs: RetAttrs);
376 // TODO: How to preserve metadata?
377
378 // Move the body of the function into the new rewritten function, and replace
379 // this function with a stub.
380 NewFunc->splice(ToIt: NewFunc->begin(), FromF: &F);
381
382 for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
383 ReturnInst *RI = Replacement.first;
384 IRBuilder<> B(RI);
385 B.SetCurrentDebugLocation(RI->getDebugLoc());
386
387 Value *NewRetVal = PoisonValue::get(T: NewRetTy);
388
389 Value *RetVal = RI->getReturnValue();
390 if (RetVal)
391 NewRetVal = B.CreateInsertValue(Agg: NewRetVal, Val: RetVal, Idxs: 0);
392
393 // Use OutArgIndexes so body and stub agree on the field for each argument.
394 for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
395 unsigned FieldIdx = OutArgIndexes.lookup(Val: ReturnPoint.first->getArgNo());
396 NewRetVal = B.CreateInsertValue(Agg: NewRetVal, Val: ReturnPoint.second, Idxs: FieldIdx);
397 }
398
399 if (RetVal)
400 RI->setOperand(i_nocapture: 0, Val_nocapture: NewRetVal);
401 else {
402 B.CreateRet(V: NewRetVal);
403 RI->eraseFromParent();
404 }
405 }
406
407 SmallVector<Value *, 16> StubCallArgs;
408 for (Argument &Arg : F.args()) {
409 if (OutArgIndexes.count(Val: Arg.getArgNo())) {
410 // It's easier to preserve the type of the argument list. We rely on
411 // DeadArgumentElimination to take care of these.
412 StubCallArgs.push_back(Elt: PoisonValue::get(T: Arg.getType()));
413 } else {
414 StubCallArgs.push_back(Elt: &Arg);
415 }
416 }
417
418 BasicBlock *StubBB = BasicBlock::Create(Context&: Ctx, Name: "", Parent: &F);
419 IRBuilder<> B(StubBB);
420 CallInst *StubCall = B.CreateCall(Callee: NewFunc, Args: StubCallArgs);
421
422 for (Argument &Arg : F.args()) {
423 auto It = OutArgIndexes.find(Val: Arg.getArgNo());
424 if (It == OutArgIndexes.end())
425 continue;
426
427 unsigned FieldIdx = It->second;
428 Type *EltTy = NewRetTy->getElementType(N: FieldIdx);
429 const auto Align =
430 DL->getValueOrABITypeAlignment(Alignment: Arg.getParamAlign(), Ty: EltTy);
431
432 Value *Val = B.CreateExtractValue(Agg: StubCall, Idxs: FieldIdx);
433 B.CreateAlignedStore(Val, Ptr: &Arg, Align);
434 }
435
436 if (!RetTy->isVoidTy()) {
437 B.CreateRet(V: B.CreateExtractValue(Agg: StubCall, Idxs: 0));
438 } else {
439 B.CreateRetVoid();
440 }
441
442 // The function is now a stub we want to inline.
443 F.addFnAttr(Kind: Attribute::AlwaysInline);
444
445 ++NumOutArgumentFunctionsReplaced;
446 return true;
447}
448
449FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
450 return new AMDGPURewriteOutArguments();
451}
452