shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
llvm-project/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
Rahul Joshi fa789dffb1
[NFC] Rename Intrinsic::getDeclaration to getOrInsertDeclaration (#111752) 
Rename the function to reflect its correct behavior and to be consistent
with `Module::getOrInsertFunction`. This is also in preparation of
adding a new `Intrinsic::getDeclaration` that will have behavior similar
to `Module::getFunction` (i.e, just lookup, no creation).
2024-10-11 05:26:03 -07:00

733 lines
28 KiB
C++
Raw Blame History

//===-- NVPTXLowerArgs.cpp - Lower arguments ------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
//
// Arguments to kernel and device functions are passed via param space,
// which imposes certain restrictions:
// http://docs.nvidia.com/cuda/parallel-thread-execution/#state-spaces
//
// Kernel parameters are read-only and accessible only via ld.param
// instruction, directly or via a pointer.
//
// Device function parameters are directly accessible via
// ld.param/st.param, but taking the address of one returns a pointer
// to a copy created in local space which *can't* be used with
// ld.param/st.param.
//
// Copying a byval struct into local memory in IR allows us to enforce
// the param space restrictions, gives the rest of IR a pointer w/o
// param space restrictions, and gives us an opportunity to eliminate
// the copy.
//
// Pointer arguments to kernel functions need more work to be lowered:
//
// 1. Convert non-byval pointer arguments of CUDA kernels to pointers in the
// global address space. This allows later optimizations to emit
// ld.global.*/st.global.* for accessing these pointer arguments. For
// example,
//
// define void @foo(float* %input) {
// %v = load float, float* %input, align 4
// ...
// }
//
// becomes
//
// define void @foo(float* %input) {
// %input2 = addrspacecast float* %input to float addrspace(1)*
// %input3 = addrspacecast float addrspace(1)* %input2 to float*
// %v = load float, float* %input3, align 4
// ...
// }
//
// Later, NVPTXInferAddressSpaces will optimize it to
//
// define void @foo(float* %input) {
// %input2 = addrspacecast float* %input to float addrspace(1)*
// %v = load float, float addrspace(1)* %input2, align 4
// ...
// }
//
// 2. Convert byval kernel parameters to pointers in the param address space
// (so that NVPTX emits ld/st.param). Convert pointers *within* a byval
// kernel parameter to pointers in the global address space. This allows
// NVPTX to emit ld/st.global.
//
// struct S {
// int *x;
// int *y;
// };
// __global__ void foo(S s) {
// int *b = s.y;
// // use b
// }
//
// "b" points to the global address space. In the IR level,
//
// define void @foo(ptr byval %input) {
// %b_ptr = getelementptr {ptr, ptr}, ptr %input, i64 0, i32 1
// %b = load ptr, ptr %b_ptr
// ; use %b
// }
//
// becomes
//
// define void @foo({i32*, i32*}* byval %input) {
// %b_param = addrspacecat ptr %input to ptr addrspace(101)
// %b_ptr = getelementptr {ptr, ptr}, ptr addrspace(101) %b_param, i64 0, i32 1
// %b = load ptr, ptr addrspace(101) %b_ptr
// %b_global = addrspacecast ptr %b to ptr addrspace(1)
// ; use %b_generic
// }
//
// Create a local copy of kernel byval parameters used in a way that *might* mutate
// the parameter, by storing it in an alloca. Mutations to "grid_constant" parameters
// are undefined behaviour, and don't require local copies.
//
// define void @foo(ptr byval(%struct.s) align 4 %input) {
// store i32 42, ptr %input
// ret void
// }
//
// becomes
//
// define void @foo(ptr byval(%struct.s) align 4 %input) #1 {
// %input1 = alloca %struct.s, align 4
// %input2 = addrspacecast ptr %input to ptr addrspace(101)
// %input3 = load %struct.s, ptr addrspace(101) %input2, align 4
// store %struct.s %input3, ptr %input1, align 4
// store i32 42, ptr %input1, align 4
// ret void
// }
//
// If %input were passed to a device function, or written to memory,
// conservatively assume that %input gets mutated, and create a local copy.
//
// Convert param pointers to grid_constant byval kernel parameters that are
// passed into calls (device functions, intrinsics, inline asm), or otherwise
// "escape" (into stores/ptrtoints) to the generic address space, using the
// `nvvm.ptr.param.to.gen` intrinsic, so that NVPTX emits cvta.param
// (available for sm70+)
//
// define void @foo(ptr byval(%struct.s) %input) {
// ; %input is a grid_constant
// %call = call i32 @escape(ptr %input)
// ret void
// }
//
// becomes
//
// define void @foo(ptr byval(%struct.s) %input) {
// %input1 = addrspacecast ptr %input to ptr addrspace(101)
// ; the following intrinsic converts pointer to generic. We don't use an addrspacecast
// ; to prevent generic -> param -> generic from getting cancelled out
// %input1.gen = call ptr @llvm.nvvm.ptr.param.to.gen.p0.p101(ptr addrspace(101) %input1)
// %call = call i32 @escape(ptr %input1.gen)
// ret void
// }
//
// TODO: merge this pass with NVPTXInferAddressSpaces so that other passes don't
// cancel the addrspacecast pair this pass emits.
//===----------------------------------------------------------------------===//
#include "MCTargetDesc/NVPTXBaseInfo.h"
#include "NVPTX.h"
#include "NVPTXTargetMachine.h"
#include "NVPTXUtilities.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Analysis/PtrUseVisitor.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/IntrinsicsNVPTX.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Type.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include <numeric>
#include <queue>
#define DEBUG_TYPE "nvptx-lower-args"
using namespace llvm;
namespace llvm {
void initializeNVPTXLowerArgsPass(PassRegistry &);
}
namespace {
class NVPTXLowerArgs : public FunctionPass {
bool runOnFunction(Function &F) override;
bool runOnKernelFunction(const NVPTXTargetMachine &TM, Function &F);
bool runOnDeviceFunction(const NVPTXTargetMachine &TM, Function &F);
// handle byval parameters
void handleByValParam(const NVPTXTargetMachine &TM, Argument *Arg);
// Knowing Ptr must point to the global address space, this function
// addrspacecasts Ptr to global and then back to generic. This allows
// NVPTXInferAddressSpaces to fold the global-to-generic cast into
// loads/stores that appear later.
void markPointerAsGlobal(Value *Ptr);
public:
static char ID; // Pass identification, replacement for typeid
NVPTXLowerArgs() : FunctionPass(ID) {}
StringRef getPassName() const override {
return "Lower pointer arguments of CUDA kernels";
}
void getAnalysisUsage(AnalysisUsage &AU) const override {
AU.addRequired<TargetPassConfig>();
}
};
} // namespace
char NVPTXLowerArgs::ID = 1;
INITIALIZE_PASS_BEGIN(NVPTXLowerArgs, "nvptx-lower-args",
"Lower arguments (NVPTX)", false, false)
INITIALIZE_PASS_DEPENDENCY(TargetPassConfig)
INITIALIZE_PASS_END(NVPTXLowerArgs, "nvptx-lower-args",
"Lower arguments (NVPTX)", false, false)
// =============================================================================
// If the function had a byval struct ptr arg, say foo(%struct.x* byval %d),
// and we can't guarantee that the only accesses are loads,
// then add the following instructions to the first basic block:
//
// %temp = alloca %struct.x, align 8
// %tempd = addrspacecast %struct.x* %d to %struct.x addrspace(101)*
// %tv = load %struct.x addrspace(101)* %tempd
// store %struct.x %tv, %struct.x* %temp, align 8
//
// The above code allocates some space in the stack and copies the incoming
// struct from param space to local space.
// Then replace all occurrences of %d by %temp.
//
// In case we know that all users are GEPs or Loads, replace them with the same
// ones in parameter AS, so we can access them using ld.param.
// =============================================================================
// For Loads, replaces the \p OldUse of the pointer with a Use of the same
// pointer in parameter AS.
// For "escapes" (to memory, a function call, or a ptrtoint), cast the OldUse to
// generic using cvta.param.
static void convertToParamAS(Use *OldUse, Value *Param, bool HasCvtaParam,
bool IsGridConstant) {
Instruction *I = dyn_cast<Instruction>(OldUse->getUser());
assert(I && "OldUse must be in an instruction");
struct IP {
Use *OldUse;
Instruction *OldInstruction;
Value *NewParam;
};
SmallVector<IP> ItemsToConvert = {{OldUse, I, Param}};
SmallVector<Instruction *> InstructionsToDelete;
auto CloneInstInParamAS = [HasCvtaParam,
IsGridConstant](const IP &I) -> Value * {
if (auto *LI = dyn_cast<LoadInst>(I.OldInstruction)) {
LI->setOperand(0, I.NewParam);
return LI;
}
if (auto *GEP = dyn_cast<GetElementPtrInst>(I.OldInstruction)) {
SmallVector<Value *, 4> Indices(GEP->indices());
auto *NewGEP = GetElementPtrInst::Create(
GEP->getSourceElementType(), I.NewParam, Indices, GEP->getName(),
GEP->getIterator());
NewGEP->setIsInBounds(GEP->isInBounds());
return NewGEP;
}
if (auto *BC = dyn_cast<BitCastInst>(I.OldInstruction)) {
auto *NewBCType = PointerType::get(BC->getContext(), ADDRESS_SPACE_PARAM);
return BitCastInst::Create(BC->getOpcode(), I.NewParam, NewBCType,
BC->getName(), BC->getIterator());
}
if (auto *ASC = dyn_cast<AddrSpaceCastInst>(I.OldInstruction)) {
assert(ASC->getDestAddressSpace() == ADDRESS_SPACE_PARAM);
(void)ASC;
// Just pass through the argument, the old ASC is no longer needed.
return I.NewParam;
}
if (auto *MI = dyn_cast<MemTransferInst>(I.OldInstruction)) {
if (MI->getRawSource() == I.OldUse->get()) {
// convert to memcpy/memmove from param space.
IRBuilder<> Builder(I.OldInstruction);
Intrinsic::ID ID = MI->getIntrinsicID();
CallInst *B = Builder.CreateMemTransferInst(
ID, MI->getRawDest(), MI->getDestAlign(), I.NewParam,
MI->getSourceAlign(), MI->getLength(), MI->isVolatile());
for (unsigned I : {0, 1})
if (uint64_t Bytes = MI->getParamDereferenceableBytes(I))
B->addDereferenceableParamAttr(I, Bytes);
return B;
}
// We may be able to handle other cases if the argument is
// __grid_constant__
}
if (HasCvtaParam) {
auto GetParamAddrCastToGeneric =
[](Value *Addr, Instruction *OriginalUser) -> Value * {
PointerType *ReturnTy =
PointerType::get(OriginalUser->getContext(), ADDRESS_SPACE_GENERIC);
Function *CvtToGen = Intrinsic::getOrInsertDeclaration(
OriginalUser->getModule(), Intrinsic::nvvm_ptr_param_to_gen,
{ReturnTy, PointerType::get(OriginalUser->getContext(),
ADDRESS_SPACE_PARAM)});
// Cast param address to generic address space
Value *CvtToGenCall =
CallInst::Create(CvtToGen, Addr, Addr->getName() + ".gen",
OriginalUser->getIterator());
return CvtToGenCall;
};
auto *ParamInGenericAS =
GetParamAddrCastToGeneric(I.NewParam, I.OldInstruction);
// phi/select could use generic arg pointers w/o __grid_constant__
if (auto *PHI = dyn_cast<PHINode>(I.OldInstruction)) {
for (auto [Idx, V] : enumerate(PHI->incoming_values())) {
if (V.get() == I.OldUse->get())
PHI->setIncomingValue(Idx, ParamInGenericAS);
}
}
if (auto *SI = dyn_cast<SelectInst>(I.OldInstruction)) {
if (SI->getTrueValue() == I.OldUse->get())
SI->setTrueValue(ParamInGenericAS);
if (SI->getFalseValue() == I.OldUse->get())
SI->setFalseValue(ParamInGenericAS);
}
// Escapes or writes can only use generic param pointers if
// __grid_constant__ is in effect.
if (IsGridConstant) {
if (auto *CI = dyn_cast<CallInst>(I.OldInstruction)) {
I.OldUse->set(ParamInGenericAS);
return CI;
}
if (auto *SI = dyn_cast<StoreInst>(I.OldInstruction)) {
// byval address is being stored, cast it to generic
if (SI->getValueOperand() == I.OldUse->get())
SI->setOperand(0, ParamInGenericAS);
return SI;
}
if (auto *PI = dyn_cast<PtrToIntInst>(I.OldInstruction)) {
if (PI->getPointerOperand() == I.OldUse->get())
PI->setOperand(0, ParamInGenericAS);
return PI;
}
// TODO: iIf we allow stores, we should allow memcpy/memset to
// parameter, too.
}
}
llvm_unreachable("Unsupported instruction");
};
while (!ItemsToConvert.empty()) {
IP I = ItemsToConvert.pop_back_val();
Value *NewInst = CloneInstInParamAS(I);
if (NewInst && NewInst != I.OldInstruction) {
// We've created a new instruction. Queue users of the old instruction to
// be converted and the instruction itself to be deleted. We can't delete
// the old instruction yet, because it's still in use by a load somewhere.
for (Use &U : I.OldInstruction->uses())
ItemsToConvert.push_back({&U, cast<Instruction>(U.getUser()), NewInst});
InstructionsToDelete.push_back(I.OldInstruction);
}
}
// Now we know that all argument loads are using addresses in parameter space
// and we can finally remove the old instructions in generic AS. Instructions
// scheduled for removal should be processed in reverse order so the ones
// closest to the load are deleted first. Otherwise they may still be in use.
// E.g if we have Value = Load(BitCast(GEP(arg))), InstructionsToDelete will
// have {GEP,BitCast}. GEP can't be deleted first, because it's still used by
// the BitCast.
for (Instruction *I : llvm::reverse(InstructionsToDelete))
I->eraseFromParent();
}
// Adjust alignment of arguments passed byval in .param address space. We can
// increase alignment of such arguments in a way that ensures that we can
// effectively vectorize their loads. We should also traverse all loads from
// byval pointer and adjust their alignment, if those were using known offset.
// Such alignment changes must be conformed with parameter store and load in
// NVPTXTargetLowering::LowerCall.
static void adjustByValArgAlignment(Argument *Arg, Value *ArgInParamAS,
const NVPTXTargetLowering *TLI) {
Function *Func = Arg->getParent();
Type *StructType = Arg->getParamByValType();
const DataLayout &DL = Func->getDataLayout();
uint64_t NewArgAlign =
TLI->getFunctionParamOptimizedAlign(Func, StructType, DL).value();
uint64_t CurArgAlign =
Arg->getAttribute(Attribute::Alignment).getValueAsInt();
if (CurArgAlign >= NewArgAlign)
return;
LLVM_DEBUG(dbgs() << "Try to use alignment " << NewArgAlign << " instead of "
<< CurArgAlign << " for " << *Arg << '\n');
auto NewAlignAttr =
Attribute::get(Func->getContext(), Attribute::Alignment, NewArgAlign);
Arg->removeAttr(Attribute::Alignment);
Arg->addAttr(NewAlignAttr);
struct Load {
LoadInst *Inst;
uint64_t Offset;
};
struct LoadContext {
Value *InitialVal;
uint64_t Offset;
};
SmallVector<Load> Loads;
std::queue<LoadContext> Worklist;
Worklist.push({ArgInParamAS, 0});
bool IsGridConstant = isParamGridConstant(*Arg);
while (!Worklist.empty()) {
LoadContext Ctx = Worklist.front();
Worklist.pop();
for (User *CurUser : Ctx.InitialVal->users()) {
if (auto *I = dyn_cast<LoadInst>(CurUser)) {
Loads.push_back({I, Ctx.Offset});
continue;
}
if (auto *I = dyn_cast<BitCastInst>(CurUser)) {
Worklist.push({I, Ctx.Offset});
continue;
}
if (auto *I = dyn_cast<GetElementPtrInst>(CurUser)) {
APInt OffsetAccumulated =
APInt::getZero(DL.getIndexSizeInBits(ADDRESS_SPACE_PARAM));
if (!I->accumulateConstantOffset(DL, OffsetAccumulated))
continue;
uint64_t OffsetLimit = -1;
uint64_t Offset = OffsetAccumulated.getLimitedValue(OffsetLimit);
assert(Offset != OffsetLimit && "Expect Offset less than UINT64_MAX");
Worklist.push({I, Ctx.Offset + Offset});
continue;
}
// supported for grid_constant
if (IsGridConstant &&
(isa<CallInst>(CurUser) || isa<StoreInst>(CurUser) ||
isa<PtrToIntInst>(CurUser)))
continue;
llvm_unreachable("All users must be one of: load, "
"bitcast, getelementptr, call, store, ptrtoint");
}
}
for (Load &CurLoad : Loads) {
Align NewLoadAlign(std::gcd(NewArgAlign, CurLoad.Offset));
Align CurLoadAlign(CurLoad.Inst->getAlign());
CurLoad.Inst->setAlignment(std::max(NewLoadAlign, CurLoadAlign));
}
}
namespace {
struct ArgUseChecker : PtrUseVisitor<ArgUseChecker> {
using Base = PtrUseVisitor<ArgUseChecker>;
bool IsGridConstant;
// Set of phi/select instructions using the Arg
SmallPtrSet<Instruction *, 4> Conditionals;
ArgUseChecker(const DataLayout &DL, bool IsGridConstant)
: PtrUseVisitor(DL), IsGridConstant(IsGridConstant) {}
PtrInfo visitArgPtr(Argument &A) {
assert(A.getType()->isPointerTy());
IntegerType *IntIdxTy = cast<IntegerType>(DL.getIndexType(A.getType()));
IsOffsetKnown = false;
Offset = APInt(IntIdxTy->getBitWidth(), 0);
PI.reset();
Conditionals.clear();
LLVM_DEBUG(dbgs() << "Checking Argument " << A << "\n");
// Enqueue the uses of this pointer.
enqueueUsers(A);
// Visit all the uses off the worklist until it is empty.
// Note that unlike PtrUseVisitor we intentionally do not track offsets.
// We're only interested in how we use the pointer.
while (!(Worklist.empty() || PI.isAborted())) {
UseToVisit ToVisit = Worklist.pop_back_val();
U = ToVisit.UseAndIsOffsetKnown.getPointer();
Instruction *I = cast<Instruction>(U->getUser());
if (isa<PHINode>(I) || isa<SelectInst>(I))
Conditionals.insert(I);
LLVM_DEBUG(dbgs() << "Processing " << *I << "\n");
Base::visit(I);
}
if (PI.isEscaped())
LLVM_DEBUG(dbgs() << "Argument pointer escaped: " << *PI.getEscapingInst()
<< "\n");
else if (PI.isAborted())
LLVM_DEBUG(dbgs() << "Pointer use needs a copy: " << *PI.getAbortingInst()
<< "\n");
LLVM_DEBUG(dbgs() << "Traversed " << Conditionals.size()
<< " conditionals\n");
return PI;
}
void visitStoreInst(StoreInst &SI) {
// Storing the pointer escapes it.
if (U->get() == SI.getValueOperand())
return PI.setEscapedAndAborted(&SI);
// Writes to the pointer are UB w/ __grid_constant__, but do not force a
// copy.
if (!IsGridConstant)
return PI.setAborted(&SI);
}
void visitAddrSpaceCastInst(AddrSpaceCastInst &ASC) {
// ASC to param space are no-ops and do not need a copy
if (ASC.getDestAddressSpace() != ADDRESS_SPACE_PARAM)
return PI.setEscapedAndAborted(&ASC);
Base::visitAddrSpaceCastInst(ASC);
}
void visitPtrToIntInst(PtrToIntInst &I) {
if (IsGridConstant)
return;
Base::visitPtrToIntInst(I);
}
void visitPHINodeOrSelectInst(Instruction &I) {
assert(isa<PHINode>(I) || isa<SelectInst>(I));
}
// PHI and select just pass through the pointers.
void visitPHINode(PHINode &PN) { enqueueUsers(PN); }
void visitSelectInst(SelectInst &SI) { enqueueUsers(SI); }
void visitMemTransferInst(MemTransferInst &II) {
if (*U == II.getRawDest() && !IsGridConstant)
PI.setAborted(&II);
// memcpy/memmove are OK when the pointer is source. We can convert them to
// AS-specific memcpy.
}
void visitMemSetInst(MemSetInst &II) {
if (!IsGridConstant)
PI.setAborted(&II);
}
}; // struct ArgUseChecker
} // namespace
void NVPTXLowerArgs::handleByValParam(const NVPTXTargetMachine &TM,
Argument *Arg) {
Function *Func = Arg->getParent();
bool HasCvtaParam = TM.getSubtargetImpl(*Func)->hasCvtaParam();
bool IsGridConstant = HasCvtaParam && isParamGridConstant(*Arg);
const DataLayout &DL = Func->getDataLayout();
BasicBlock::iterator FirstInst = Func->getEntryBlock().begin();
Type *StructType = Arg->getParamByValType();
assert(StructType && "Missing byval type");
ArgUseChecker AUC(DL, IsGridConstant);
ArgUseChecker::PtrInfo PI = AUC.visitArgPtr(*Arg);
bool ArgUseIsReadOnly = !(PI.isEscaped() || PI.isAborted());
// Easy case, accessing parameter directly is fine.
if (ArgUseIsReadOnly && AUC.Conditionals.empty()) {
// Convert all loads and intermediate operations to use parameter AS and
// skip creation of a local copy of the argument.
SmallVector<Use *, 16> UsesToUpdate;
for (Use &U : Arg->uses())
UsesToUpdate.push_back(&U);
Value *ArgInParamAS = new AddrSpaceCastInst(
Arg, PointerType::get(StructType, ADDRESS_SPACE_PARAM), Arg->getName(),
FirstInst);
for (Use *U : UsesToUpdate)
convertToParamAS(U, ArgInParamAS, HasCvtaParam, IsGridConstant);
LLVM_DEBUG(dbgs() << "No need to copy or cast " << *Arg << "\n");
const auto *TLI =
cast<NVPTXTargetLowering>(TM.getSubtargetImpl()->getTargetLowering());
adjustByValArgAlignment(Arg, ArgInParamAS, TLI);
return;
}
// We can't access byval arg directly and need a pointer. on sm_70+ we have
// ability to take a pointer to the argument without making a local copy.
// However, we're still not allowed to write to it. If the user specified
// `__grid_constant__` for the argument, we'll consider escaped pointer as
// read-only.
unsigned AS = DL.getAllocaAddrSpace();
if (HasCvtaParam && (ArgUseIsReadOnly || IsGridConstant)) {
LLVM_DEBUG(dbgs() << "Using non-copy pointer to " << *Arg << "\n");
// Replace all argument pointer uses (which might include a device function
// call) with a cast to the generic address space using cvta.param
// instruction, which avoids a local copy.
IRBuilder<> IRB(&Func->getEntryBlock().front());
// Cast argument to param address space
auto *CastToParam = cast<AddrSpaceCastInst>(IRB.CreateAddrSpaceCast(
Arg, IRB.getPtrTy(ADDRESS_SPACE_PARAM), Arg->getName() + ".param"));
// Cast param address to generic address space. We do not use an
// addrspacecast to generic here, because, LLVM considers `Arg` to be in the
// generic address space, and a `generic -> param` cast followed by a `param
// -> generic` cast will be folded away. The `param -> generic` intrinsic
// will be correctly lowered to `cvta.param`.
Value *CvtToGenCall = IRB.CreateIntrinsic(
IRB.getPtrTy(ADDRESS_SPACE_GENERIC), Intrinsic::nvvm_ptr_param_to_gen,
CastToParam, nullptr, CastToParam->getName() + ".gen");
Arg->replaceAllUsesWith(CvtToGenCall);
// Do not replace Arg in the cast to param space
CastToParam->setOperand(0, Arg);
} else {
LLVM_DEBUG(dbgs() << "Creating a local copy of " << *Arg << "\n");
// Otherwise we have to create a temporary copy.
AllocaInst *AllocA =
new AllocaInst(StructType, AS, Arg->getName(), FirstInst);
// Set the alignment to alignment of the byval parameter. This is because,
// later load/stores assume that alignment, and we are going to replace
// the use of the byval parameter with this alloca instruction.
AllocA->setAlignment(Func->getParamAlign(Arg->getArgNo())
.value_or(DL.getPrefTypeAlign(StructType)));
Arg->replaceAllUsesWith(AllocA);
Value *ArgInParam = new AddrSpaceCastInst(
Arg, PointerType::get(Arg->getContext(), ADDRESS_SPACE_PARAM),
Arg->getName(), FirstInst);
// Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
// addrspacecast preserves alignment. Since params are constant, this load
// is definitely not volatile.
LoadInst *LI =
new LoadInst(StructType, ArgInParam, Arg->getName(),
/*isVolatile=*/false, AllocA->getAlign(), FirstInst);
new StoreInst(LI, AllocA, FirstInst);
}
}
void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {
if (Ptr->getType()->getPointerAddressSpace() != ADDRESS_SPACE_GENERIC)
return;
// Deciding where to emit the addrspacecast pair.
BasicBlock::iterator InsertPt;
if (Argument *Arg = dyn_cast<Argument>(Ptr)) {
// Insert at the functon entry if Ptr is an argument.
InsertPt = Arg->getParent()->getEntryBlock().begin();
} else {
// Insert right after Ptr if Ptr is an instruction.
InsertPt = ++cast<Instruction>(Ptr)->getIterator();
assert(InsertPt != InsertPt->getParent()->end() &&
"We don't call this function with Ptr being a terminator.");
}
Instruction *PtrInGlobal = new AddrSpaceCastInst(
Ptr, PointerType::get(Ptr->getContext(), ADDRESS_SPACE_GLOBAL),
Ptr->getName(), InsertPt);
Value *PtrInGeneric = new AddrSpaceCastInst(PtrInGlobal, Ptr->getType(),
Ptr->getName(), InsertPt);
// Replace with PtrInGeneric all uses of Ptr except PtrInGlobal.
Ptr->replaceAllUsesWith(PtrInGeneric);
PtrInGlobal->setOperand(0, Ptr);
}
// =============================================================================
// Main function for this pass.
// =============================================================================
bool NVPTXLowerArgs::runOnKernelFunction(const NVPTXTargetMachine &TM,
Function &F) {
// Copying of byval aggregates + SROA may result in pointers being loaded as
// integers, followed by intotoptr. We may want to mark those as global, too,
// but only if the loaded integer is used exclusively for conversion to a
// pointer with inttoptr.
auto HandleIntToPtr = [this](Value &V) {
if (llvm::all_of(V.users(), [](User *U) { return isa<IntToPtrInst>(U); })) {
SmallVector<User *, 16> UsersToUpdate(V.users());
for (User *U : UsersToUpdate)
markPointerAsGlobal(U);
}
};
if (TM.getDrvInterface() == NVPTX::CUDA) {
// Mark pointers in byval structs as global.
for (auto &B : F) {
for (auto &I : B) {
if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
if (LI->getType()->isPointerTy() || LI->getType()->isIntegerTy()) {
Value *UO = getUnderlyingObject(LI->getPointerOperand());
if (Argument *Arg = dyn_cast<Argument>(UO)) {
if (Arg->hasByValAttr()) {
// LI is a load from a pointer within a byval kernel parameter.
if (LI->getType()->isPointerTy())
markPointerAsGlobal(LI);
else
HandleIntToPtr(*LI);
}
}
}
}
}
}
}
LLVM_DEBUG(dbgs() << "Lowering kernel args of " << F.getName() << "\n");
for (Argument &Arg : F.args()) {
if (Arg.getType()->isPointerTy()) {
if (Arg.hasByValAttr())
handleByValParam(TM, &Arg);
else if (TM.getDrvInterface() == NVPTX::CUDA)
markPointerAsGlobal(&Arg);
} else if (Arg.getType()->isIntegerTy() &&
TM.getDrvInterface() == NVPTX::CUDA) {
HandleIntToPtr(Arg);
}
}
return true;
}
// Device functions only need to copy byval args into local memory.
bool NVPTXLowerArgs::runOnDeviceFunction(const NVPTXTargetMachine &TM,
Function &F) {
LLVM_DEBUG(dbgs() << "Lowering function args of " << F.getName() << "\n");
for (Argument &Arg : F.args())
if (Arg.getType()->isPointerTy() && Arg.hasByValAttr())
handleByValParam(TM, &Arg);
return true;
}
bool NVPTXLowerArgs::runOnFunction(Function &F) {
auto &TM = getAnalysis<TargetPassConfig>().getTM<NVPTXTargetMachine>();
return isKernelFunction(F) ? runOnKernelFunction(TM, F)
: runOnDeviceFunction(TM, F);
}
FunctionPass *llvm::createNVPTXLowerArgsPass() { return new NVPTXLowerArgs(); }
Reference in New Issue View Git Blame Copy Permalink