LCOV - code coverage report
Current view: top level - lib/Target/AMDGPU - AMDGPURewriteOutArguments.cpp (source / functions) Hit Total Coverage
Test: llvm-toolchain.info Lines: 174 176 98.9 %
Date: 2017-09-14 15:23:50 Functions: 10 12 83.3 %
Legend: Lines: hit not hit

          Line data    Source code
       1             : //===- AMDGPURewriteOutArgumentsPass.cpp - Create struct returns ----------===//
       2             : //
       3             : //                     The LLVM Compiler Infrastructure
       4             : //
       5             : // This file is distributed under the University of Illinois Open Source
       6             : // License. See LICENSE.TXT for details.
       7             : //
       8             : //===----------------------------------------------------------------------===//
       9             : //
      10             : /// \file This pass attempts to replace out argument usage with a return of a
      11             : /// struct.
      12             : ///
      13             : /// We can support returning a lot of values directly in registers, but
      14             : /// idiomatic C code frequently uses a pointer argument to return a second value
      15             : /// rather than returning a struct by value. GPU stack access is also quite
      16             : /// painful, so we want to avoid that if possible. Passing a stack object
      17             : /// pointer to a function also requires an additional address expansion code
      18             : /// sequence to convert the pointer to be relative to the kernel's scratch wave
      19             : /// offset register since the callee doesn't know what stack frame the incoming
      20             : /// pointer is relative to.
      21             : ///
      22             : /// The goal is to try rewriting code that looks like this:
      23             : ///
      24             : ///  int foo(int a, int b, int* out) {
      25             : ///     *out = bar();
      26             : ///     return a + b;
      27             : /// }
      28             : ///
      29             : /// into something like this:
      30             : ///
      31             : ///  std::pair<int, int> foo(int a, int b) {
      32             : ///     return std::make_pair(a + b, bar());
      33             : /// }
      34             : ///
      35             : /// Typically the incoming pointer is a simple alloca for a temporary variable
      36             : /// to use the API, which if replaced with a struct return will be easily SROA'd
      37             : /// out when the stub function we create is inlined
      38             : ///
      39             : /// This pass introduces the struct return, but leaves the unused pointer
      40             : /// arguments and introduces a new stub function calling the struct returning
      41             : /// body. DeadArgumentElimination should be run after this to clean these up.
      42             : //
      43             : //===----------------------------------------------------------------------===//
      44             : 
      45             : #include "AMDGPU.h"
      46             : #include "Utils/AMDGPUBaseInfo.h"
      47             : #include "llvm/Analysis/MemoryDependenceAnalysis.h"
      48             : #include "llvm/ADT/DenseMap.h"
      49             : #include "llvm/ADT/STLExtras.h"
      50             : #include "llvm/ADT/SmallSet.h"
      51             : #include "llvm/ADT/SmallVector.h"
      52             : #include "llvm/ADT/Statistic.h"
      53             : #include "llvm/Analysis/MemoryLocation.h"
      54             : #include "llvm/IR/Argument.h"
      55             : #include "llvm/IR/Attributes.h"
      56             : #include "llvm/IR/BasicBlock.h"
      57             : #include "llvm/IR/Constants.h"
      58             : #include "llvm/IR/DataLayout.h"
      59             : #include "llvm/IR/DerivedTypes.h"
      60             : #include "llvm/IR/Function.h"
      61             : #include "llvm/IR/IRBuilder.h"
      62             : #include "llvm/IR/Instructions.h"
      63             : #include "llvm/IR/Module.h"
      64             : #include "llvm/IR/Type.h"
      65             : #include "llvm/IR/Use.h"
      66             : #include "llvm/IR/User.h"
      67             : #include "llvm/IR/Value.h"
      68             : #include "llvm/Pass.h"
      69             : #include "llvm/Support/Casting.h"
      70             : #include "llvm/Support/CommandLine.h"
      71             : #include "llvm/Support/Debug.h"
      72             : #include "llvm/Support/raw_ostream.h"
      73             : #include <cassert>
      74             : #include <utility>
      75             : 
      76             : #define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
      77             : 
      78             : using namespace llvm;
      79             : 
      80       72306 : static cl::opt<bool> AnyAddressSpace(
      81             :   "amdgpu-any-address-space-out-arguments",
      82      216918 :   cl::desc("Replace pointer out arguments with "
      83             :            "struct returns for non-private address space"),
      84             :   cl::Hidden,
      85      289224 :   cl::init(false));
      86             : 
      87       72306 : static cl::opt<unsigned> MaxNumRetRegs(
      88             :   "amdgpu-max-return-arg-num-regs",
      89      216918 :   cl::desc("Approximately limit number of return registers for replacing out arguments"),
      90             :   cl::Hidden,
      91      289224 :   cl::init(16));
      92             : 
      93             : STATISTIC(NumOutArgumentsReplaced,
      94             :           "Number out arguments moved to struct return values");
      95             : STATISTIC(NumOutArgumentFunctionsReplaced,
      96             :           "Number of functions with out arguments moved to struct return values");
      97             : 
      98             : namespace {
      99             : 
     100           4 : class AMDGPURewriteOutArguments : public FunctionPass {
     101             : private:
     102             :   const DataLayout *DL = nullptr;
     103             :   MemoryDependenceResults *MDA = nullptr;
     104             : 
     105             :   bool checkArgumentUses(Value &Arg) const;
     106             :   bool isOutArgumentCandidate(Argument &Arg) const;
     107             : 
     108             : #ifndef NDEBUG
     109             :   bool isVec3ToVec4Shuffle(Type *Ty0, Type* Ty1) const;
     110             : #endif
     111             : 
     112             : public:
     113             :   static char ID;
     114             : 
     115           4 :   AMDGPURewriteOutArguments() : FunctionPass(ID) {}
     116             : 
     117           2 :   void getAnalysisUsage(AnalysisUsage &AU) const override {
     118           2 :     AU.addRequired<MemoryDependenceWrapperPass>();
     119           2 :     FunctionPass::getAnalysisUsage(AU);
     120           2 :   }
     121             : 
     122             :   bool doInitialization(Module &M) override;
     123             :   bool runOnFunction(Function &F) override;
     124             : };
     125             : 
     126             : } // end anonymous namespace
     127             : 
     128       53042 : INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
     129             :                       "AMDGPU Rewrite Out Arguments", false, false)
     130       53042 : INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
     131      312538 : INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
     132             :                     "AMDGPU Rewrite Out Arguments", false, false)
     133             : 
     134             : char AMDGPURewriteOutArguments::ID = 0;
     135             : 
     136          56 : bool AMDGPURewriteOutArguments::checkArgumentUses(Value &Arg) const {
     137          68 :   const int MaxUses = 10;
     138          68 :   int UseCount = 0;
     139             : 
     140         183 :   for (Use &U : Arg.uses()) {
     141         138 :     StoreInst *SI = dyn_cast<StoreInst>(U.getUser());
     142          69 :     if (UseCount > MaxUses)
     143             :       return false;
     144             : 
     145          69 :     if (!SI) {
     146          40 :       auto *BCI = dyn_cast<BitCastInst>(U.getUser());
     147          38 :       if (!BCI || !BCI->hasOneUse())
     148             :         return false;
     149             : 
     150             :       // We don't handle multiple stores currently, so stores to aggregate
     151             :       // pointers aren't worth the trouble since they are canonically split up.
     152          38 :       Type *DestEltTy = BCI->getType()->getPointerElementType();
     153          17 :       if (DestEltTy->isAggregateType())
     154             :         return false;
     155             : 
     156             :       // We could handle these if we had a convenient way to bitcast between
     157             :       // them.
     158          34 :       Type *SrcEltTy = Arg.getType()->getPointerElementType();
     159          17 :       if (SrcEltTy->isArrayTy())
     160             :         return false;
     161             : 
     162             :       // Special case handle structs with single members. It is useful to handle
     163             :       // some casts between structs and non-structs, but we can't bitcast
     164             :       // directly between them.  directly bitcast between them.  Blender uses
     165             :       // some casts that look like { <3 x float> }* to <4 x float>*
     166          17 :       if ((SrcEltTy->isStructTy() && (SrcEltTy->getNumContainedTypes() != 1)))
     167             :         return false;
     168             : 
     169             :       // Clang emits OpenCL 3-vector type accesses with a bitcast to the
     170             :       // equivalent 4-element vector and accesses that, and we're looking for
     171             :       // this pointer cast.
     172          15 :       if (DL->getTypeAllocSize(SrcEltTy) != DL->getTypeAllocSize(DestEltTy))
     173             :         return false;
     174             : 
     175          12 :       return checkArgumentUses(*BCI);
     176             :     }
     177             : 
     178          96 :     if (!SI->isSimple() ||
     179          48 :         U.getOperandNo() != StoreInst::getPointerOperandIndex())
     180             :       return false;
     181             : 
     182          47 :     ++UseCount;
     183             :   }
     184             : 
     185             :   // Skip unused arguments.
     186          46 :   return UseCount > 0;
     187             : }
     188             : 
     189          81 : bool AMDGPURewriteOutArguments::isOutArgumentCandidate(Argument &Arg) const {
     190          81 :   const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
     191         142 :   PointerType *ArgTy = dyn_cast<PointerType>(Arg.getType());
     192             : 
     193             :   // TODO: It might be useful for any out arguments, not just privates.
     194          64 :   if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
     195          63 :                  !AnyAddressSpace) ||
     196         174 :       Arg.hasByValAttr() || Arg.hasStructRetAttr() ||
     197         171 :       DL->getTypeStoreSize(ArgTy->getPointerElementType()) > MaxOutArgSizeBytes) {
     198             :     return false;
     199             :   }
     200             : 
     201          56 :   return checkArgumentUses(Arg);
     202             : }
     203             : 
     204           2 : bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
     205           2 :   DL = &M.getDataLayout();
     206           2 :   return false;
     207             : }
     208             : 
     209             : #ifndef NDEBUG
     210             : bool AMDGPURewriteOutArguments::isVec3ToVec4Shuffle(Type *Ty0, Type* Ty1) const {
     211             :   VectorType *VT0 = dyn_cast<VectorType>(Ty0);
     212             :   VectorType *VT1 = dyn_cast<VectorType>(Ty1);
     213             :   if (!VT0 || !VT1)
     214             :     return false;
     215             : 
     216             :   if (VT0->getNumElements() != 3 ||
     217             :       VT1->getNumElements() != 4)
     218             :     return false;
     219             : 
     220             :   return DL->getTypeSizeInBits(VT0->getElementType()) ==
     221             :          DL->getTypeSizeInBits(VT1->getElementType());
     222             : }
     223             : #endif
     224             : 
     225          60 : bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
     226          60 :   if (skipFunction(F))
     227             :     return false;
     228             : 
     229             :   // TODO: Could probably handle variadic functions.
     230         117 :   if (F.isVarArg() || F.hasStructRetAttr() ||
     231          58 :       AMDGPU::isEntryFunctionCC(F.getCallingConv()))
     232             :     return false;
     233             : 
     234         116 :   MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
     235             : 
     236          58 :   unsigned ReturnNumRegs = 0;
     237         116 :   SmallSet<int, 4> OutArgIndexes;
     238         116 :   SmallVector<Type *, 4> ReturnTypes;
     239          58 :   Type *RetTy = F.getReturnType();
     240         116 :   if (!RetTy->isVoidTy()) {
     241          18 :     ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
     242             : 
     243           9 :     if (ReturnNumRegs >= MaxNumRetRegs)
     244             :       return false;
     245             : 
     246           8 :     ReturnTypes.push_back(RetTy);
     247             :   }
     248             : 
     249          57 :   SmallVector<Argument *, 4> OutArgs;
     250         138 :   for (Argument &Arg : F.args()) {
     251          81 :     if (isOutArgumentCandidate(Arg)) {
     252             :       DEBUG(dbgs() << "Found possible out argument " << Arg
     253             :             << " in function " << F.getName() << '\n');
     254          44 :       OutArgs.push_back(&Arg);
     255             :     }
     256             :   }
     257             : 
     258          57 :   if (OutArgs.empty())
     259             :     return false;
     260             : 
     261             :   using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
     262             : 
     263          40 :   DenseMap<ReturnInst *, ReplacementVec> Replacements;
     264             : 
     265          80 :   SmallVector<ReturnInst *, 4> Returns;
     266         169 :   for (BasicBlock &BB : F) {
     267          49 :     if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
     268          43 :       Returns.push_back(RI);
     269             :   }
     270             : 
     271          40 :   if (Returns.empty())
     272             :     return false;
     273             : 
     274             :   bool Changing;
     275             : 
     276          78 :   do {
     277          78 :     Changing = false;
     278             : 
     279             :     // Keep retrying if we are able to successfully eliminate an argument. This
     280             :     // helps with cases with multiple arguments which may alias, such as in a
     281             :     // sincos implemntation. If we have 2 stores to arguments, on the first
     282             :     // attempt the MDA query will succeed for the second store but not the
     283             :     // first. On the second iteration we've removed that out clobbering argument
     284             :     // (by effectively moving it into another function) and will find the second
     285             :     // argument is OK to move.
     286         322 :     for (Argument *OutArg : OutArgs) {
     287          88 :       bool ThisReplaceable = true;
     288         129 :       SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
     289             : 
     290         176 :       Type *ArgTy = OutArg->getType()->getPointerElementType();
     291             : 
     292             :       // Skip this argument if converting it will push us over the register
     293             :       // count to return limit.
     294             : 
     295             :       // TODO: This is an approximation. When legalized this could be more. We
     296             :       // can ask TLI for exactly how many.
     297         176 :       unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
     298         176 :       if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
     299          47 :         continue;
     300             : 
     301             :       // An argument is convertible only if all exit blocks are able to replace
     302             :       // it.
     303         307 :       for (ReturnInst *RI : Returns) {
     304          90 :         BasicBlock *BB = RI->getParent();
     305             : 
     306         360 :         MemDepResult Q = MDA->getPointerDependencyFrom(MemoryLocation(OutArg),
     307         270 :                                                        true, BB->end(), BB, RI);
     308          90 :         StoreInst *SI = nullptr;
     309          90 :         if (Q.isDef())
     310          86 :           SI = dyn_cast<StoreInst>(Q.getInst());
     311             : 
     312          90 :         if (SI) {
     313             :           DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
     314          43 :           ReplaceableStores.emplace_back(RI, SI);
     315             :         } else {
     316          47 :           ThisReplaceable = false;
     317          47 :           break;
     318             :         }
     319             :       }
     320             : 
     321          47 :       if (!ThisReplaceable)
     322          47 :         continue; // Try the next argument candidate.
     323             : 
     324         165 :       for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
     325          86 :         Value *ReplVal = Store.second->getValueOperand();
     326             : 
     327          43 :         auto &ValVec = Replacements[Store.first];
     328         129 :         if (llvm::find_if(ValVec,
     329             :               [OutArg](const std::pair<Argument *, Value *> &Entry) {
     330          43 :                  return Entry.first == OutArg;}) != ValVec.end()) {
     331             :           DEBUG(dbgs() << "Saw multiple out arg stores" << *OutArg << '\n');
     332             :           // It is possible to see stores to the same argument multiple times,
     333             :           // but we expect these would have been optimized out already.
     334           1 :           ThisReplaceable = false;
     335           1 :           break;
     336             :         }
     337             : 
     338          42 :         ValVec.emplace_back(OutArg, ReplVal);
     339          42 :         Store.second->eraseFromParent();
     340             :       }
     341             : 
     342             :       if (ThisReplaceable) {
     343          40 :         ReturnTypes.push_back(ArgTy);
     344          40 :         OutArgIndexes.insert(OutArg->getArgNo());
     345          40 :         ++NumOutArgumentsReplaced;
     346          40 :         Changing = true;
     347             :       }
     348             :     }
     349             :   } while (Changing);
     350             : 
     351          40 :   if (Replacements.empty())
     352             :     return false;
     353             : 
     354          36 :   LLVMContext &Ctx = F.getParent()->getContext();
     355          72 :   StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
     356             : 
     357         108 :   FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
     358             :                                               F.getFunctionType()->params(),
     359          72 :                                               F.isVarArg());
     360             : 
     361             :   DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
     362             : 
     363             :   Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
     364         108 :                                        F.getName() + ".body");
     365         108 :   F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
     366          36 :   NewFunc->copyAttributesFrom(&F);
     367          72 :   NewFunc->setComdat(F.getComdat());
     368             : 
     369             :   // We want to preserve the function and param attributes, but need to strip
     370             :   // off any return attributes, e.g. zeroext doesn't make sense with a struct.
     371          36 :   NewFunc->stealArgumentListFrom(F);
     372             : 
     373          36 :   AttrBuilder RetAttrs;
     374          36 :   RetAttrs.addAttribute(Attribute::SExt);
     375          36 :   RetAttrs.addAttribute(Attribute::ZExt);
     376          36 :   RetAttrs.addAttribute(Attribute::NoAlias);
     377          36 :   NewFunc->removeAttributes(AttributeList::ReturnIndex, RetAttrs);
     378             :   // TODO: How to preserve metadata?
     379             : 
     380             :   // Move the body of the function into the new rewritten function, and replace
     381             :   // this function with a stub.
     382         108 :   NewFunc->getBasicBlockList().splice(NewFunc->begin(), F.getBasicBlockList());
     383             : 
     384         146 :   for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
     385          38 :     ReturnInst *RI = Replacement.first;
     386          76 :     IRBuilder<> B(RI);
     387         152 :     B.SetCurrentDebugLocation(RI->getDebugLoc());
     388             : 
     389          38 :     int RetIdx = 0;
     390          38 :     Value *NewRetVal = UndefValue::get(NewRetTy);
     391             : 
     392           7 :     Value *RetVal = RI->getReturnValue();
     393             :     if (RetVal)
     394          14 :       NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
     395             : 
     396         156 :     for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
     397          42 :       Argument *Arg = ReturnPoint.first;
     398          42 :       Value *Val = ReturnPoint.second;
     399          84 :       Type *EltTy = Arg->getType()->getPointerElementType();
     400          42 :       if (Val->getType() != EltTy) {
     401          13 :         Type *EffectiveEltTy = EltTy;
     402           8 :         if (StructType *CT = dyn_cast<StructType>(EltTy)) {
     403             :           assert(CT->getNumContainedTypes() == 1);
     404          16 :           EffectiveEltTy = CT->getContainedType(0);
     405             :         }
     406             : 
     407          26 :         if (DL->getTypeSizeInBits(EffectiveEltTy) !=
     408          13 :             DL->getTypeSizeInBits(Val->getType())) {
     409             :           assert(isVec3ToVec4Shuffle(EffectiveEltTy, Val->getType()));
     410          16 :           Val = B.CreateShuffleVector(Val, UndefValue::get(Val->getType()),
     411             :                                       { 0, 1, 2 });
     412             :         }
     413             : 
     414          26 :         Val = B.CreateBitCast(Val, EffectiveEltTy);
     415             : 
     416             :         // Re-create single element composite.
     417          13 :         if (EltTy != EffectiveEltTy)
     418          16 :           Val = B.CreateInsertValue(UndefValue::get(EltTy), Val, 0);
     419             :       }
     420             : 
     421          84 :       NewRetVal = B.CreateInsertValue(NewRetVal, Val, RetIdx++);
     422             :     }
     423             : 
     424          38 :     if (RetVal)
     425             :       RI->setOperand(0, NewRetVal);
     426             :     else {
     427          31 :       B.CreateRet(NewRetVal);
     428          31 :       RI->eraseFromParent();
     429             :     }
     430             :   }
     431             : 
     432          72 :   SmallVector<Value *, 16> StubCallArgs;
     433          90 :   for (Argument &Arg : F.args()) {
     434          54 :     if (OutArgIndexes.count(Arg.getArgNo())) {
     435             :       // It's easier to preserve the type of the argument list. We rely on
     436             :       // DeadArgumentElimination to take care of these.
     437          40 :       StubCallArgs.push_back(UndefValue::get(Arg.getType()));
     438             :     } else {
     439          14 :       StubCallArgs.push_back(&Arg);
     440             :     }
     441             :   }
     442             : 
     443          72 :   BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
     444         108 :   IRBuilder<> B(StubBB);
     445         108 :   CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
     446             : 
     447          72 :   int RetIdx = RetTy->isVoidTy() ? 0 : 1;
     448          90 :   for (Argument &Arg : F.args()) {
     449          54 :     if (!OutArgIndexes.count(Arg.getArgNo()))
     450          14 :       continue;
     451             : 
     452          80 :     PointerType *ArgType = cast<PointerType>(Arg.getType());
     453             : 
     454          40 :     auto *EltTy = ArgType->getElementType();
     455          40 :     unsigned Align = Arg.getParamAlignment();
     456          40 :     if (Align == 0)
     457          39 :       Align = DL->getABITypeAlignment(EltTy);
     458             : 
     459          80 :     Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
     460          80 :     Type *PtrTy = Val->getType()->getPointerTo(ArgType->getAddressSpace());
     461             : 
     462             :     // We can peek through bitcasts, so the type may not match.
     463          80 :     Value *PtrVal = B.CreateBitCast(&Arg, PtrTy);
     464             : 
     465             :     B.CreateAlignedStore(Val, PtrVal, Align);
     466             :   }
     467             : 
     468          72 :   if (!RetTy->isVoidTy()) {
     469          14 :     B.CreateRet(B.CreateExtractValue(StubCall, 0));
     470             :   } else {
     471          29 :     B.CreateRetVoid();
     472             :   }
     473             : 
     474             :   // The function is now a stub we want to inline.
     475          36 :   F.addFnAttr(Attribute::AlwaysInline);
     476             : 
     477          36 :   ++NumOutArgumentFunctionsReplaced;
     478          36 :   return true;
     479             : }
     480             : 
     481           0 : FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
     482           0 :   return new AMDGPURewriteOutArguments();
     483      216918 : }

Generated by: LCOV version 1.13