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