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/MemoryDependenceAnalysis.h"
48#include "llvm/IR/AttributeMask.h"
49#include "llvm/IR/IRBuilder.h"
50#include "llvm/IR/Instructions.h"
51#include "llvm/InitializePasses.h"
52#include "llvm/Pass.h"
53#include "llvm/Support/CommandLine.h"
54#include "llvm/Support/Debug.h"
55#include "llvm/Support/raw_ostream.h"
56
57#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
58
59using namespace llvm;
60
61static cl::opt<bool> AnyAddressSpace(
62 "amdgpu-any-address-space-out-arguments",
63 cl::desc("Replace pointer out arguments with "
64 "struct returns for non-private address space"),
65 cl::Hidden,
66 cl::init(Val: false));
67
68static cl::opt<unsigned> MaxNumRetRegs(
69 "amdgpu-max-return-arg-num-regs",
70 cl::desc("Approximately limit number of return registers for replacing out arguments"),
71 cl::Hidden,
72 cl::init(Val: 16));
73
74STATISTIC(NumOutArgumentsReplaced,
75 "Number out arguments moved to struct return values");
76STATISTIC(NumOutArgumentFunctionsReplaced,
77 "Number of functions with out arguments moved to struct return values");
78
79namespace {
80
81class AMDGPURewriteOutArguments : public FunctionPass {
82private:
83 const DataLayout *DL = nullptr;
84 MemoryDependenceResults *MDA = nullptr;
85
86 Type *getStoredType(Value &Arg) const;
87 Type *getOutArgumentType(Argument &Arg) const;
88
89public:
90 static char ID;
91
92 AMDGPURewriteOutArguments() : FunctionPass(ID) {}
93
94 void getAnalysisUsage(AnalysisUsage &AU) const override {
95 AU.addRequired<MemoryDependenceWrapperPass>();
96 FunctionPass::getAnalysisUsage(AU);
97 }
98
99 bool doInitialization(Module &M) override;
100 bool runOnFunction(Function &F) override;
101};
102
103} // end anonymous namespace
104
105INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
106 "AMDGPU Rewrite Out Arguments", false, false)
107INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
108INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
109 "AMDGPU Rewrite Out Arguments", false, false)
110
111char AMDGPURewriteOutArguments::ID = 0;
112
113Type *AMDGPURewriteOutArguments::getStoredType(Value &Arg) const {
114 const int MaxUses = 10;
115 int UseCount = 0;
116
117 SmallVector<Use *> Worklist(llvm::make_pointer_range(Range: Arg.uses()));
118
119 Type *StoredType = nullptr;
120 while (!Worklist.empty()) {
121 Use *U = Worklist.pop_back_val();
122
123 if (auto *BCI = dyn_cast<BitCastInst>(Val: U->getUser())) {
124 for (Use &U : BCI->uses())
125 Worklist.push_back(Elt: &U);
126 continue;
127 }
128
129 if (auto *SI = dyn_cast<StoreInst>(Val: U->getUser())) {
130 if (UseCount++ > MaxUses)
131 return nullptr;
132
133 if (!SI->isSimple() ||
134 U->getOperandNo() != StoreInst::getPointerOperandIndex())
135 return nullptr;
136
137 if (StoredType && StoredType != SI->getValueOperand()->getType())
138 return nullptr; // More than one type.
139 StoredType = SI->getValueOperand()->getType();
140 continue;
141 }
142
143 // Unsupported user.
144 return nullptr;
145 }
146
147 return StoredType;
148}
149
150Type *AMDGPURewriteOutArguments::getOutArgumentType(Argument &Arg) const {
151 const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
152 PointerType *ArgTy = dyn_cast<PointerType>(Val: Arg.getType());
153
154 // TODO: It might be useful for any out arguments, not just privates.
155 if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
156 !AnyAddressSpace) ||
157 Arg.hasByValAttr() || Arg.hasStructRetAttr()) {
158 return nullptr;
159 }
160
161 Type *StoredType = getStoredType(Arg);
162 if (!StoredType || DL->getTypeStoreSize(Ty: StoredType) > MaxOutArgSizeBytes)
163 return nullptr;
164
165 return StoredType;
166}
167
168bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
169 DL = &M.getDataLayout();
170 return false;
171}
172
173bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
174 if (skipFunction(F))
175 return false;
176
177 // TODO: Could probably handle variadic functions.
178 if (F.isVarArg() || F.hasStructRetAttr() ||
179 AMDGPU::isEntryFunctionCC(CC: F.getCallingConv()))
180 return false;
181
182 MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
183
184 unsigned ReturnNumRegs = 0;
185 // Maps an out-argument number to its field index in the return struct.
186 // Fields are in processing order, which the retry loop below can reorder
187 // relative to argument order, so the index must be tracked, not recomputed.
188 SmallDenseMap<unsigned, unsigned, 4> OutArgIndexes;
189 SmallVector<Type *, 4> ReturnTypes;
190 Type *RetTy = F.getReturnType();
191 if (!RetTy->isVoidTy()) {
192 ReturnNumRegs = DL->getTypeStoreSize(Ty: RetTy) / 4;
193
194 if (ReturnNumRegs >= MaxNumRetRegs)
195 return false;
196
197 ReturnTypes.push_back(Elt: RetTy);
198 }
199
200 SmallVector<std::pair<Argument *, Type *>, 4> OutArgs;
201 for (Argument &Arg : F.args()) {
202 if (Type *Ty = getOutArgumentType(Arg)) {
203 LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg
204 << " in function " << F.getName() << '\n');
205 OutArgs.push_back(Elt: {&Arg, Ty});
206 }
207 }
208
209 if (OutArgs.empty())
210 return false;
211
212 using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
213
214 DenseMap<ReturnInst *, ReplacementVec> Replacements;
215
216 SmallVector<ReturnInst *, 4> Returns;
217 for (BasicBlock &BB : F) {
218 if (ReturnInst *RI = dyn_cast<ReturnInst>(Val: &BB.back()))
219 Returns.push_back(Elt: RI);
220 }
221
222 if (Returns.empty())
223 return false;
224
225 bool Changing;
226
227 do {
228 Changing = false;
229
230 // Keep retrying if we are able to successfully eliminate an argument. This
231 // helps with cases with multiple arguments which may alias, such as in a
232 // sincos implementation. With 2 stores to may-aliasing arguments, MDA
233 // returns the second store for the first argument too; the identity guard
234 // below rejects it, and a later iteration folds the first argument once the
235 // second store has been removed.
236 for (const auto &Pair : OutArgs) {
237 bool ThisReplaceable = true;
238 SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
239
240 Argument *OutArg = Pair.first;
241 Type *ArgTy = Pair.second;
242
243 // Skip this argument if converting it will push us over the register
244 // count to return limit.
245
246 // TODO: This is an approximation. When legalized this could be more. We
247 // can ask TLI for exactly how many.
248 unsigned ArgNumRegs = DL->getTypeStoreSize(Ty: ArgTy) / 4;
249 if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
250 continue;
251
252 // An argument is convertible only if all exit blocks are able to replace
253 // it.
254 for (ReturnInst *RI : Returns) {
255 BasicBlock *BB = RI->getParent();
256
257 MemDepResult Q = MDA->getPointerDependencyFrom(
258 Loc: MemoryLocation::getBeforeOrAfter(Ptr: OutArg), isLoad: true, ScanIt: BB->end(), BB, QueryInst: RI);
259 StoreInst *SI = nullptr;
260 if (Q.isDef())
261 SI = dyn_cast<StoreInst>(Val: Q.getInst());
262
263 // MDA stops at the first may-aliasing store, which need not be to this
264 // argument; only fold a store whose pointer is exactly OutArg.
265 if (SI && SI->getPointerOperand() != OutArg)
266 SI = nullptr;
267
268 if (SI) {
269 LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
270 ReplaceableStores.emplace_back(Args&: RI, Args&: SI);
271 } else {
272 ThisReplaceable = false;
273 break;
274 }
275 }
276
277 if (!ThisReplaceable)
278 continue; // Try the next argument candidate.
279
280 for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
281 Value *ReplVal = Store.second->getValueOperand();
282
283 auto &ValVec = Replacements[Store.first];
284 if (llvm::is_contained(Range: llvm::make_first_range(c&: ValVec), Element: OutArg)) {
285 LLVM_DEBUG(dbgs()
286 << "Saw multiple out arg stores" << *OutArg << '\n');
287 // It is possible to see stores to the same argument multiple times,
288 // but we expect these would have been optimized out already.
289 ThisReplaceable = false;
290 break;
291 }
292
293 ValVec.emplace_back(Args&: OutArg, Args&: ReplVal);
294 Store.second->eraseFromParent();
295 }
296
297 if (ThisReplaceable) {
298 OutArgIndexes.insert(KV: {OutArg->getArgNo(), ReturnTypes.size()});
299 ReturnTypes.push_back(Elt: ArgTy);
300 ++NumOutArgumentsReplaced;
301 Changing = true;
302 }
303 }
304 } while (Changing);
305
306 if (Replacements.empty())
307 return false;
308
309 LLVMContext &Ctx = F.getContext();
310 StructType *NewRetTy = StructType::create(Context&: Ctx, Elements: ReturnTypes, Name: F.getName());
311
312 FunctionType *NewFuncTy = FunctionType::get(Result: NewRetTy,
313 Params: F.getFunctionType()->params(),
314 isVarArg: F.isVarArg());
315
316 LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
317
318 Function *NewFunc = Function::Create(Ty: NewFuncTy, Linkage: Function::PrivateLinkage,
319 N: F.getName() + ".body");
320 F.getParent()->getFunctionList().insert(where: F.getIterator(), New: NewFunc);
321 NewFunc->copyAttributesFrom(Src: &F);
322 NewFunc->setComdat(F.getComdat());
323
324 // We want to preserve the function and param attributes, but need to strip
325 // off any return attributes, e.g. zeroext doesn't make sense with a struct.
326 NewFunc->stealArgumentListFrom(Src&: F);
327
328 AttributeMask RetAttrs;
329 RetAttrs.addAttribute(Val: Attribute::SExt);
330 RetAttrs.addAttribute(Val: Attribute::ZExt);
331 RetAttrs.addAttribute(Val: Attribute::NoAlias);
332 NewFunc->removeRetAttrs(Attrs: RetAttrs);
333 // TODO: How to preserve metadata?
334
335 // Move the body of the function into the new rewritten function, and replace
336 // this function with a stub.
337 NewFunc->splice(ToIt: NewFunc->begin(), FromF: &F);
338
339 for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
340 ReturnInst *RI = Replacement.first;
341 IRBuilder<> B(RI);
342 B.SetCurrentDebugLocation(RI->getDebugLoc());
343
344 Value *NewRetVal = PoisonValue::get(T: NewRetTy);
345
346 Value *RetVal = RI->getReturnValue();
347 if (RetVal)
348 NewRetVal = B.CreateInsertValue(Agg: NewRetVal, Val: RetVal, Idxs: 0);
349
350 // Use OutArgIndexes so body and stub agree on the field for each argument.
351 for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
352 unsigned FieldIdx = OutArgIndexes.lookup(Val: ReturnPoint.first->getArgNo());
353 NewRetVal = B.CreateInsertValue(Agg: NewRetVal, Val: ReturnPoint.second, Idxs: FieldIdx);
354 }
355
356 if (RetVal)
357 RI->setOperand(i_nocapture: 0, Val_nocapture: NewRetVal);
358 else {
359 B.CreateRet(V: NewRetVal);
360 RI->eraseFromParent();
361 }
362 }
363
364 SmallVector<Value *, 16> StubCallArgs;
365 for (Argument &Arg : F.args()) {
366 if (OutArgIndexes.count(Val: Arg.getArgNo())) {
367 // It's easier to preserve the type of the argument list. We rely on
368 // DeadArgumentElimination to take care of these.
369 StubCallArgs.push_back(Elt: PoisonValue::get(T: Arg.getType()));
370 } else {
371 StubCallArgs.push_back(Elt: &Arg);
372 }
373 }
374
375 BasicBlock *StubBB = BasicBlock::Create(Context&: Ctx, Name: "", Parent: &F);
376 IRBuilder<> B(StubBB);
377 CallInst *StubCall = B.CreateCall(Callee: NewFunc, Args: StubCallArgs);
378
379 for (Argument &Arg : F.args()) {
380 auto It = OutArgIndexes.find(Val: Arg.getArgNo());
381 if (It == OutArgIndexes.end())
382 continue;
383
384 unsigned FieldIdx = It->second;
385 Type *EltTy = NewRetTy->getElementType(N: FieldIdx);
386 const auto Align =
387 DL->getValueOrABITypeAlignment(Alignment: Arg.getParamAlign(), Ty: EltTy);
388
389 Value *Val = B.CreateExtractValue(Agg: StubCall, Idxs: FieldIdx);
390 B.CreateAlignedStore(Val, Ptr: &Arg, Align);
391 }
392
393 if (!RetTy->isVoidTy()) {
394 B.CreateRet(V: B.CreateExtractValue(Agg: StubCall, Idxs: 0));
395 } else {
396 B.CreateRetVoid();
397 }
398
399 // The function is now a stub we want to inline.
400 F.addFnAttr(Kind: Attribute::AlwaysInline);
401
402 ++NumOutArgumentFunctionsReplaced;
403 return true;
404}
405
406FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
407 return new AMDGPURewriteOutArguments();
408}
409