LLVM 22.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/Statistic.h"
49#include "llvm/IR/IRBuilder.h"
52#include "llvm/Pass.h"
54#include "llvm/Support/Debug.h"
56
57#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
58
59using namespace llvm;
60
62 "amdgpu-any-address-space-out-arguments",
63 cl::desc("Replace pointer out arguments with "
64 "struct returns for non-private address space"),
66 cl::init(false));
67
69 "amdgpu-max-return-arg-num-regs",
70 cl::desc("Approximately limit number of return registers for replacing out arguments"),
72 cl::init(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)
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(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>(U->getUser())) {
124 for (Use &U : BCI->uses())
125 Worklist.push_back(&U);
126 continue;
127 }
128
129 if (auto *SI = dyn_cast<StoreInst>(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;
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(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(F.getCallingConv()))
180 return false;
181
182 MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
183
184 unsigned ReturnNumRegs = 0;
185 SmallDenseMap<int, Type *, 4> OutArgIndexes;
186 SmallVector<Type *, 4> ReturnTypes;
187 Type *RetTy = F.getReturnType();
188 if (!RetTy->isVoidTy()) {
189 ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
190
191 if (ReturnNumRegs >= MaxNumRetRegs)
192 return false;
193
194 ReturnTypes.push_back(RetTy);
195 }
196
198 for (Argument &Arg : F.args()) {
199 if (Type *Ty = getOutArgumentType(Arg)) {
200 LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg
201 << " in function " << F.getName() << '\n');
202 OutArgs.push_back({&Arg, Ty});
203 }
204 }
205
206 if (OutArgs.empty())
207 return false;
208
209 using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
210
211 DenseMap<ReturnInst *, ReplacementVec> Replacements;
212
214 for (BasicBlock &BB : F) {
215 if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
216 Returns.push_back(RI);
217 }
218
219 if (Returns.empty())
220 return false;
221
222 bool Changing;
223
224 do {
225 Changing = false;
226
227 // Keep retrying if we are able to successfully eliminate an argument. This
228 // helps with cases with multiple arguments which may alias, such as in a
229 // sincos implementation. If we have 2 stores to arguments, on the first
230 // attempt the MDA query will succeed for the second store but not the
231 // first. On the second iteration we've removed that out clobbering argument
232 // (by effectively moving it into another function) and will find the second
233 // argument is OK to move.
234 for (const auto &Pair : OutArgs) {
235 bool ThisReplaceable = true;
237
238 Argument *OutArg = Pair.first;
239 Type *ArgTy = Pair.second;
240
241 // Skip this argument if converting it will push us over the register
242 // count to return limit.
243
244 // TODO: This is an approximation. When legalized this could be more. We
245 // can ask TLI for exactly how many.
246 unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
247 if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
248 continue;
249
250 // An argument is convertible only if all exit blocks are able to replace
251 // it.
252 for (ReturnInst *RI : Returns) {
253 BasicBlock *BB = RI->getParent();
254
255 MemDepResult Q = MDA->getPointerDependencyFrom(
256 MemoryLocation::getBeforeOrAfter(OutArg), true, BB->end(), BB, RI);
257 StoreInst *SI = nullptr;
258 if (Q.isDef())
260
261 if (SI) {
262 LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
263 ReplaceableStores.emplace_back(RI, SI);
264 } else {
265 ThisReplaceable = false;
266 break;
267 }
268 }
269
270 if (!ThisReplaceable)
271 continue; // Try the next argument candidate.
272
273 for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
274 Value *ReplVal = Store.second->getValueOperand();
275
276 auto &ValVec = Replacements[Store.first];
277 if (llvm::is_contained(llvm::make_first_range(ValVec), OutArg)) {
279 << "Saw multiple out arg stores" << *OutArg << '\n');
280 // It is possible to see stores to the same argument multiple times,
281 // but we expect these would have been optimized out already.
282 ThisReplaceable = false;
283 break;
284 }
285
286 ValVec.emplace_back(OutArg, ReplVal);
287 Store.second->eraseFromParent();
288 }
289
290 if (ThisReplaceable) {
291 ReturnTypes.push_back(ArgTy);
292 OutArgIndexes.insert({OutArg->getArgNo(), ArgTy});
293 ++NumOutArgumentsReplaced;
294 Changing = true;
295 }
296 }
297 } while (Changing);
298
299 if (Replacements.empty())
300 return false;
301
302 LLVMContext &Ctx = F.getParent()->getContext();
303 StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
304
305 FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
306 F.getFunctionType()->params(),
307 F.isVarArg());
308
309 LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
310
311 Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
312 F.getName() + ".body");
313 F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
314 NewFunc->copyAttributesFrom(&F);
315 NewFunc->setComdat(F.getComdat());
316
317 // We want to preserve the function and param attributes, but need to strip
318 // off any return attributes, e.g. zeroext doesn't make sense with a struct.
319 NewFunc->stealArgumentListFrom(F);
320
321 AttributeMask RetAttrs;
322 RetAttrs.addAttribute(Attribute::SExt);
323 RetAttrs.addAttribute(Attribute::ZExt);
324 RetAttrs.addAttribute(Attribute::NoAlias);
325 NewFunc->removeRetAttrs(RetAttrs);
326 // TODO: How to preserve metadata?
327
328 // Move the body of the function into the new rewritten function, and replace
329 // this function with a stub.
330 NewFunc->splice(NewFunc->begin(), &F);
331
332 for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
333 ReturnInst *RI = Replacement.first;
334 IRBuilder<> B(RI);
335 B.SetCurrentDebugLocation(RI->getDebugLoc());
336
337 int RetIdx = 0;
338 Value *NewRetVal = PoisonValue::get(NewRetTy);
339
340 Value *RetVal = RI->getReturnValue();
341 if (RetVal)
342 NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
343
344 for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second)
345 NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, RetIdx++);
346
347 if (RetVal)
348 RI->setOperand(0, NewRetVal);
349 else {
350 B.CreateRet(NewRetVal);
351 RI->eraseFromParent();
352 }
353 }
354
355 SmallVector<Value *, 16> StubCallArgs;
356 for (Argument &Arg : F.args()) {
357 if (OutArgIndexes.count(Arg.getArgNo())) {
358 // It's easier to preserve the type of the argument list. We rely on
359 // DeadArgumentElimination to take care of these.
360 StubCallArgs.push_back(PoisonValue::get(Arg.getType()));
361 } else {
362 StubCallArgs.push_back(&Arg);
363 }
364 }
365
366 BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
367 IRBuilder<> B(StubBB);
368 CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
369
370 int RetIdx = RetTy->isVoidTy() ? 0 : 1;
371 for (Argument &Arg : F.args()) {
372 auto It = OutArgIndexes.find(Arg.getArgNo());
373 if (It == OutArgIndexes.end())
374 continue;
375
376 Type *EltTy = It->second;
377 const auto Align =
378 DL->getValueOrABITypeAlignment(Arg.getParamAlign(), EltTy);
379
380 Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
381 B.CreateAlignedStore(Val, &Arg, Align);
382 }
383
384 if (!RetTy->isVoidTy()) {
385 B.CreateRet(B.CreateExtractValue(StubCall, 0));
386 } else {
387 B.CreateRetVoid();
388 }
389
390 // The function is now a stub we want to inline.
391 F.addFnAttr(Attribute::AlwaysInline);
392
393 ++NumOutArgumentFunctionsReplaced;
394 return true;
395}
396
398 return new AMDGPURewriteOutArguments();
399}
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))
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))
MachineBasicBlock MachineBasicBlock::iterator DebugLoc DL
static GCRegistry::Add< OcamlGC > B("ocaml", "ocaml 3.10-compatible GC")
static bool runOnFunction(Function &F, bool PostInlining)
#define DEBUG_TYPE
#define F(x, y, z)
Definition MD5.cpp:55
Machine Check Debug Module
#define INITIALIZE_PASS_DEPENDENCY(depName)
Definition PassSupport.h:42
#define INITIALIZE_PASS_END(passName, arg, name, cfg, analysis)
Definition PassSupport.h:44
#define INITIALIZE_PASS_BEGIN(passName, arg, name, cfg, analysis)
Definition PassSupport.h:39
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:171
#define LLVM_DEBUG(...)
Definition Debug.h:114
AnalysisUsage & addRequired()
This class represents an incoming formal argument to a Function.
Definition Argument.h:32
LLVM_ABI bool hasByValAttr() const
Return true if this argument has the byval attribute.
Definition Function.cpp:128
unsigned getArgNo() const
Return the index of this formal argument in its containing function.
Definition Argument.h:50
LLVM_ABI MaybeAlign getParamAlign() const
If this is a byval or inalloca argument, return its alignment.
Definition Function.cpp:216
LLVM_ABI bool hasStructRetAttr() const
Return true if this argument has the sret attribute.
Definition Function.cpp:288
AttributeMask & addAttribute(Attribute::AttrKind Val)
Add an attribute to the mask.
iterator end()
Definition BasicBlock.h:472
static BasicBlock * Create(LLVMContext &Context, const Twine &Name="", Function *Parent=nullptr, BasicBlock *InsertBefore=nullptr)
Creates a new BasicBlock.
Definition BasicBlock.h:206
A parsed version of the target data layout string in and methods for querying it.
Definition DataLayout.h:63
iterator find(const_arg_type_t< KeyT > Val)
Definition DenseMap.h:167
bool empty() const
Definition DenseMap.h:109
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:163
iterator end()
Definition DenseMap.h:81
std::pair< iterator, bool > insert(const std::pair< KeyT, ValueT > &KV)
Definition DenseMap.h:222
FunctionPass class - This class is used to implement most global optimizations.
Definition Pass.h:314
static Function * Create(FunctionType *Ty, LinkageTypes Linkage, unsigned AddrSpace, const Twine &N="", Module *M=nullptr)
Definition Function.h:166
void splice(Function::iterator ToIt, Function *FromF)
Transfer all blocks from FromF to this function at ToIt.
Definition Function.h:759
iterator begin()
Definition Function.h:851
void stealArgumentListFrom(Function &Src)
Steal arguments from another function.
Definition Function.cpp:566
void removeRetAttrs(const AttributeMask &Attrs)
removes the attributes from the return value list of attributes.
Definition Function.cpp:705
void copyAttributesFrom(const Function *Src)
copyAttributesFrom - copy all additional attributes (those not needed to create a Function) from the ...
Definition Function.cpp:856
LLVM_ABI void setComdat(Comdat *C)
Definition Globals.cpp:214
const DebugLoc & getDebugLoc() const
Return the debug location for this node as a DebugLoc.
LLVM_ABI InstListType::iterator eraseFromParent()
This method unlinks 'this' from the containing basic block and deletes it.
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,...
MemDepResult getPointerDependencyFrom(const MemoryLocation &Loc, bool isLoad, BasicBlock::iterator ScanIt, BasicBlock *BB, Instruction *QueryInst=nullptr, unsigned *Limit=nullptr)
Returns the instruction on which a memory location depends.
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...
static LLVM_ABI PoisonValue * get(Type *T)
Static factory methods - Return an 'poison' object of the specified type.
Value * getReturnValue() const
Convenience accessor. Returns null if there is no return value.
reference emplace_back(ArgTypes &&... Args)
void push_back(const T &Elt)
This is a 'vector' (really, a variable-sized array), optimized for the case when the array is small.
static unsigned getPointerOperandIndex()
static LLVM_ABI StructType * create(LLVMContext &Context, StringRef Name)
This creates an identified struct.
Definition Type.cpp:620
The instances of the Type class are immutable: once they are created, they are never changed.
Definition Type.h:45
bool isVoidTy() const
Return true if this is 'void'.
Definition Type.h:139
A Use represents the edge between a Value definition and its users.
Definition Use.h:35
void setOperand(unsigned i, Value *Val)
Definition User.h:237
LLVM Value Representation.
Definition Value.h:75
Type * getType() const
All values are typed, get the type of this value.
Definition Value.h:256
iterator_range< use_iterator > uses()
Definition Value.h:380
constexpr char Align[]
Key for Kernel::Arg::Metadata::mAlign.
LLVM_READNONE constexpr bool isEntryFunctionCC(CallingConv::ID CC)
unsigned ID
LLVM IR allows to use arbitrary numbers as calling convention identifiers.
Definition CallingConv.h:24
@ BasicBlock
Various leaf nodes.
Definition ISDOpcodes.h:81
initializer< Ty > init(const Ty &Val)
This is an optimization pass for GlobalISel generic memory operations.
FunctionAddr VTableAddr Value
Definition InstrProf.h:137
decltype(auto) dyn_cast(const From &Val)
dyn_cast<X> - Return the argument parameter cast to the specified type.
Definition Casting.h:649
FunctionPass * createAMDGPURewriteOutArgumentsPass()
LLVM_ABI raw_ostream & dbgs()
dbgs() - This returns a reference to a raw_ostream for debugging messages.
Definition Debug.cpp:207
auto make_first_range(ContainerTy &&c)
Given a container of pairs, return a range over the first elements.
Definition STLExtras.h:1399
class LLVM_GSL_OWNER SmallVector
Forward declaration of SmallVector so that calculateSmallVectorDefaultInlinedElements can reference s...
IRBuilder(LLVMContext &, FolderTy, InserterTy, MDNode *, ArrayRef< OperandBundleDef >) -> IRBuilder< FolderTy, InserterTy >
iterator_range< pointer_iterator< WrappedIteratorT > > make_pointer_range(RangeT &&Range)
Definition iterator.h:363
bool is_contained(R &&Range, const E &Element)
Returns true if Element is found in Range.
Definition STLExtras.h:1877