1 //===-- AMDGPUOpenCLImageTypeLoweringPass.cpp -----------------------------===// 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 /// \file 11 /// This pass resolves calls to OpenCL image attribute, image resource ID and 12 /// sampler resource ID getter functions. 13 /// 14 /// Image attributes (size and format) are expected to be passed to the kernel 15 /// as kernel arguments immediately following the image argument itself, 16 /// therefore this pass adds image size and format arguments to the kernel 17 /// functions in the module. The kernel functions with image arguments are 18 /// re-created using the new signature. The new arguments are added to the 19 /// kernel metadata with kernel_arg_type set to "image_size" or "image_format". 20 /// Note: this pass may invalidate pointers to functions. 21 /// 22 /// Resource IDs of read-only images, write-only images and samplers are 23 /// defined to be their index among the kernel arguments of the same 24 /// type and access qualifier. 25 //===----------------------------------------------------------------------===// 26 27 #include "AMDGPU.h" 28 #include "llvm/ADT/DenseSet.h" 29 #include "llvm/ADT/STLExtras.h" 30 #include "llvm/ADT/SmallVector.h" 31 #include "llvm/Analysis/Passes.h" 32 #include "llvm/IR/Constants.h" 33 #include "llvm/IR/Function.h" 34 #include "llvm/IR/Instructions.h" 35 #include "llvm/IR/Module.h" 36 #include "llvm/Transforms/Utils/Cloning.h" 37 38 using namespace llvm; 39 40 namespace { 41 42 StringRef GetImageSizeFunc = "llvm.OpenCL.image.get.size"; 43 StringRef GetImageFormatFunc = "llvm.OpenCL.image.get.format"; 44 StringRef GetImageResourceIDFunc = "llvm.OpenCL.image.get.resource.id"; 45 StringRef GetSamplerResourceIDFunc = "llvm.OpenCL.sampler.get.resource.id"; 46 47 StringRef ImageSizeArgMDType = "__llvm_image_size"; 48 StringRef ImageFormatArgMDType = "__llvm_image_format"; 49 50 StringRef KernelsMDNodeName = "opencl.kernels"; 51 StringRef KernelArgMDNodeNames[] = { 52 "kernel_arg_addr_space", 53 "kernel_arg_access_qual", 54 "kernel_arg_type", 55 "kernel_arg_base_type", 56 "kernel_arg_type_qual"}; 57 const unsigned NumKernelArgMDNodes = 5; 58 59 typedef SmallVector<Metadata *, 8> MDVector; 60 struct KernelArgMD { 61 MDVector ArgVector[NumKernelArgMDNodes]; 62 }; 63 64 } // end anonymous namespace 65 66 static inline bool 67 IsImageType(StringRef TypeString) { 68 return TypeString == "image2d_t" || TypeString == "image3d_t"; 69 } 70 71 static inline bool 72 IsSamplerType(StringRef TypeString) { 73 return TypeString == "sampler_t"; 74 } 75 76 static Function * 77 GetFunctionFromMDNode(MDNode *Node) { 78 if (!Node) 79 return nullptr; 80 81 size_t NumOps = Node->getNumOperands(); 82 if (NumOps != NumKernelArgMDNodes + 1) 83 return nullptr; 84 85 auto F = mdconst::dyn_extract<Function>(Node->getOperand(0)); 86 if (!F) 87 return nullptr; 88 89 // Sanity checks. 90 size_t ExpectNumArgNodeOps = F->arg_size() + 1; 91 for (size_t i = 0; i < NumKernelArgMDNodes; ++i) { 92 MDNode *ArgNode = dyn_cast_or_null<MDNode>(Node->getOperand(i + 1)); 93 if (ArgNode->getNumOperands() != ExpectNumArgNodeOps) 94 return nullptr; 95 if (!ArgNode->getOperand(0)) 96 return nullptr; 97 98 // FIXME: It should be possible to do image lowering when some metadata 99 // args missing or not in the expected order. 100 MDString *StringNode = dyn_cast<MDString>(ArgNode->getOperand(0)); 101 if (!StringNode || StringNode->getString() != KernelArgMDNodeNames[i]) 102 return nullptr; 103 } 104 105 return F; 106 } 107 108 static StringRef 109 AccessQualFromMD(MDNode *KernelMDNode, unsigned ArgIdx) { 110 MDNode *ArgAQNode = cast<MDNode>(KernelMDNode->getOperand(2)); 111 return cast<MDString>(ArgAQNode->getOperand(ArgIdx + 1))->getString(); 112 } 113 114 static StringRef 115 ArgTypeFromMD(MDNode *KernelMDNode, unsigned ArgIdx) { 116 MDNode *ArgTypeNode = cast<MDNode>(KernelMDNode->getOperand(3)); 117 return cast<MDString>(ArgTypeNode->getOperand(ArgIdx + 1))->getString(); 118 } 119 120 static MDVector 121 GetArgMD(MDNode *KernelMDNode, unsigned OpIdx) { 122 MDVector Res; 123 for (unsigned i = 0; i < NumKernelArgMDNodes; ++i) { 124 MDNode *Node = cast<MDNode>(KernelMDNode->getOperand(i + 1)); 125 Res.push_back(Node->getOperand(OpIdx)); 126 } 127 return Res; 128 } 129 130 static void 131 PushArgMD(KernelArgMD &MD, const MDVector &V) { 132 assert(V.size() == NumKernelArgMDNodes); 133 for (unsigned i = 0; i < NumKernelArgMDNodes; ++i) { 134 MD.ArgVector[i].push_back(V[i]); 135 } 136 } 137 138 namespace { 139 140 class AMDGPUOpenCLImageTypeLoweringPass : public ModulePass { 141 static char ID; 142 143 LLVMContext *Context; 144 Type *Int32Type; 145 Type *ImageSizeType; 146 Type *ImageFormatType; 147 SmallVector<Instruction *, 4> InstsToErase; 148 149 bool replaceImageUses(Argument &ImageArg, uint32_t ResourceID, 150 Argument &ImageSizeArg, 151 Argument &ImageFormatArg) { 152 bool Modified = false; 153 154 for (auto &Use : ImageArg.uses()) { 155 auto Inst = dyn_cast<CallInst>(Use.getUser()); 156 if (!Inst) { 157 continue; 158 } 159 160 Function *F = Inst->getCalledFunction(); 161 if (!F) 162 continue; 163 164 Value *Replacement = nullptr; 165 StringRef Name = F->getName(); 166 if (Name.startswith(GetImageResourceIDFunc)) { 167 Replacement = ConstantInt::get(Int32Type, ResourceID); 168 } else if (Name.startswith(GetImageSizeFunc)) { 169 Replacement = &ImageSizeArg; 170 } else if (Name.startswith(GetImageFormatFunc)) { 171 Replacement = &ImageFormatArg; 172 } else { 173 continue; 174 } 175 176 Inst->replaceAllUsesWith(Replacement); 177 InstsToErase.push_back(Inst); 178 Modified = true; 179 } 180 181 return Modified; 182 } 183 184 bool replaceSamplerUses(Argument &SamplerArg, uint32_t ResourceID) { 185 bool Modified = false; 186 187 for (const auto &Use : SamplerArg.uses()) { 188 auto Inst = dyn_cast<CallInst>(Use.getUser()); 189 if (!Inst) { 190 continue; 191 } 192 193 Function *F = Inst->getCalledFunction(); 194 if (!F) 195 continue; 196 197 Value *Replacement = nullptr; 198 StringRef Name = F->getName(); 199 if (Name == GetSamplerResourceIDFunc) { 200 Replacement = ConstantInt::get(Int32Type, ResourceID); 201 } else { 202 continue; 203 } 204 205 Inst->replaceAllUsesWith(Replacement); 206 InstsToErase.push_back(Inst); 207 Modified = true; 208 } 209 210 return Modified; 211 } 212 213 bool replaceImageAndSamplerUses(Function *F, MDNode *KernelMDNode) { 214 uint32_t NumReadOnlyImageArgs = 0; 215 uint32_t NumWriteOnlyImageArgs = 0; 216 uint32_t NumSamplerArgs = 0; 217 218 bool Modified = false; 219 InstsToErase.clear(); 220 for (auto ArgI = F->arg_begin(); ArgI != F->arg_end(); ++ArgI) { 221 Argument &Arg = *ArgI; 222 StringRef Type = ArgTypeFromMD(KernelMDNode, Arg.getArgNo()); 223 224 // Handle image types. 225 if (IsImageType(Type)) { 226 StringRef AccessQual = AccessQualFromMD(KernelMDNode, Arg.getArgNo()); 227 uint32_t ResourceID; 228 if (AccessQual == "read_only") { 229 ResourceID = NumReadOnlyImageArgs++; 230 } else if (AccessQual == "write_only") { 231 ResourceID = NumWriteOnlyImageArgs++; 232 } else { 233 llvm_unreachable("Wrong image access qualifier."); 234 } 235 236 Argument &SizeArg = *(++ArgI); 237 Argument &FormatArg = *(++ArgI); 238 Modified |= replaceImageUses(Arg, ResourceID, SizeArg, FormatArg); 239 240 // Handle sampler type. 241 } else if (IsSamplerType(Type)) { 242 uint32_t ResourceID = NumSamplerArgs++; 243 Modified |= replaceSamplerUses(Arg, ResourceID); 244 } 245 } 246 for (unsigned i = 0; i < InstsToErase.size(); ++i) { 247 InstsToErase[i]->eraseFromParent(); 248 } 249 250 return Modified; 251 } 252 253 std::tuple<Function *, MDNode *> 254 addImplicitArgs(Function *F, MDNode *KernelMDNode) { 255 bool Modified = false; 256 257 FunctionType *FT = F->getFunctionType(); 258 SmallVector<Type *, 8> ArgTypes; 259 260 // Metadata operands for new MDNode. 261 KernelArgMD NewArgMDs; 262 PushArgMD(NewArgMDs, GetArgMD(KernelMDNode, 0)); 263 264 // Add implicit arguments to the signature. 265 for (unsigned i = 0; i < FT->getNumParams(); ++i) { 266 ArgTypes.push_back(FT->getParamType(i)); 267 MDVector ArgMD = GetArgMD(KernelMDNode, i + 1); 268 PushArgMD(NewArgMDs, ArgMD); 269 270 if (!IsImageType(ArgTypeFromMD(KernelMDNode, i))) 271 continue; 272 273 // Add size implicit argument. 274 ArgTypes.push_back(ImageSizeType); 275 ArgMD[2] = ArgMD[3] = MDString::get(*Context, ImageSizeArgMDType); 276 PushArgMD(NewArgMDs, ArgMD); 277 278 // Add format implicit argument. 279 ArgTypes.push_back(ImageFormatType); 280 ArgMD[2] = ArgMD[3] = MDString::get(*Context, ImageFormatArgMDType); 281 PushArgMD(NewArgMDs, ArgMD); 282 283 Modified = true; 284 } 285 if (!Modified) { 286 return std::make_tuple(nullptr, nullptr); 287 } 288 289 // Create function with new signature and clone the old body into it. 290 auto NewFT = FunctionType::get(FT->getReturnType(), ArgTypes, false); 291 auto NewF = Function::Create(NewFT, F->getLinkage(), F->getName()); 292 ValueToValueMapTy VMap; 293 auto NewFArgIt = NewF->arg_begin(); 294 for (auto &Arg: F->args()) { 295 auto ArgName = Arg.getName(); 296 NewFArgIt->setName(ArgName); 297 VMap[&Arg] = &(*NewFArgIt++); 298 if (IsImageType(ArgTypeFromMD(KernelMDNode, Arg.getArgNo()))) { 299 (NewFArgIt++)->setName(Twine("__size_") + ArgName); 300 (NewFArgIt++)->setName(Twine("__format_") + ArgName); 301 } 302 } 303 SmallVector<ReturnInst*, 8> Returns; 304 CloneFunctionInto(NewF, F, VMap, /*ModuleLevelChanges=*/false, Returns); 305 306 // Build new MDNode. 307 SmallVector<llvm::Metadata *, 6> KernelMDArgs; 308 KernelMDArgs.push_back(ConstantAsMetadata::get(NewF)); 309 for (unsigned i = 0; i < NumKernelArgMDNodes; ++i) 310 KernelMDArgs.push_back(MDNode::get(*Context, NewArgMDs.ArgVector[i])); 311 MDNode *NewMDNode = MDNode::get(*Context, KernelMDArgs); 312 313 return std::make_tuple(NewF, NewMDNode); 314 } 315 316 bool transformKernels(Module &M) { 317 NamedMDNode *KernelsMDNode = M.getNamedMetadata(KernelsMDNodeName); 318 if (!KernelsMDNode) 319 return false; 320 321 bool Modified = false; 322 for (unsigned i = 0; i < KernelsMDNode->getNumOperands(); ++i) { 323 MDNode *KernelMDNode = KernelsMDNode->getOperand(i); 324 Function *F = GetFunctionFromMDNode(KernelMDNode); 325 if (!F) 326 continue; 327 328 Function *NewF; 329 MDNode *NewMDNode; 330 std::tie(NewF, NewMDNode) = addImplicitArgs(F, KernelMDNode); 331 if (NewF) { 332 // Replace old function and metadata with new ones. 333 F->eraseFromParent(); 334 M.getFunctionList().push_back(NewF); 335 M.getOrInsertFunction(NewF->getName(), NewF->getFunctionType(), 336 NewF->getAttributes()); 337 KernelsMDNode->setOperand(i, NewMDNode); 338 339 F = NewF; 340 KernelMDNode = NewMDNode; 341 Modified = true; 342 } 343 344 Modified |= replaceImageAndSamplerUses(F, KernelMDNode); 345 } 346 347 return Modified; 348 } 349 350 public: 351 AMDGPUOpenCLImageTypeLoweringPass() : ModulePass(ID) {} 352 353 bool runOnModule(Module &M) override { 354 Context = &M.getContext(); 355 Int32Type = Type::getInt32Ty(M.getContext()); 356 ImageSizeType = ArrayType::get(Int32Type, 3); 357 ImageFormatType = ArrayType::get(Int32Type, 2); 358 359 return transformKernels(M); 360 } 361 362 const char *getPassName() const override { 363 return "AMDGPU OpenCL Image Type Pass"; 364 } 365 }; 366 367 char AMDGPUOpenCLImageTypeLoweringPass::ID = 0; 368 369 } // end anonymous namespace 370 371 ModulePass *llvm::createAMDGPUOpenCLImageTypeLoweringPass() { 372 return new AMDGPUOpenCLImageTypeLoweringPass(); 373 } 374