400 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9/// \file This pass attempts to replace out argument usage with a return of a10/// struct.11///12/// We can support returning a lot of values directly in registers, but13/// idiomatic C code frequently uses a pointer argument to return a second value14/// rather than returning a struct by value. GPU stack access is also quite15/// painful, so we want to avoid that if possible. Passing a stack object16/// pointer to a function also requires an additional address expansion code17/// sequence to convert the pointer to be relative to the kernel's scratch wave18/// offset register since the callee doesn't know what stack frame the incoming19/// 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 variable35/// to use the API, which if replaced with a struct return will be easily SROA'd36/// out when the stub function we create is inlined37///38/// This pass introduces the struct return, but leaves the unused pointer39/// arguments and introduces a new stub function calling the struct returning40/// 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(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(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 namespace104 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(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;152 PointerType *ArgTy = dyn_cast<PointerType>(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(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 197 SmallVector<std::pair<Argument *, Type *>, 4> OutArgs;198 for (Argument &Arg : F.args()) {199 if (Type *Ty = getOutArgumentType(Arg)) {200 LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg201 << " 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 213 SmallVector<ReturnInst *, 4> Returns;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. This228 // helps with cases with multiple arguments which may alias, such as in a229 // sincos implementation. If we have 2 stores to arguments, on the first230 // attempt the MDA query will succeed for the second store but not the231 // first. On the second iteration we've removed that out clobbering argument232 // (by effectively moving it into another function) and will find the second233 // argument is OK to move.234 for (const auto &Pair : OutArgs) {235 bool ThisReplaceable = true;236 SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;237 238 Argument *OutArg = Pair.first;239 Type *ArgTy = Pair.second;240 241 // Skip this argument if converting it will push us over the register242 // count to return limit.243 244 // TODO: This is an approximation. When legalized this could be more. We245 // 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 replace251 // 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())259 SI = dyn_cast<StoreInst>(Q.getInst());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)) {278 LLVM_DEBUG(dbgs()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.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 strip318 // 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 replace329 // 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 on359 // 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 397FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {398 return new AMDGPURewriteOutArguments();399}400