llvm/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::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/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
  59 using namespace llvm;
  60
  61 static 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
  68 static 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
  74 STATISTIC(NumOutArgumentsReplaced,
  75           "Number out arguments moved to struct return values");
  76 STATISTIC(NumOutArgumentFunctionsReplaced,
  77           "Number of functions with out arguments moved to struct return values");
  78
  79 namespace {
  80
  81 class AMDGPURewriteOutArguments : public FunctionPass {
  82 private:
  83   const DataLayout *DL = nullptr;
  84   MemoryDependenceResults *MDA = nullptr;
  85
  86   Type *getStoredType(Value &Arg) const;
  87   Type *getOutArgumentType(Argument &Arg) const;
  88
  89 public:
  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
 105 INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
 106                       "AMDGPU Rewrite Out Arguments", false, false)
 107 INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
 108 INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
 109                     "AMDGPU Rewrite Out Arguments", false, false)
 110
 111 char AMDGPURewriteOutArguments::ID = 0;
 112
 113 Type *AMDGPURewriteOutArguments::getStoredType(Value &Arg) const {
 114   const int MaxUses = 10;
 115   int UseCount = 0;
 116
 117   SmallVector<Use *> Worklist;
 118   for (Use &U : Arg.uses())
 119     Worklist.push_back(&U);
 120
 121   Type *StoredType = nullptr;
 122   while (!Worklist.empty()) {
 123     Use *U = Worklist.pop_back_val();
 124
 125     if (auto *BCI = dyn_cast<BitCastInst>(U->getUser())) {
 126       for (Use &U : BCI->uses())
 127         Worklist.push_back(&U);
 128       continue;
 129     }
 130
 131     if (auto *SI = dyn_cast<StoreInst>(U->getUser())) {
 132       if (UseCount++ > MaxUses)
 133         return nullptr;
 134
 135       if (!SI->isSimple() ||
 136           U->getOperandNo() != StoreInst::getPointerOperandIndex())
 137         return nullptr;
 138
 139       if (StoredType && StoredType != SI->getValueOperand()->getType())
 140         return nullptr; // More than one type.
 141       StoredType = SI->getValueOperand()->getType();
 142       continue;
 143     }
 144
 145     // Unsupported user.
 146     return nullptr;
 147   }
 148
 149   return StoredType;
 150 }
 151
 152 Type *AMDGPURewriteOutArguments::getOutArgumentType(Argument &Arg) const {
 153   const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
 154   PointerType *ArgTy = dyn_cast<PointerType>(Arg.getType());
 155
 156   // TODO: It might be useful for any out arguments, not just privates.
 157   if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
 158                  !AnyAddressSpace) ||
 159       Arg.hasByValAttr() || Arg.hasStructRetAttr()) {
 160     return nullptr;
 161   }
 162
 163   Type *StoredType = getStoredType(Arg);
 164   if (!StoredType || DL->getTypeStoreSize(StoredType) > MaxOutArgSizeBytes)
 165     return nullptr;
 166
 167   return StoredType;
 168 }
 169
 170 bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
 171   DL = &M.getDataLayout();
 172   return false;
 173 }
 174
 175 bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
 176   if (skipFunction(F))
 177     return false;
 178
 179   // TODO: Could probably handle variadic functions.
 180   if (F.isVarArg() || F.hasStructRetAttr() ||
 181       AMDGPU::isEntryFunctionCC(F.getCallingConv()))
 182     return false;
 183
 184   MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
 185
 186   unsigned ReturnNumRegs = 0;
 187   SmallDenseMap<int, Type *, 4> OutArgIndexes;
 188   SmallVector<Type *, 4> ReturnTypes;
 189   Type *RetTy = F.getReturnType();
 190   if (!RetTy->isVoidTy()) {
 191     ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
 192
 193     if (ReturnNumRegs >= MaxNumRetRegs)
 194       return false;
 195
 196     ReturnTypes.push_back(RetTy);
 197   }
 198
 199   SmallVector<std::pair<Argument *, Type *>, 4> OutArgs;
 200   for (Argument &Arg : F.args()) {
 201     if (Type *Ty = getOutArgumentType(Arg)) {
 202       LLVM_DEBUG(dbgs() << "Found possible out argument " << Arg
 203                         << " in function " << F.getName() << '\n');
 204       OutArgs.push_back({&Arg, Ty});
 205     }
 206   }
 207
 208   if (OutArgs.empty())
 209     return false;
 210
 211   using ReplacementVec = SmallVector<std::pair<Argument *, Value *>, 4>;
 212
 213   DenseMap<ReturnInst *, ReplacementVec> Replacements;
 214
 215   SmallVector<ReturnInst *, 4> Returns;
 216   for (BasicBlock &BB : F) {
 217     if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
 218       Returns.push_back(RI);
 219   }
 220
 221   if (Returns.empty())
 222     return false;
 223
 224   bool Changing;
 225
 226   do {
 227     Changing = false;
 228
 229     // Keep retrying if we are able to successfully eliminate an argument. This
 230     // helps with cases with multiple arguments which may alias, such as in a
 231     // sincos implementation. If we have 2 stores to arguments, on the first
 232     // attempt the MDA query will succeed for the second store but not the
 233     // first. On the second iteration we've removed that out clobbering argument
 234     // (by effectively moving it into another function) and will find the second
 235     // argument is OK to move.
 236     for (const auto &Pair : OutArgs) {
 237       bool ThisReplaceable = true;
 238       SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
 239
 240       Argument *OutArg = Pair.first;
 241       Type *ArgTy = Pair.second;
 242
 243       // Skip this argument if converting it will push us over the register
 244       // count to return limit.
 245
 246       // TODO: This is an approximation. When legalized this could be more. We
 247       // can ask TLI for exactly how many.
 248       unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
 249       if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
 250         continue;
 251
 252       // An argument is convertible only if all exit blocks are able to replace
 253       // it.
 254       for (ReturnInst *RI : Returns) {
 255         BasicBlock *BB = RI->getParent();
 256
 257         MemDepResult Q = MDA->getPointerDependencyFrom(
 258             MemoryLocation::getBeforeOrAfter(OutArg), true, BB->end(), BB, RI);
 259         StoreInst *SI = nullptr;
 260         if (Q.isDef())
 261           SI = dyn_cast<StoreInst>(Q.getInst());
 262
 263         if (SI) {
 264           LLVM_DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
 265           ReplaceableStores.emplace_back(RI, SI);
 266         } else {
 267           ThisReplaceable = false;
 268           break;
 269         }
 270       }
 271
 272       if (!ThisReplaceable)
 273         continue; // Try the next argument candidate.
 274
 275       for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
 276         Value *ReplVal = Store.second->getValueOperand();
 277
 278         auto &ValVec = Replacements[Store.first];
 279         if (llvm::any_of(ValVec,
 280                          [OutArg](const std::pair<Argument *, Value *> &Entry) {
 281                            return Entry.first == OutArg;
 282                          })) {
 283           LLVM_DEBUG(dbgs()
 284                      << "Saw multiple out arg stores" << *OutArg << '\n');
 285           // It is possible to see stores to the same argument multiple times,
 286           // but we expect these would have been optimized out already.
 287           ThisReplaceable = false;
 288           break;
 289         }
 290
 291         ValVec.emplace_back(OutArg, ReplVal);
 292         Store.second->eraseFromParent();
 293       }
 294
 295       if (ThisReplaceable) {
 296         ReturnTypes.push_back(ArgTy);
 297         OutArgIndexes.insert({OutArg->getArgNo(), ArgTy});
 298         ++NumOutArgumentsReplaced;
 299         Changing = true;
 300       }
 301     }
 302   } while (Changing);
 303
 304   if (Replacements.empty())
 305     return false;
 306
 307   LLVMContext &Ctx = F.getParent()->getContext();
 308   StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
 309
 310   FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
 311                                               F.getFunctionType()->params(),
 312                                               F.isVarArg());
 313
 314   LLVM_DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
 315
 316   Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
 317                                        F.getName() + ".body");
 318   F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
 319   NewFunc->copyAttributesFrom(&F);
 320   NewFunc->setComdat(F.getComdat());
 321
 322   // We want to preserve the function and param attributes, but need to strip
 323   // off any return attributes, e.g. zeroext doesn't make sense with a struct.
 324   NewFunc->stealArgumentListFrom(F);
 325
 326   AttributeMask RetAttrs;
 327   RetAttrs.addAttribute(Attribute::SExt);
 328   RetAttrs.addAttribute(Attribute::ZExt);
 329   RetAttrs.addAttribute(Attribute::NoAlias);
 330   NewFunc->removeRetAttrs(RetAttrs);
 331   // TODO: How to preserve metadata?
 332
 333   NewFunc->setIsNewDbgInfoFormat(F.IsNewDbgInfoFormat);
 334
 335   // Move the body of the function into the new rewritten function, and replace
 336   // this function with a stub.
 337   NewFunc->splice(NewFunc->begin(), &F);
 338
 339   for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
 340     ReturnInst *RI = Replacement.first;
 341     IRBuilder<> B(RI);
 342     B.SetCurrentDebugLocation(RI->getDebugLoc());
 343
 344     int RetIdx = 0;
 345     Value *NewRetVal = PoisonValue::get(NewRetTy);
 346
 347     Value *RetVal = RI->getReturnValue();
 348     if (RetVal)
 349       NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
 350
 351     for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second)
 352       NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, RetIdx++);
 353
 354     if (RetVal)
 355       RI->setOperand(0, NewRetVal);
 356     else {
 357       B.CreateRet(NewRetVal);
 358       RI->eraseFromParent();
 359     }
 360   }
 361
 362   SmallVector<Value *, 16> StubCallArgs;
 363   for (Argument &Arg : F.args()) {
 364     if (OutArgIndexes.count(Arg.getArgNo())) {
 365       // It's easier to preserve the type of the argument list. We rely on
 366       // DeadArgumentElimination to take care of these.
 367       StubCallArgs.push_back(PoisonValue::get(Arg.getType()));
 368     } else {
 369       StubCallArgs.push_back(&Arg);
 370     }
 371   }
 372
 373   BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
 374   IRBuilder<> B(StubBB);
 375   CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
 376
 377   int RetIdx = RetTy->isVoidTy() ? 0 : 1;
 378   for (Argument &Arg : F.args()) {
 379     if (!OutArgIndexes.count(Arg.getArgNo()))
 380       continue;
 381
 382     Type *EltTy = OutArgIndexes[Arg.getArgNo()];
 383     const auto Align =
 384         DL->getValueOrABITypeAlignment(Arg.getParamAlign(), EltTy);
 385
 386     Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
 387     B.CreateAlignedStore(Val, &Arg, Align);
 388   }
 389
 390   if (!RetTy->isVoidTy()) {
 391     B.CreateRet(B.CreateExtractValue(StubCall, 0));
 392   } else {
 393     B.CreateRetVoid();
 394   }
 395
 396   // The function is now a stub we want to inline.
 397   F.addFnAttr(Attribute::AlwaysInline);
 398
 399   ++NumOutArgumentFunctionsReplaced;
 400   return true;
 401 }
 402
 403 FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
 404   return new AMDGPURewriteOutArguments();
 405 }