| //===-- NVPTXLowerKernelArgs.cpp - Lower kernel arguments -----------------===// |
| // |
| // The LLVM Compiler Infrastructure |
| // |
| // This file is distributed under the University of Illinois Open Source |
| // License. See LICENSE.TXT for details. |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Pointer arguments to kernel functions need to be lowered specially. |
| // |
| // 1. Copy byval struct args to local memory. This is a preparation for handling |
| // cases like |
| // |
| // kernel void foo(struct A arg, ...) |
| // { |
| // struct A *p = &arg; |
| // ... |
| // ... = p->filed1 ... (this is no generic address for .param) |
| // p->filed2 = ... (this is no write access to .param) |
| // } |
| // |
| // 2. 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, NVPTXFavorNonGenericAddrSpaces 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 |
| // ... |
| // } |
| // |
| // 3. Convert pointers in a byval kernel parameter to pointers in the global |
| // address space. As #2, it allows NVPTX to emit more ld/st.global. E.g., |
| // |
| // 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({i32*, i32*}* byval %input) { |
| // %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1 |
| // %b = load i32*, i32** %b_ptr |
| // ; use %b |
| // } |
| // |
| // becomes |
| // |
| // define void @foo({i32*, i32*}* byval %input) { |
| // %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1 |
| // %b = load i32*, i32** %b_ptr |
| // %b_global = addrspacecast i32* %b to i32 addrspace(1)* |
| // %b_generic = addrspacecast i32 addrspace(1)* %b_global to i32* |
| // ; use %b_generic |
| // } |
| // |
| // TODO: merge this pass with NVPTXFavorNonGenericAddrSpace so that other passes |
| // don't cancel the addrspacecast pair this pass emits. |
| //===----------------------------------------------------------------------===// |
| |
| #include "NVPTX.h" |
| #include "NVPTXUtilities.h" |
| #include "NVPTXTargetMachine.h" |
| #include "llvm/Analysis/ValueTracking.h" |
| #include "llvm/IR/Function.h" |
| #include "llvm/IR/Instructions.h" |
| #include "llvm/IR/Module.h" |
| #include "llvm/IR/Type.h" |
| #include "llvm/Pass.h" |
| |
| using namespace llvm; |
| |
| namespace llvm { |
| void initializeNVPTXLowerKernelArgsPass(PassRegistry &); |
| } |
| |
| namespace { |
| class NVPTXLowerKernelArgs : public FunctionPass { |
| bool runOnFunction(Function &F) override; |
| |
| // handle byval parameters |
| void handleByValParam(Argument *Arg); |
| // Knowing Ptr must point to the global address space, this function |
| // addrspacecasts Ptr to global and then back to generic. This allows |
| // NVPTXFavorNonGenericAddrSpace 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 |
| NVPTXLowerKernelArgs(const NVPTXTargetMachine *TM = nullptr) |
| : FunctionPass(ID), TM(TM) {} |
| const char *getPassName() const override { |
| return "Lower pointer arguments of CUDA kernels"; |
| } |
| |
| private: |
| const NVPTXTargetMachine *TM; |
| }; |
| } // namespace |
| |
| char NVPTXLowerKernelArgs::ID = 1; |
| |
| INITIALIZE_PASS(NVPTXLowerKernelArgs, "nvptx-lower-kernel-args", |
| "Lower kernel arguments (NVPTX)", false, false) |
| |
| // ============================================================================= |
| // If the function had a byval struct ptr arg, say foo(%struct.x *byval %d), |
| // 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. |
| // ============================================================================= |
| void NVPTXLowerKernelArgs::handleByValParam(Argument *Arg) { |
| Function *Func = Arg->getParent(); |
| Instruction *FirstInst = &(Func->getEntryBlock().front()); |
| PointerType *PType = dyn_cast<PointerType>(Arg->getType()); |
| |
| assert(PType && "Expecting pointer type in handleByValParam"); |
| |
| Type *StructType = PType->getElementType(); |
| AllocaInst *AllocA = new AllocaInst(StructType, 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->getParamAlignment(Arg->getArgNo() + 1)); |
| Arg->replaceAllUsesWith(AllocA); |
| |
| Value *ArgInParam = new AddrSpaceCastInst( |
| Arg, PointerType::get(StructType, ADDRESS_SPACE_PARAM), Arg->getName(), |
| FirstInst); |
| LoadInst *LI = new LoadInst(ArgInParam, Arg->getName(), FirstInst); |
| new StoreInst(LI, AllocA, FirstInst); |
| } |
| |
| void NVPTXLowerKernelArgs::markPointerAsGlobal(Value *Ptr) { |
| if (Ptr->getType()->getPointerAddressSpace() == ADDRESS_SPACE_GLOBAL) |
| 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->getType()->getPointerElementType(), |
| 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 NVPTXLowerKernelArgs::runOnFunction(Function &F) { |
| // Skip non-kernels. See the comments at the top of this file. |
| if (!isKernelFunction(F)) |
| return false; |
| |
| if (TM && 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()) { |
| Value *UO = GetUnderlyingObject(LI->getPointerOperand(), |
| F.getParent()->getDataLayout()); |
| if (Argument *Arg = dyn_cast<Argument>(UO)) { |
| if (Arg->hasByValAttr()) { |
| // LI is a load from a pointer within a byval kernel parameter. |
| markPointerAsGlobal(LI); |
| } |
| } |
| } |
| } |
| } |
| } |
| } |
| |
| for (Argument &Arg : F.args()) { |
| if (Arg.getType()->isPointerTy()) { |
| if (Arg.hasByValAttr()) |
| handleByValParam(&Arg); |
| else if (TM && TM->getDrvInterface() == NVPTX::CUDA) |
| markPointerAsGlobal(&Arg); |
| } |
| } |
| return true; |
| } |
| |
| FunctionPass * |
| llvm::createNVPTXLowerKernelArgsPass(const NVPTXTargetMachine *TM) { |
| return new NVPTXLowerKernelArgs(TM); |
| } |