LLVM  mainline
NVPTXLowerKernelArgs.cpp
Go to the documentation of this file.
00001 //===-- NVPTXLowerKernelArgs.cpp - Lower kernel arguments -----------------===//
00002 //
00003 //                     The LLVM Compiler Infrastructure
00004 //
00005 // This file is distributed under the University of Illinois Open Source
00006 // License. See LICENSE.TXT for details.
00007 //
00008 //===----------------------------------------------------------------------===//
00009 //
00010 // Pointer arguments to kernel functions need to be lowered specially.
00011 //
00012 // 1. Copy byval struct args to local memory. This is a preparation for handling
00013 //    cases like
00014 //
00015 //    kernel void foo(struct A arg, ...)
00016 //    {
00017 //      struct A *p = &arg;
00018 //      ...
00019 //      ... = p->filed1 ...  (this is no generic address for .param)
00020 //      p->filed2 = ...      (this is no write access to .param)
00021 //    }
00022 //
00023 // 2. Convert non-byval pointer arguments of CUDA kernels to pointers in the
00024 //    global address space. This allows later optimizations to emit
00025 //    ld.global.*/st.global.* for accessing these pointer arguments. For
00026 //    example,
00027 //
00028 //    define void @foo(float* %input) {
00029 //      %v = load float, float* %input, align 4
00030 //      ...
00031 //    }
00032 //
00033 //    becomes
00034 //
00035 //    define void @foo(float* %input) {
00036 //      %input2 = addrspacecast float* %input to float addrspace(1)*
00037 //      %input3 = addrspacecast float addrspace(1)* %input2 to float*
00038 //      %v = load float, float* %input3, align 4
00039 //      ...
00040 //    }
00041 //
00042 //    Later, NVPTXFavorNonGenericAddrSpaces will optimize it to
00043 //
00044 //    define void @foo(float* %input) {
00045 //      %input2 = addrspacecast float* %input to float addrspace(1)*
00046 //      %v = load float, float addrspace(1)* %input2, align 4
00047 //      ...
00048 //    }
00049 //
00050 // 3. Convert pointers in a byval kernel parameter to pointers in the global
00051 //    address space. As #2, it allows NVPTX to emit more ld/st.global. E.g.,
00052 //
00053 //    struct S {
00054 //      int *x;
00055 //      int *y;
00056 //    };
00057 //    __global__ void foo(S s) {
00058 //      int *b = s.y;
00059 //      // use b
00060 //    }
00061 //
00062 //    "b" points to the global address space. In the IR level,
00063 //
00064 //    define void @foo({i32*, i32*}* byval %input) {
00065 //      %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1
00066 //      %b = load i32*, i32** %b_ptr
00067 //      ; use %b
00068 //    }
00069 //
00070 //    becomes
00071 //
00072 //    define void @foo({i32*, i32*}* byval %input) {
00073 //      %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1
00074 //      %b = load i32*, i32** %b_ptr
00075 //      %b_global = addrspacecast i32* %b to i32 addrspace(1)*
00076 //      %b_generic = addrspacecast i32 addrspace(1)* %b_global to i32*
00077 //      ; use %b_generic
00078 //    }
00079 //
00080 // TODO: merge this pass with NVPTXFavorNonGenericAddrSpace so that other passes
00081 // don't cancel the addrspacecast pair this pass emits.
00082 //===----------------------------------------------------------------------===//
00083 
00084 #include "NVPTX.h"
00085 #include "NVPTXUtilities.h"
00086 #include "NVPTXTargetMachine.h"
00087 #include "llvm/Analysis/ValueTracking.h"
00088 #include "llvm/IR/Function.h"
00089 #include "llvm/IR/Instructions.h"
00090 #include "llvm/IR/Module.h"
00091 #include "llvm/IR/Type.h"
00092 #include "llvm/Pass.h"
00093 
00094 using namespace llvm;
00095 
00096 namespace llvm {
00097 void initializeNVPTXLowerKernelArgsPass(PassRegistry &);
00098 }
00099 
00100 namespace {
00101 class NVPTXLowerKernelArgs : public FunctionPass {
00102   bool runOnFunction(Function &F) override;
00103 
00104   // handle byval parameters
00105   void handleByValParam(Argument *Arg);
00106   // Knowing Ptr must point to the global address space, this function
00107   // addrspacecasts Ptr to global and then back to generic. This allows
00108   // NVPTXFavorNonGenericAddrSpace to fold the global-to-generic cast into
00109   // loads/stores that appear later.
00110   void markPointerAsGlobal(Value *Ptr);
00111 
00112 public:
00113   static char ID; // Pass identification, replacement for typeid
00114   NVPTXLowerKernelArgs(const NVPTXTargetMachine *TM = nullptr)
00115       : FunctionPass(ID), TM(TM) {}
00116   const char *getPassName() const override {
00117     return "Lower pointer arguments of CUDA kernels";
00118   }
00119 
00120 private:
00121   const NVPTXTargetMachine *TM;
00122 };
00123 } // namespace
00124 
00125 char NVPTXLowerKernelArgs::ID = 1;
00126 
00127 INITIALIZE_PASS(NVPTXLowerKernelArgs, "nvptx-lower-kernel-args",
00128                 "Lower kernel arguments (NVPTX)", false, false)
00129 
00130 // =============================================================================
00131 // If the function had a byval struct ptr arg, say foo(%struct.x* byval %d),
00132 // then add the following instructions to the first basic block:
00133 //
00134 // %temp = alloca %struct.x, align 8
00135 // %tempd = addrspacecast %struct.x* %d to %struct.x addrspace(101)*
00136 // %tv = load %struct.x addrspace(101)* %tempd
00137 // store %struct.x %tv, %struct.x* %temp, align 8
00138 //
00139 // The above code allocates some space in the stack and copies the incoming
00140 // struct from param space to local space.
00141 // Then replace all occurrences of %d by %temp.
00142 // =============================================================================
00143 void NVPTXLowerKernelArgs::handleByValParam(Argument *Arg) {
00144   Function *Func = Arg->getParent();
00145   Instruction *FirstInst = &(Func->getEntryBlock().front());
00146   PointerType *PType = dyn_cast<PointerType>(Arg->getType());
00147 
00148   assert(PType && "Expecting pointer type in handleByValParam");
00149 
00150   Type *StructType = PType->getElementType();
00151   AllocaInst *AllocA = new AllocaInst(StructType, Arg->getName(), FirstInst);
00152   // Set the alignment to alignment of the byval parameter. This is because,
00153   // later load/stores assume that alignment, and we are going to replace
00154   // the use of the byval parameter with this alloca instruction.
00155   AllocA->setAlignment(Func->getParamAlignment(Arg->getArgNo() + 1));
00156   Arg->replaceAllUsesWith(AllocA);
00157 
00158   Value *ArgInParam = new AddrSpaceCastInst(
00159       Arg, PointerType::get(StructType, ADDRESS_SPACE_PARAM), Arg->getName(),
00160       FirstInst);
00161   LoadInst *LI = new LoadInst(ArgInParam, Arg->getName(), FirstInst);
00162   new StoreInst(LI, AllocA, FirstInst);
00163 }
00164 
00165 void NVPTXLowerKernelArgs::markPointerAsGlobal(Value *Ptr) {
00166   if (Ptr->getType()->getPointerAddressSpace() == ADDRESS_SPACE_GLOBAL)
00167     return;
00168 
00169   // Deciding where to emit the addrspacecast pair.
00170   BasicBlock::iterator InsertPt;
00171   if (Argument *Arg = dyn_cast<Argument>(Ptr)) {
00172     // Insert at the functon entry if Ptr is an argument.
00173     InsertPt = Arg->getParent()->getEntryBlock().begin();
00174   } else {
00175     // Insert right after Ptr if Ptr is an instruction.
00176     InsertPt = ++cast<Instruction>(Ptr)->getIterator();
00177     assert(InsertPt != InsertPt->getParent()->end() &&
00178            "We don't call this function with Ptr being a terminator.");
00179   }
00180 
00181   Instruction *PtrInGlobal = new AddrSpaceCastInst(
00182       Ptr, PointerType::get(Ptr->getType()->getPointerElementType(),
00183                             ADDRESS_SPACE_GLOBAL),
00184       Ptr->getName(), &*InsertPt);
00185   Value *PtrInGeneric = new AddrSpaceCastInst(PtrInGlobal, Ptr->getType(),
00186                                               Ptr->getName(), &*InsertPt);
00187   // Replace with PtrInGeneric all uses of Ptr except PtrInGlobal.
00188   Ptr->replaceAllUsesWith(PtrInGeneric);
00189   PtrInGlobal->setOperand(0, Ptr);
00190 }
00191 
00192 // =============================================================================
00193 // Main function for this pass.
00194 // =============================================================================
00195 bool NVPTXLowerKernelArgs::runOnFunction(Function &F) {
00196   // Skip non-kernels. See the comments at the top of this file.
00197   if (!isKernelFunction(F))
00198     return false;
00199 
00200   if (TM && TM->getDrvInterface() == NVPTX::CUDA) {
00201     // Mark pointers in byval structs as global.
00202     for (auto &B : F) {
00203       for (auto &I : B) {
00204         if (LoadInst *LI = dyn_cast<LoadInst>(&I)) {
00205           if (LI->getType()->isPointerTy()) {
00206             Value *UO = GetUnderlyingObject(LI->getPointerOperand(),
00207                                             F.getParent()->getDataLayout());
00208             if (Argument *Arg = dyn_cast<Argument>(UO)) {
00209               if (Arg->hasByValAttr()) {
00210                 // LI is a load from a pointer within a byval kernel parameter.
00211                 markPointerAsGlobal(LI);
00212               }
00213             }
00214           }
00215         }
00216       }
00217     }
00218   }
00219 
00220   for (Argument &Arg : F.args()) {
00221     if (Arg.getType()->isPointerTy()) {
00222       if (Arg.hasByValAttr())
00223         handleByValParam(&Arg);
00224       else if (TM && TM->getDrvInterface() == NVPTX::CUDA)
00225         markPointerAsGlobal(&Arg);
00226     }
00227   }
00228   return true;
00229 }
00230 
00231 FunctionPass *
00232 llvm::createNVPTXLowerKernelArgsPass(const NVPTXTargetMachine *TM) {
00233   return new NVPTXLowerKernelArgs(TM);
00234 }