1 //===-- NVPTXLowerKernelArgs.cpp - Lower kernel arguments -----------------===// 2 // 3 // The LLVM Compiler Infrastructure 4 // 5 // This file is distributed under the University of Illinois Open Source 6 // License. See LICENSE.TXT for details. 7 // 8 //===----------------------------------------------------------------------===// 9 // 10 // Pointer arguments to kernel functions need to be lowered specially. 11 // 12 // 1. Copy byval struct args to local memory. This is a preparation for handling 13 // cases like 14 // 15 // kernel void foo(struct A arg, ...) 16 // { 17 // struct A *p = &arg; 18 // ... 19 // ... = p->filed1 ... (this is no generic address for .param) 20 // p->filed2 = ... (this is no write access to .param) 21 // } 22 // 23 // 2. Convert non-byval pointer arguments of CUDA kernels to pointers in the 24 // global address space. This allows later optimizations to emit 25 // ld.global.*/st.global.* for accessing these pointer arguments. For 26 // example, 27 // 28 // define void @foo(float* %input) { 29 // %v = load float, float* %input, align 4 30 // ... 31 // } 32 // 33 // becomes 34 // 35 // define void @foo(float* %input) { 36 // %input2 = addrspacecast float* %input to float addrspace(1)* 37 // %input3 = addrspacecast float addrspace(1)* %input2 to float* 38 // %v = load float, float* %input3, align 4 39 // ... 40 // } 41 // 42 // Later, NVPTXFavorNonGenericAddrSpaces will optimize it to 43 // 44 // define void @foo(float* %input) { 45 // %input2 = addrspacecast float* %input to float addrspace(1)* 46 // %v = load float, float addrspace(1)* %input2, align 4 47 // ... 48 // } 49 // 50 // 3. Convert pointers in a byval kernel parameter to pointers in the global 51 // address space. As #2, it allows NVPTX to emit more ld/st.global. E.g., 52 // 53 // struct S { 54 // int *x; 55 // int *y; 56 // }; 57 // __global__ void foo(S s) { 58 // int *b = s.y; 59 // // use b 60 // } 61 // 62 // "b" points to the global address space. In the IR level, 63 // 64 // define void @foo({i32*, i32*}* byval %input) { 65 // %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1 66 // %b = load i32*, i32** %b_ptr 67 // ; use %b 68 // } 69 // 70 // becomes 71 // 72 // define void @foo({i32*, i32*}* byval %input) { 73 // %b_ptr = getelementptr {i32*, i32*}, {i32*, i32*}* %input, i64 0, i32 1 74 // %b = load i32*, i32** %b_ptr 75 // %b_global = addrspacecast i32* %b to i32 addrspace(1)* 76 // %b_generic = addrspacecast i32 addrspace(1)* %b_global to i32* 77 // ; use %b_generic 78 // } 79 // 80 // TODO: merge this pass with NVPTXFavorNonGenericAddrSpace so that other passes 81 // don't cancel the addrspacecast pair this pass emits. 82 //===----------------------------------------------------------------------===// 83 84 #include "NVPTX.h" 85 #include "NVPTXUtilities.h" 86 #include "NVPTXTargetMachine.h" 87 #include "llvm/Analysis/ValueTracking.h" 88 #include "llvm/IR/Function.h" 89 #include "llvm/IR/Instructions.h" 90 #include "llvm/IR/Module.h" 91 #include "llvm/IR/Type.h" 92 #include "llvm/Pass.h" 93 94 using namespace llvm; 95 96 namespace llvm { 97 void initializeNVPTXLowerKernelArgsPass(PassRegistry &); 98 } 99 100 namespace { 101 class NVPTXLowerKernelArgs : public FunctionPass { 102 bool runOnFunction(Function &F) override; 103 104 // handle byval parameters 105 void handleByValParam(Argument *Arg); 106 // Knowing Ptr must point to the global address space, this function 107 // addrspacecasts Ptr to global and then back to generic. This allows 108 // NVPTXFavorNonGenericAddrSpace to fold the global-to-generic cast into 109 // loads/stores that appear later. 110 void markPointerAsGlobal(Value *Ptr); 111 112 public: 113 static char ID; // Pass identification, replacement for typeid 114 NVPTXLowerKernelArgs(const NVPTXTargetMachine *TM = nullptr) 115 : FunctionPass(ID), TM(TM) {} 116 const char *getPassName() const override { 117 return "Lower pointer arguments of CUDA kernels"; 118 } 119 120 private: 121 const NVPTXTargetMachine *TM; 122 }; 123 } // namespace 124 125 char NVPTXLowerKernelArgs::ID = 1; 126 127 INITIALIZE_PASS(NVPTXLowerKernelArgs, "nvptx-lower-kernel-args", 128 "Lower kernel arguments (NVPTX)", false, false) 129 130 // ============================================================================= 131 // If the function had a byval struct ptr arg, say foo(%struct.x *byval %d), 132 // then add the following instructions to the first basic block: 133 // 134 // %temp = alloca %struct.x, align 8 135 // %tempd = addrspacecast %struct.x* %d to %struct.x addrspace(101)* 136 // %tv = load %struct.x addrspace(101)* %tempd 137 // store %struct.x %tv, %struct.x* %temp, align 8 138 // 139 // The above code allocates some space in the stack and copies the incoming 140 // struct from param space to local space. 141 // Then replace all occurrences of %d by %temp. 142 // ============================================================================= 143 void NVPTXLowerKernelArgs::handleByValParam(Argument *Arg) { 144 Function *Func = Arg->getParent(); 145 Instruction *FirstInst = &(Func->getEntryBlock().front()); 146 PointerType *PType = dyn_cast<PointerType>(Arg->getType()); 147 148 assert(PType && "Expecting pointer type in handleByValParam"); 149 150 Type *StructType = PType->getElementType(); 151 AllocaInst *AllocA = new AllocaInst(StructType, Arg->getName(), FirstInst); 152 // Set the alignment to alignment of the byval parameter. This is because, 153 // later load/stores assume that alignment, and we are going to replace 154 // the use of the byval parameter with this alloca instruction. 155 AllocA->setAlignment(Func->getParamAlignment(Arg->getArgNo() + 1)); 156 Arg->replaceAllUsesWith(AllocA); 157 158 Value *ArgInParam = new AddrSpaceCastInst( 159 Arg, PointerType::get(StructType, ADDRESS_SPACE_PARAM), Arg->getName(), 160 FirstInst); 161 LoadInst *LI = new LoadInst(ArgInParam, Arg->getName(), FirstInst); 162 new StoreInst(LI, AllocA, FirstInst); 163 } 164 165 void NVPTXLowerKernelArgs::markPointerAsGlobal(Value *Ptr) { 166 if (Ptr->getType()->getPointerAddressSpace() == ADDRESS_SPACE_GLOBAL) 167 return; 168 169 // Deciding where to emit the addrspacecast pair. 170 BasicBlock::iterator InsertPt; 171 if (Argument *Arg = dyn_cast<Argument>(Ptr)) { 172 // Insert at the functon entry if Ptr is an argument. 173 InsertPt = Arg->getParent()->getEntryBlock().begin(); 174 } else { 175 // Insert right after Ptr if Ptr is an instruction. 176 InsertPt = ++cast<Instruction>(Ptr)->getIterator(); 177 assert(InsertPt != InsertPt->getParent()->end() && 178 "We don't call this function with Ptr being a terminator."); 179 } 180 181 Instruction *PtrInGlobal = new AddrSpaceCastInst( 182 Ptr, PointerType::get(Ptr->getType()->getPointerElementType(), 183 ADDRESS_SPACE_GLOBAL), 184 Ptr->getName(), &*InsertPt); 185 Value *PtrInGeneric = new AddrSpaceCastInst(PtrInGlobal, Ptr->getType(), 186 Ptr->getName(), &*InsertPt); 187 // Replace with PtrInGeneric all uses of Ptr except PtrInGlobal. 188 Ptr->replaceAllUsesWith(PtrInGeneric); 189 PtrInGlobal->setOperand(0, Ptr); 190 } 191 192 // ============================================================================= 193 // Main function for this pass. 194 // ============================================================================= 195 bool NVPTXLowerKernelArgs::runOnFunction(Function &F) { 196 // Skip non-kernels. See the comments at the top of this file. 197 if (!isKernelFunction(F)) 198 return false; 199 200 if (TM && TM->getDrvInterface() == NVPTX::CUDA) { 201 // Mark pointers in byval structs as global. 202 for (auto &B : F) { 203 for (auto &I : B) { 204 if (LoadInst *LI = dyn_cast<LoadInst>(&I)) { 205 if (LI->getType()->isPointerTy()) { 206 Value *UO = GetUnderlyingObject(LI->getPointerOperand(), 207 F.getParent()->getDataLayout()); 208 if (Argument *Arg = dyn_cast<Argument>(UO)) { 209 if (Arg->hasByValAttr()) { 210 // LI is a load from a pointer within a byval kernel parameter. 211 markPointerAsGlobal(LI); 212 } 213 } 214 } 215 } 216 } 217 } 218 } 219 220 for (Argument &Arg : F.args()) { 221 if (Arg.getType()->isPointerTy()) { 222 if (Arg.hasByValAttr()) 223 handleByValParam(&Arg); 224 else if (TM && TM->getDrvInterface() == NVPTX::CUDA) 225 markPointerAsGlobal(&Arg); 226 } 227 } 228 return true; 229 } 230 231 FunctionPass * 232 llvm::createNVPTXLowerKernelArgsPass(const NVPTXTargetMachine *TM) { 233 return new NVPTXLowerKernelArgs(TM); 234 } 235