AMDGPU: Add pass to replace out arguments
It is better to return arguments directly in registers
if we are making a call rather than introducing expensive
stack usage. In one of sample compile from one of
Blender's many kernel variants, this fires on about
~20 different functions. Future improvements may be to
recognize simple cases where the pointer is indexing a small
array. This also fails when the store to the out argument
is in a separate block from the return, which happens in
a few of the Blender functions. This should also probably
be using MemorySSA which might help with that.
I'm not sure this is correct as a FunctionPass, but
MemoryDependenceAnalysis seems to not work with
a ModulePass.
I'm also not sure where it should run.I think it should
run before DeadArgumentElimination, so maybe either
EP_CGSCCOptimizerLate or EP_ScalarOptimizerLate.
llvm-svn: 309416
diff --git a/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
new file mode 100644
index 0000000..4d03040
--- /dev/null
+++ b/llvm/lib/Target/AMDGPU/AMDGPURewriteOutArguments.cpp
@@ -0,0 +1,375 @@
+//===-- AMDGPURewriteOutArgumentsPass.cpp - Create struct returns ---------===//
+//
+// The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file This pass attempts to replace out argument usage with a return of a
+/// struct.
+///
+/// We can support returning a lot of values directly in registers, but
+/// idiomatic C code frequently uses a pointer argument to return a second value
+/// rather than returning a struct by value. GPU stack access is also quite
+/// painful, so we want to avoid that if possible. Passing a stack object
+/// pointer to a function also requires an additional address expansion code
+/// sequence to convert the pointer to be relative to the kernel's scratch wave
+/// offset register since the callee doesn't know what stack frame the incoming
+/// pointer is relative to.
+///
+/// The goal is to try rewriting code that looks like this:
+///
+/// int foo(int a, int b, int* out) {
+/// *out = bar();
+/// return a + b;
+/// }
+///
+/// into something like this:
+///
+/// std::pair<int, int> foo(int a, int b) {
+/// return std::make_pair(a + b, bar());
+/// }
+///
+/// Typically the incoming pointer is a simple alloca for a temporary variable
+/// to use the API, which if replaced with a struct return will be easily SROA'd
+/// out when the stub function we create is inlined
+///
+/// This pass introduces the struct return, but leaves the unused pointer
+/// arguments and introduces a new stub function calling the struct returning
+/// body. DeadArgumentElimination should be run after this to clean these up.
+//
+//===----------------------------------------------------------------------===//
+
+#include "AMDGPU.h"
+#include "Utils/AMDGPUBaseInfo.h"
+
+#include "llvm/Analysis/MemoryDependenceAnalysis.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "amdgpu-rewrite-out-arguments"
+
+using namespace llvm;
+
+namespace {
+
+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));
+
+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));
+
+STATISTIC(NumOutArgumentsReplaced,
+ "Number out arguments moved to struct return values");
+STATISTIC(NumOutArgumentFunctionsReplaced,
+ "Number of functions with out arguments moved to struct return values");
+
+class AMDGPURewriteOutArguments : public FunctionPass {
+private:
+ const DataLayout *DL = nullptr;
+ MemoryDependenceResults *MDA = nullptr;
+
+ bool isOutArgumentCandidate(Argument &Arg) const;
+
+public:
+ static char ID;
+
+ AMDGPURewriteOutArguments() :
+ FunctionPass(ID) {}
+
+ void getAnalysisUsage(AnalysisUsage &AU) const override {
+ AU.addRequired<MemoryDependenceWrapperPass>();
+ FunctionPass::getAnalysisUsage(AU);
+ }
+
+ bool doInitialization(Module &M) override;
+ bool runOnFunction(Function &M) override;
+};
+
+} // End anonymous namespace
+
+INITIALIZE_PASS_BEGIN(AMDGPURewriteOutArguments, DEBUG_TYPE,
+ "AMDGPU Rewrite Out Arguments", false, false)
+INITIALIZE_PASS_DEPENDENCY(MemoryDependenceWrapperPass)
+INITIALIZE_PASS_END(AMDGPURewriteOutArguments, DEBUG_TYPE,
+ "AMDGPU Rewrite Out Arguments", false, false)
+
+char AMDGPURewriteOutArguments::ID = 0;
+
+bool AMDGPURewriteOutArguments::isOutArgumentCandidate(Argument &Arg) const {
+ const int MaxUses = 10;
+ const unsigned MaxOutArgSizeBytes = 4 * MaxNumRetRegs;
+ int UseCount = 0;
+
+ PointerType *ArgTy = dyn_cast<PointerType>(Arg.getType());
+
+ // TODO: It might be useful for any out arguments, not just privates.
+ if (!ArgTy || (ArgTy->getAddressSpace() != DL->getAllocaAddrSpace() &&
+ !AnyAddressSpace) ||
+ Arg.hasByValAttr() || Arg.hasStructRetAttr() ||
+ DL->getTypeStoreSize(ArgTy->getPointerElementType()) > MaxOutArgSizeBytes) {
+ return false;
+ }
+
+ for (Use &U : Arg.uses()) {
+ StoreInst *SI = dyn_cast<StoreInst>(U.getUser());
+ if (UseCount > MaxUses)
+ return false;
+
+ if (!SI || !SI->isSimple() ||
+ U.getOperandNo() != StoreInst::getPointerOperandIndex())
+ return false;
+
+ ++UseCount;
+ }
+
+ // Skip unused arguments.
+ return UseCount > 0;
+}
+
+bool AMDGPURewriteOutArguments::doInitialization(Module &M) {
+ DL = &M.getDataLayout();
+ return false;
+}
+
+bool AMDGPURewriteOutArguments::runOnFunction(Function &F) {
+ if (skipFunction(F))
+ return false;
+
+ // TODO: Could probably handle variadic functions.
+ if (F.isVarArg() || F.hasStructRetAttr() ||
+ AMDGPU::isEntryFunctionCC(F.getCallingConv()))
+ return false;
+
+ MDA = &getAnalysis<MemoryDependenceWrapperPass>().getMemDep();
+
+ unsigned ReturnNumRegs = 0;
+ SmallSet<int, 4> OutArgIndexes;
+ SmallVector<Type *, 4> ReturnTypes;
+ Type *RetTy = F.getReturnType();
+ if (!RetTy->isVoidTy()) {
+ ReturnNumRegs = DL->getTypeStoreSize(RetTy) / 4;
+
+ if (ReturnNumRegs >= MaxNumRetRegs)
+ return false;
+
+ ReturnTypes.push_back(RetTy);
+ }
+
+ SmallVector<Argument *, 4> OutArgs;
+ for (Argument &Arg : F.args()) {
+ if (isOutArgumentCandidate(Arg)) {
+ DEBUG(dbgs() << "Found possible out argument " << Arg
+ << " in function " << F.getName() << '\n');
+ OutArgs.push_back(&Arg);
+ }
+ }
+
+ if (OutArgs.empty())
+ return false;
+
+ typedef SmallVector<std::pair<Argument *, Value *>, 4> ReplacementVec;
+ DenseMap<ReturnInst *, ReplacementVec> Replacements;
+
+ SmallVector<ReturnInst *, 4> Returns;
+ for (BasicBlock &BB : F) {
+ if (ReturnInst *RI = dyn_cast<ReturnInst>(&BB.back()))
+ Returns.push_back(RI);
+ }
+
+ if (Returns.empty())
+ return false;
+
+ bool Changing;
+
+ do {
+ Changing = false;
+
+ // Keep retrying if we are able to successfully eliminate an argument. This
+ // helps with cases with multiple arguments which may alias, such as in a
+ // sincos implemntation. If we have 2 stores to arguments, on the first
+ // attempt the MDA query will succeed for the second store but not the
+ // first. On the second iteration we've removed that out clobbering argument
+ // (by effectively moving it into another function) and will find the second
+ // argument is OK to move.
+ for (Argument *OutArg : OutArgs) {
+ bool ThisReplaceable = true;
+ SmallVector<std::pair<ReturnInst *, StoreInst *>, 4> ReplaceableStores;
+
+ Type *ArgTy = OutArg->getType()->getPointerElementType();
+
+ // Skip this argument if converting it will push us over the register
+ // count to return limit.
+
+ // TODO: This is an approximation. When legalized this could be more. We
+ // can ask TLI for exactly how many.
+ unsigned ArgNumRegs = DL->getTypeStoreSize(ArgTy) / 4;
+ if (ArgNumRegs + ReturnNumRegs > MaxNumRetRegs)
+ continue;
+
+ // An argument is convertible only if all exit blocks are able to replace
+ // it.
+ for (ReturnInst *RI : Returns) {
+ BasicBlock *BB = RI->getParent();
+
+ MemDepResult Q = MDA->getPointerDependencyFrom(MemoryLocation(OutArg),
+ true, BB->end(), BB, RI);
+ StoreInst *SI = nullptr;
+ if (Q.isDef())
+ SI = dyn_cast<StoreInst>(Q.getInst());
+
+ if (SI) {
+ DEBUG(dbgs() << "Found out argument store: " << *SI << '\n');
+ ReplaceableStores.emplace_back(RI, SI);
+ } else {
+ ThisReplaceable = false;
+ break;
+ }
+ }
+
+ if (!ThisReplaceable)
+ continue; // Try the next argument candidate.
+
+ for (std::pair<ReturnInst *, StoreInst *> Store : ReplaceableStores) {
+ Value *ReplVal = Store.second->getValueOperand();
+
+ auto &ValVec = Replacements[Store.first];
+ if (llvm::find_if(ValVec,
+ [OutArg](const std::pair<Argument *, Value *> &Entry) {
+ return Entry.first == OutArg;}) != ValVec.end()) {
+ DEBUG(dbgs() << "Saw multiple out arg stores" << *OutArg << '\n');
+ // It is possible to see stores to the same argument multiple times,
+ // but we expect these would have been optimized out already.
+ ThisReplaceable = false;
+ break;
+ }
+
+ ValVec.emplace_back(OutArg, ReplVal);
+ Store.second->eraseFromParent();
+ }
+
+ if (ThisReplaceable) {
+ ReturnTypes.push_back(ArgTy);
+ OutArgIndexes.insert(OutArg->getArgNo());
+ ++NumOutArgumentsReplaced;
+ Changing = true;
+ }
+ }
+ } while (Changing);
+
+ if (Replacements.empty())
+ return false;
+
+ LLVMContext &Ctx = F.getParent()->getContext();
+ StructType *NewRetTy = StructType::create(Ctx, ReturnTypes, F.getName());
+
+ FunctionType *NewFuncTy = FunctionType::get(NewRetTy,
+ F.getFunctionType()->params(),
+ F.isVarArg());
+
+ DEBUG(dbgs() << "Computed new return type: " << *NewRetTy << '\n');
+
+ Function *NewFunc = Function::Create(NewFuncTy, Function::PrivateLinkage,
+ F.getName() + ".body");
+ F.getParent()->getFunctionList().insert(F.getIterator(), NewFunc);
+ NewFunc->copyAttributesFrom(&F);
+ NewFunc->setComdat(F.getComdat());
+
+ // We want to preserve the function and param attributes, but need to strip
+ // off any return attributes, e.g. zeroext doesn't make sense with a struct.
+ NewFunc->stealArgumentListFrom(F);
+
+ AttrBuilder RetAttrs;
+ RetAttrs.addAttribute(Attribute::SExt);
+ RetAttrs.addAttribute(Attribute::ZExt);
+ RetAttrs.addAttribute(Attribute::NoAlias);
+ NewFunc->removeAttributes(AttributeList::ReturnIndex, RetAttrs);
+ // TODO: How to preserve metadata?
+
+ // Move the body of the function into the new rewritten function, and replace
+ // this function with a stub.
+ NewFunc->getBasicBlockList().splice(NewFunc->begin(), F.getBasicBlockList());
+
+ for (std::pair<ReturnInst *, ReplacementVec> &Replacement : Replacements) {
+ ReturnInst *RI = Replacement.first;
+ IRBuilder<> B(RI);
+ B.SetCurrentDebugLocation(RI->getDebugLoc());
+
+ int RetIdx = 0;
+ Value *NewRetVal = UndefValue::get(NewRetTy);
+
+ Value *RetVal = RI->getReturnValue();
+ if (RetVal)
+ NewRetVal = B.CreateInsertValue(NewRetVal, RetVal, RetIdx++);
+
+ for (std::pair<Argument *, Value *> ReturnPoint : Replacement.second) {
+ NewRetVal = B.CreateInsertValue(NewRetVal, ReturnPoint.second, RetIdx++);
+ }
+
+ if (RetVal)
+ RI->setOperand(0, NewRetVal);
+ else {
+ B.CreateRet(NewRetVal);
+ RI->eraseFromParent();
+ }
+ }
+
+ SmallVector<Value *, 16> StubCallArgs;
+ for (Argument &Arg : F.args()) {
+ if (OutArgIndexes.count(Arg.getArgNo())) {
+ // It's easier to preserve the type of the argument list. We rely on
+ // DeadArgumentElimination to take care of these.
+ StubCallArgs.push_back(UndefValue::get(Arg.getType()));
+ } else {
+ StubCallArgs.push_back(&Arg);
+ }
+ }
+
+ BasicBlock *StubBB = BasicBlock::Create(Ctx, "", &F);
+ IRBuilder<> B(StubBB);
+ CallInst *StubCall = B.CreateCall(NewFunc, StubCallArgs);
+
+ int RetIdx = RetTy->isVoidTy() ? 0 : 1;
+ for (Argument &Arg : F.args()) {
+ if (!OutArgIndexes.count(Arg.getArgNo()))
+ continue;
+
+ auto *EltTy = Arg.getType()->getPointerElementType();
+ unsigned Align = Arg.getParamAlignment();
+ if (Align == 0)
+ Align = DL->getABITypeAlignment(EltTy);
+
+ Value *Val = B.CreateExtractValue(StubCall, RetIdx++);
+ B.CreateAlignedStore(Val, &Arg, Align);
+ }
+
+ if (!RetTy->isVoidTy()) {
+ B.CreateRet(B.CreateExtractValue(StubCall, 0));
+ } else {
+ B.CreateRetVoid();
+ }
+
+ // The function is now a stub we want to inline.
+ F.addFnAttr(Attribute::AlwaysInline);
+
+ ++NumOutArgumentFunctionsReplaced;
+ return true;
+}
+
+FunctionPass *llvm::createAMDGPURewriteOutArgumentsPass() {
+ return new AMDGPURewriteOutArguments();
+}