LLVM 18.0.0git
AMDGPURewriteOutArguments.cpp
Go to the documentation of this file.
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"
46#include "llvm/ADT/SmallSet.h"
47#include "llvm/ADT/Statistic.h"
50#include "llvm/IR/IRBuilder.h"
53#include "llvm/Pass.h"
55#include "llvm/Support/Debug.h"
57
58#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
59
60using namespace llvm;
61
63 "amdgpu-any-address-space-out-arguments",
64 cl::desc("Replace pointer out arguments with "
65 "struct returns for non-private address space"),
67 cl::init(false));
68
70 "amdgpu-max-return-arg-num-regs",
71 cl::desc("Approximately limit number of return registers for replacing out arguments"),
73 cl::init(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 MemoryDependenceResults *MDA = nullptr;
86
87 Type *getStoredType(Value &Arg) const;
88 Type *getOutArgumentType(Argument &Arg) const;
89
90public:
91 static char ID;
92
93 AMDGPURewriteOutArguments() : FunctionPass(ID) {}
94
95 void getAnalysisUsage(AnalysisUsage &AU) const override {
98 }
99
100 bool doInitialization(Module &M) override;
101 bool runOnFunction(Function &F) override;
102};
103
104} // end anonymous namespace
105
106INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
107 "AMDGPU Rewrite Out Arguments", false, false)
109INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
110 "AMDGPU Rewrite Out Arguments", false, false)
111
112char AMDGPURewriteOutArguments::ID = 0;
113
114Type *AMDGPURewriteOutArguments::getStoredType(Value &Arg) const {
115 const int MaxUses = 10;
116 int UseCount = 0;
117
118 SmallVector<Use *> Worklist;
119 for (Use &U : Arg.uses())
120 Worklist.push_back(&U);
121
122 Type *StoredType = nullptr;
123 while (!Worklist.empty()) {
124 Use *U = Worklist.pop_back_val();
125
126 if (auto *BCI = dyn_cast<BitCastInst>(U->getUser())) {
127 for (Use &U : BCI->uses())
128 Worklist.push_back(&U);
129 continue;
130 }
131
132 if (auto *SI = dyn_cast<StoreInst>(U->getUser())) {
133 if (UseCount++ > MaxUses)
134 return nullptr;
135
136 if (!SI->isSimple() ||
137 U->getOperandNo() != StoreInst::getPointerOperandIndex())
138 return nullptr;
139
140 if (StoredType && StoredType != SI->getValueOperand()->getType())
141 return nullptr; // More than one type.
142 StoredType = SI->getValueOperand()->getType();
143 continue;
144 }
145
146 // Unsupported user.
147 return nullptr;
148 }
149
150 return StoredType;
151}
152
153Type *AMDGPURewriteOutArguments::getOutArgumentType(Argument &Arg) const {
154 const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
155 PointerType *ArgTy = dyn_cast<PointerType>(Arg.getType());
156
157 // TODO: It might be useful for any out arguments, not just privates.
158 if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
159 !AnyAddressSpace) ||
160 Arg.hasByValAttr() || Arg.hasStructRetAttr()) {
161 return nullptr;
162 }
163
164 Type *StoredType = getStoredType(Arg);
165 if (!StoredType || DL->getTypeStoreSize(StoredType) > MaxOutArgSizeBytes)
166 return nullptr;
167
168 return StoredType;
169}
170
171bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
172 DL = &M.getDataLayout();
173 return false;
174}
175
176bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
177 if (skipFunction(F))
178 return false;
179
180 // TODO: Could probably handle variadic functions.
181 if (F.isVarArg() || F.hasStructRetAttr() ||
182 AMDGPU::isEntryFunctionCC(F.getCallingConv()))
183 return false;
184
185 MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
186
187 unsigned ReturnNumRegs = 0;
188 SmallDenseMap<int, Type *, 4> OutArgIndexes;
189 SmallVector<Type *, 4> ReturnTypes;
190 Type *RetTy = F.getReturnType();
191 if (!RetTy->isVoidTy()) {
192 ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
193
194 if (ReturnNumRegs >= MaxNumRetRegs)
195 return false;
196
197 ReturnTypes.push_back(RetTy);
198 }
199
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({&Arg, Ty});
206 }
207 }
208
209 if (OutArgs.empty())
210 return false;
211
212 using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
213
215
217 for (BasicBlock &BB : F) {
218 if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
219 Returns.push_back(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. If we have 2 stores to arguments, on the first
233 // attempt the MDA query will succeed for the second store but not the
234 // first. On the second iteration we've removed that out clobbering argument
235 // (by effectively moving it into another function) and will find the second
236 // argument is OK to move.
237 for (const auto &Pair : OutArgs) {
238 bool ThisReplaceable = true;
240
241 Argument *OutArg = Pair.first;
242 Type *ArgTy = Pair.second;
243
244 // Skip this argument if converting it will push us over the register
245 // count to return limit.
246
247 // TODO: This is an approximation. When legalized this could be more. We
248 // can ask TLI for exactly how many.
249 unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
250 if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
251 continue;
252
253 // An argument is convertible only if all exit blocks are able to replace
254 // it.
255 for (ReturnInst *RI : Returns) {
256 BasicBlock *BB = RI->getParent();
257
258 MemDepResult Q = MDA->getPointerDependencyFrom(
259 MemoryLocation::getBeforeOrAfter(OutArg), true, BB->end(), BB, RI);
260 StoreInst *SI = nullptr;
261 if (Q.isDef())
262 SI = dyn_cast<StoreInst>(Q.getInst());
263
264 if (SI) {
265 LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
266 ReplaceableStores.emplace_back(RI, SI);
267 } else {
268 ThisReplaceable = false;
269 break;
270 }
271 }
272
273 if (!ThisReplaceable)
274 continue; // Try the next argument candidate.
275
276 for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
277 Value *ReplVal = Store.second->getValueOperand();
278
279 auto &ValVec = Replacements[Store.first];
280 if (llvm::any_of(ValVec,
281 [OutArg](const std::pair<Argument *, Value *> &Entry) {
282 return Entry.first == OutArg;
283 })) {
285 << "Saw multiple out arg stores" << *OutArg << '\n');
286 // It is possible to see stores to the same argument multiple times,
287 // but we expect these would have been optimized out already.
288 ThisReplaceable = false;
289 break;
290 }
291
292 ValVec.emplace_back(OutArg, ReplVal);
293 Store.second->eraseFromParent();
294 }
295
296 if (ThisReplaceable) {
297 ReturnTypes.push_back(ArgTy);
298 OutArgIndexes.insert({OutArg->getArgNo(), ArgTy});
299 ++NumOutArgumentsReplaced;
300 Changing = true;
301 }
302 }
303 } while (Changing);
304
305 if (Replacements.empty())
306 return false;
307
308 LLVMContext &Ctx = F.getParent()->getContext();
309 StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
310
311 FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
312 F.getFunctionType()->params(),
313 F.isVarArg());
314
315 LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
316
317 Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
318 F.getName() + ".body");
319 F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
320 NewFunc->copyAttributesFrom(&F);
321 NewFunc->setComdat(F.getComdat());
322
323 // We want to preserve the function and param attributes, but need to strip
324 // off any return attributes, e.g. zeroext doesn't make sense with a struct.
325 NewFunc->stealArgumentListFrom(F);
326
327 AttributeMask RetAttrs;
328 RetAttrs.addAttribute(Attribute::SExt);
329 RetAttrs.addAttribute(Attribute::ZExt);
330 RetAttrs.addAttribute(Attribute::NoAlias);
331 NewFunc->removeRetAttrs(RetAttrs);
332 // TODO: How to preserve metadata?
333
334 // Move the body of the function into the new rewritten function, and replace
335 // this function with a stub.
336 NewFunc->splice(NewFunc->begin(), &F);
337
338 for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
339 ReturnInst *RI = Replacement.first;
340 IRBuilder<> B(RI);
341 B.SetCurrentDebugLocation(RI->getDebugLoc());
342
343 int RetIdx = 0;
344 Value *NewRetVal = PoisonValue::get(NewRetTy);
345
346 Value *RetVal = RI->getReturnValue();
347 if (RetVal)
348 NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
349
350 for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second)
351 NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, RetIdx++);
352
353 if (RetVal)
354 RI->setOperand(0, NewRetVal);
355 else {
356 B.CreateRet(NewRetVal);
357 RI->eraseFromParent();
358 }
359 }
360
361 SmallVector<Value *, 16> StubCallArgs;
362 for (Argument &Arg : F.args()) {
363 if (OutArgIndexes.count(Arg.getArgNo())) {
364 // It's easier to preserve the type of the argument list. We rely on
365 // DeadArgumentElimination to take care of these.
366 StubCallArgs.push_back(PoisonValue::get(Arg.getType()));
367 } else {
368 StubCallArgs.push_back(&Arg);
369 }
370 }
371
372 BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
373 IRBuilder<> B(StubBB);
374 CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
375
376 int RetIdx = RetTy->isVoidTy() ? 0 : 1;
377 for (Argument &Arg : F.args()) {
378 if (!OutArgIndexes.count(Arg.getArgNo()))
379 continue;
380
381 Type *EltTy = OutArgIndexes[Arg.getArgNo()];
382 const auto Align =
383 DL->getValueOrABITypeAlignment(Arg.getParamAlign(), EltTy);
384
385 Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
386 B.CreateAlignedStore(Val, &Arg, Align);
387 }
388
389 if (!RetTy->isVoidTy()) {
390 B.CreateRet(B.CreateExtractValue(StubCall, 0));
391 } else {
392 B.CreateRetVoid();
393 }
394
395 // The function is now a stub we want to inline.
396 F.addFnAttr(Attribute::AlwaysInline);
397
398 ++NumOutArgumentFunctionsReplaced;
399 return true;
400}
401
403 return new AMDGPURewriteOutArguments();
404}
MachineBasicBlock MachineBasicBlock::iterator DebugLoc DL
AMDGPU Rewrite Out Arguments
static cl::opt< unsigned > MaxNumRetRegs("amdgpu-max-return-arg-num-regs", cl::desc("Approximately limit number of return registers for replacing out arguments"), cl::Hidden, cl::init(16))
#define DEBUG_TYPE
static cl::opt< bool > AnyAddressSpace("amdgpu-any-address-space-out-arguments", cl::desc("Replace pointer out arguments with " "struct returns for non-private address space"), cl::Hidden, cl::init(false))
static GCRegistry::Add< OcamlGC > B("ocaml", "ocaml 3.10-compatible GC")
return RetTy
#define LLVM_DEBUG(X)
Definition: Debug.h:101
#define F(x, y, z)
Definition: MD5.cpp:55
#define INITIALIZE_PASS_DEPENDENCY(depName)
Definition: PassSupport.h:55
#define INITIALIZE_PASS_END(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:59
#define INITIALIZE_PASS_BEGIN(passName, arg, name, cfg, analysis)
Definition: PassSupport.h:52
This file defines the SmallSet class.
This file defines the 'Statistic' class, which is designed to be an easy way to expose various metric...
#define STATISTIC(VARNAME, DESC)
Definition: Statistic.h:167
Represent the analysis usage information of a pass.
AnalysisUsage & addRequired()
This class represents an incoming formal argument to a Function.
Definition: Argument.h:28
bool hasByValAttr() const
Return true if this argument has the byval attribute.
Definition: Function.cpp:110
unsigned getArgNo() const
Return the index of this formal argument in its containing function.
Definition: Argument.h:46
MaybeAlign getParamAlign() const
If this is a byval or inalloca argument, return its alignment.
Definition: Function.cpp:192
bool hasStructRetAttr() const
Return true if this argument has the sret attribute.
Definition: Function.cpp:257
AttributeMask & addAttribute(Attribute::AttrKind Val)
Add an attribute to the mask.
Definition: AttributeMask.h:44
LLVM Basic Block Representation.
Definition: BasicBlock.h:56
iterator end()
Definition: BasicBlock.h:337
static BasicBlock * Create(LLVMContext &Context, const Twine &Name="", Function *Parent=nullptr, BasicBlock *InsertBefore=nullptr)
Creates a new BasicBlock.
Definition: BasicBlock.h:105
const Function * getParent() const
Return the enclosing method, or null if none.
Definition: BasicBlock.h:112
This class represents a function call, abstracting a target machine's calling convention.
A parsed version of the target data layout string in and methods for querying it.
Definition: DataLayout.h:110
bool empty() const
Definition: DenseMap.h:98
size_type count(const_arg_type_t< KeyT > Val) const
Return 1 if the specified key is in the map, 0 otherwise.
Definition: DenseMap.h:151
std::pair< iterator, bool > insert(const std::pair< KeyT, ValueT > &KV)
Definition: DenseMap.h:220
FunctionPass class - This class is used to implement most global optimizations.
Definition: Pass.h:311
virtual bool runOnFunction(Function &F)=0
runOnFunction - Virtual method overriden by subclasses to do the per-function processing of the pass.
static Function * Create(FunctionType *Ty, LinkageTypes Linkage, unsigned AddrSpace, const Twine &N="", Module *M=nullptr)
Definition: Function.h:138
void splice(Function::iterator ToIt, Function *FromF)
Transfer all blocks from FromF to this function at ToIt.
Definition: Function.h:699
iterator begin()
Definition: Function.h:763
void stealArgumentListFrom(Function &Src)
Steal arguments from another function.
Definition: Function.cpp:484
void removeRetAttrs(const AttributeMask &Attrs)
removes the attributes from the return value list of attributes.
Definition: Function.cpp:623
void copyAttributesFrom(const Function *Src)
copyAttributesFrom - copy all additional attributes (those not needed to create a Function) from the ...
Definition: Function.cpp:757
void setComdat(Comdat *C)
Definition: Globals.cpp:196
This provides a uniform API for creating instructions and inserting them into a basic block: either a...
Definition: IRBuilder.h:2625
const DebugLoc & getDebugLoc() const
Return the debug location for this node as a DebugLoc.
Definition: Instruction.h:392
SymbolTableList< Instruction >::iterator eraseFromParent()
This method unlinks 'this' from the containing basic block and deletes it.
Definition: Instruction.cpp:83
This is an important class for using LLVM in a threaded context.
Definition: LLVMContext.h:67
A memory dependence query can return one of three different answers.
bool isDef() const
Tests if this MemDepResult represents a query that is an instruction definition dependency.
Instruction * getInst() const
If this is a normal dependency, returns the instruction that is depended on.
Provides a lazy, caching interface for making common memory aliasing information queries,...
A wrapper analysis pass for the legacy pass manager that exposes a MemoryDepnedenceResults instance.
static MemoryLocation getBeforeOrAfter(const Value *Ptr, const AAMDNodes &AATags=AAMDNodes())
Return a location that may access any location before or after Ptr, while remaining within the underl...
A Module instance is used to store all the information related to an LLVM module.
Definition: Module.h:65
virtual void getAnalysisUsage(AnalysisUsage &) const
getAnalysisUsage - This function should be overriden by passes that need analysis information to do t...
Definition: Pass.cpp:98
virtual bool doInitialization(Module &)
doInitialization - Virtual method overridden by subclasses to do any necessary initialization before ...
Definition: Pass.h:119
static PoisonValue * get(Type *T)
Static factory methods - Return an 'poison' object of the specified type.
Definition: Constants.cpp:1743
Return a value (possibly void), from a function.
Value * getReturnValue() const
Convenience accessor. Returns null if there is no return value.
bool empty() const
Definition: SmallVector.h:94
reference emplace_back(ArgTypes &&... Args)
Definition: SmallVector.h:941
void push_back(const T &Elt)
Definition: SmallVector.h:416
This is a 'vector' (really, a variable-sized array), optimized for the case when the array is small.
Definition: SmallVector.h:1200
An instruction for storing to memory.
Definition: Instructions.h:301
static unsigned getPointerOperandIndex()
Definition: Instructions.h:395
Class to represent struct types.
Definition: DerivedTypes.h:213
static StructType * create(LLVMContext &Context, StringRef Name)
This creates an identified struct.
Definition: Type.cpp:514
The instances of the Type class are immutable: once they are created, they are never changed.
Definition: Type.h:45
A Use represents the edge between a Value definition and its users.
Definition: Use.h:43
void setOperand(unsigned i, Value *Val)
Definition: User.h:174
LLVM Value Representation.
Definition: Value.h:74
Type * getType() const
All values are typed, get the type of this value.
Definition: Value.h:255
bool isEntryFunctionCC(CallingConv::ID CC)
unsigned ID
LLVM IR allows to use arbitrary numbers as calling convention identifiers.
Definition: CallingConv.h:24
initializer< Ty > init(const Ty &Val)
Definition: CommandLine.h:445
This is an optimization pass for GlobalISel generic memory operations.
Definition: AddressRanges.h:18
FunctionPass * createAMDGPURewriteOutArgumentsPass()
bool any_of(R &&range, UnaryPredicate P)
Provide wrappers to std::any_of which take ranges instead of having to pass begin/end explicitly.
Definition: STLExtras.h:1734
raw_ostream & dbgs()
dbgs() - This returns a reference to a raw_ostream for debugging messages.
Definition: Debug.cpp:163
This struct is a compact representation of a valid (non-zero power of two) alignment.
Definition: Alignment.h:39