clang/lib/CodeGen/CGCUDANV.cpp

   1 //===----- CGCUDANV.cpp - Interface to NVIDIA CUDA Runtime ----------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This provides a class for CUDA code generation targeting the NVIDIA CUDA
  10 // runtime library.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "CGCUDARuntime.h"
  15 #include "CGCXXABI.h"
  16 #include "CodeGenFunction.h"
  17 #include "CodeGenModule.h"
  18 #include "clang/AST/Decl.h"
  19 #include "clang/Basic/Cuda.h"
  20 #include "clang/CodeGen/CodeGenABITypes.h"
  21 #include "clang/CodeGen/ConstantInitBuilder.h"
  22 #include "llvm/IR/BasicBlock.h"
  23 #include "llvm/IR/Constants.h"
  24 #include "llvm/IR/DerivedTypes.h"
  25 #include "llvm/IR/ReplaceConstant.h"
  26 #include "llvm/Support/Format.h"
  27 #include "llvm/Support/VirtualFileSystem.h"
  28
  29 using namespace clang;
  30 using namespace CodeGen;
  31
  32 namespace {
  33 constexpr unsigned CudaFatMagic = 0x466243b1;
  34 constexpr unsigned HIPFatMagic = 0x48495046; // "HIPF"
  35
  36 class CGNVCUDARuntime : public CGCUDARuntime {
  37
  38 private:
  39   llvm::IntegerType *IntTy, *SizeTy;
  40   llvm::Type *VoidTy;
  41   llvm::PointerType *CharPtrTy, *VoidPtrTy, *VoidPtrPtrTy;
  42
  43   /// Convenience reference to LLVM Context
  44   llvm::LLVMContext &Context;
  45   /// Convenience reference to the current module
  46   llvm::Module &TheModule;
  47   /// Keeps track of kernel launch stubs and handles emitted in this module
  48   struct KernelInfo {
  49     llvm::Function *Kernel; // stub function to help launch kernel
  50     const Decl *D;
  51   };
  52   llvm::SmallVector<KernelInfo, 16> EmittedKernels;
  53   // Map a kernel mangled name to a symbol for identifying kernel in host code
  54   // For CUDA, the symbol for identifying the kernel is the same as the device
  55   // stub function. For HIP, they are different.
  56   llvm::DenseMap<StringRef, llvm::GlobalValue *> KernelHandles;
  57   // Map a kernel handle to the kernel stub.
  58   llvm::DenseMap<llvm::GlobalValue *, llvm::Function *> KernelStubs;
  59   struct VarInfo {
  60     llvm::GlobalVariable *Var;
  61     const VarDecl *D;
  62     DeviceVarFlags Flags;
  63   };
  64   llvm::SmallVector<VarInfo, 16> DeviceVars;
  65   /// Keeps track of variable containing handle of GPU binary. Populated by
  66   /// ModuleCtorFunction() and used to create corresponding cleanup calls in
  67   /// ModuleDtorFunction()
  68   llvm::GlobalVariable *GpuBinaryHandle = nullptr;
  69   /// Whether we generate relocatable device code.
  70   bool RelocatableDeviceCode;
  71   /// Mangle context for device.
  72   std::unique_ptr<MangleContext> DeviceMC;
  73   /// Some zeros used for GEPs.
  74   llvm::Constant *Zeros[2];
  75
  76   llvm::FunctionCallee getSetupArgumentFn() const;
  77   llvm::FunctionCallee getLaunchFn() const;
  78
  79   llvm::FunctionType *getRegisterGlobalsFnTy() const;
  80   llvm::FunctionType *getCallbackFnTy() const;
  81   llvm::FunctionType *getRegisterLinkedBinaryFnTy() const;
  82   std::string addPrefixToName(StringRef FuncName) const;
  83   std::string addUnderscoredPrefixToName(StringRef FuncName) const;
  84
  85   /// Creates a function to register all kernel stubs generated in this module.
  86   llvm::Function *makeRegisterGlobalsFn();
  87
  88   /// Helper function that generates a constant string and returns a pointer to
  89   /// the start of the string.  The result of this function can be used anywhere
  90   /// where the C code specifies const char*.
  91   llvm::Constant *makeConstantString(const std::string &Str,
  92                                      const std::string &Name = "") {
  93     auto ConstStr = CGM.GetAddrOfConstantCString(Str, Name.c_str());
  94     return llvm::ConstantExpr::getGetElementPtr(ConstStr.getElementType(),
  95                                                 ConstStr.getPointer(), Zeros);
  96   }
  97
  98   /// Helper function which generates an initialized constant array from Str,
  99   /// and optionally sets section name and alignment. AddNull specifies whether
 100   /// the array should nave NUL termination.
 101   llvm::Constant *makeConstantArray(StringRef Str,
 102                                     StringRef Name = "",
 103                                     StringRef SectionName = "",
 104                                     unsigned Alignment = 0,
 105                                     bool AddNull = false) {
 106     llvm::Constant *Value =
 107         llvm::ConstantDataArray::getString(Context, Str, AddNull);
 108     auto *GV = new llvm::GlobalVariable(
 109         TheModule, Value->getType(), /*isConstant=*/true,
 110         llvm::GlobalValue::PrivateLinkage, Value, Name);
 111     if (!SectionName.empty()) {
 112       GV->setSection(SectionName);
 113       // Mark the address as used which make sure that this section isn't
 114       // merged and we will really have it in the object file.
 115       GV->setUnnamedAddr(llvm::GlobalValue::UnnamedAddr::None);
 116     }
 117     if (Alignment)
 118       GV->setAlignment(llvm::Align(Alignment));
 119     return llvm::ConstantExpr::getGetElementPtr(GV->getValueType(), GV, Zeros);
 120   }
 121
 122   /// Helper function that generates an empty dummy function returning void.
 123   llvm::Function *makeDummyFunction(llvm::FunctionType *FnTy) {
 124     assert(FnTy->getReturnType()->isVoidTy() &&
 125            "Can only generate dummy functions returning void!");
 126     llvm::Function *DummyFunc = llvm::Function::Create(
 127         FnTy, llvm::GlobalValue::InternalLinkage, "dummy", &TheModule);
 128
 129     llvm::BasicBlock *DummyBlock =
 130         llvm::BasicBlock::Create(Context, "", DummyFunc);
 131     CGBuilderTy FuncBuilder(CGM, Context);
 132     FuncBuilder.SetInsertPoint(DummyBlock);
 133     FuncBuilder.CreateRetVoid();
 134
 135     return DummyFunc;
 136   }
 137
 138   void emitDeviceStubBodyLegacy(CodeGenFunction &CGF, FunctionArgList &Args);
 139   void emitDeviceStubBodyNew(CodeGenFunction &CGF, FunctionArgList &Args);
 140   std::string getDeviceSideName(const NamedDecl *ND) override;
 141
 142   void registerDeviceVar(const VarDecl *VD, llvm::GlobalVariable &Var,
 143                          bool Extern, bool Constant) {
 144     DeviceVars.push_back({&Var,
 145                           VD,
 146                           {DeviceVarFlags::Variable, Extern, Constant,
 147                            VD->hasAttr<HIPManagedAttr>(),
 148                            /*Normalized*/ false, 0}});
 149   }
 150   void registerDeviceSurf(const VarDecl *VD, llvm::GlobalVariable &Var,
 151                           bool Extern, int Type) {
 152     DeviceVars.push_back({&Var,
 153                           VD,
 154                           {DeviceVarFlags::Surface, Extern, /*Constant*/ false,
 155                            /*Managed*/ false,
 156                            /*Normalized*/ false, Type}});
 157   }
 158   void registerDeviceTex(const VarDecl *VD, llvm::GlobalVariable &Var,
 159                          bool Extern, int Type, bool Normalized) {
 160     DeviceVars.push_back({&Var,
 161                           VD,
 162                           {DeviceVarFlags::Texture, Extern, /*Constant*/ false,
 163                            /*Managed*/ false, Normalized, Type}});
 164   }
 165
 166   /// Creates module constructor function
 167   llvm::Function *makeModuleCtorFunction();
 168   /// Creates module destructor function
 169   llvm::Function *makeModuleDtorFunction();
 170   /// Transform managed variables for device compilation.
 171   void transformManagedVars();
 172   /// Create offloading entries to register globals in RDC mode.
 173   void createOffloadingEntries();
 174
 175 public:
 176   CGNVCUDARuntime(CodeGenModule &CGM);
 177
 178   llvm::GlobalValue *getKernelHandle(llvm::Function *F, GlobalDecl GD) override;
 179   llvm::Function *getKernelStub(llvm::GlobalValue *Handle) override {
 180     auto Loc = KernelStubs.find(Handle);
 181     assert(Loc != KernelStubs.end());
 182     return Loc->second;
 183   }
 184   void emitDeviceStub(CodeGenFunction &CGF, FunctionArgList &Args) override;
 185   void handleVarRegistration(const VarDecl *VD,
 186                              llvm::GlobalVariable &Var) override;
 187   void
 188   internalizeDeviceSideVar(const VarDecl *D,
 189                            llvm::GlobalValue::LinkageTypes &Linkage) override;
 190
 191   llvm::Function *finalizeModule() override;
 192 };
 193
 194 } // end anonymous namespace
 195
 196 std::string CGNVCUDARuntime::addPrefixToName(StringRef FuncName) const {
 197   if (CGM.getLangOpts().HIP)
 198     return ((Twine("hip") + Twine(FuncName)).str());
 199   return ((Twine("cuda") + Twine(FuncName)).str());
 200 }
 201 std::string
 202 CGNVCUDARuntime::addUnderscoredPrefixToName(StringRef FuncName) const {
 203   if (CGM.getLangOpts().HIP)
 204     return ((Twine("__hip") + Twine(FuncName)).str());
 205   return ((Twine("__cuda") + Twine(FuncName)).str());
 206 }
 207
 208 static std::unique_ptr<MangleContext> InitDeviceMC(CodeGenModule &CGM) {
 209   // If the host and device have different C++ ABIs, mark it as the device
 210   // mangle context so that the mangling needs to retrieve the additional
 211   // device lambda mangling number instead of the regular host one.
 212   if (CGM.getContext().getAuxTargetInfo() &&
 213       CGM.getContext().getTargetInfo().getCXXABI().isMicrosoft() &&
 214       CGM.getContext().getAuxTargetInfo()->getCXXABI().isItaniumFamily()) {
 215     return std::unique_ptr<MangleContext>(
 216         CGM.getContext().createDeviceMangleContext(
 217             *CGM.getContext().getAuxTargetInfo()));
 218   }
 219
 220   return std::unique_ptr<MangleContext>(CGM.getContext().createMangleContext(
 221       CGM.getContext().getAuxTargetInfo()));
 222 }
 223
 224 CGNVCUDARuntime::CGNVCUDARuntime(CodeGenModule &CGM)
 225     : CGCUDARuntime(CGM), Context(CGM.getLLVMContext()),
 226       TheModule(CGM.getModule()),
 227       RelocatableDeviceCode(CGM.getLangOpts().GPURelocatableDeviceCode),
 228       DeviceMC(InitDeviceMC(CGM)) {
 229   CodeGen::CodeGenTypes &Types = CGM.getTypes();
 230   ASTContext &Ctx = CGM.getContext();
 231
 232   IntTy = CGM.IntTy;
 233   SizeTy = CGM.SizeTy;
 234   VoidTy = CGM.VoidTy;
 235   Zeros[0] = llvm::ConstantInt::get(SizeTy, 0);
 236   Zeros[1] = Zeros[0];
 237
 238   CharPtrTy = llvm::PointerType::getUnqual(Types.ConvertType(Ctx.CharTy));
 239   VoidPtrTy = cast<llvm::PointerType>(Types.ConvertType(Ctx.VoidPtrTy));
 240   VoidPtrPtrTy = VoidPtrTy->getPointerTo();
 241 }
 242
 243 llvm::FunctionCallee CGNVCUDARuntime::getSetupArgumentFn() const {
 244   // cudaError_t cudaSetupArgument(void *, size_t, size_t)
 245   llvm::Type *Params[] = {VoidPtrTy, SizeTy, SizeTy};
 246   return CGM.CreateRuntimeFunction(
 247       llvm::FunctionType::get(IntTy, Params, false),
 248       addPrefixToName("SetupArgument"));
 249 }
 250
 251 llvm::FunctionCallee CGNVCUDARuntime::getLaunchFn() const {
 252   if (CGM.getLangOpts().HIP) {
 253     // hipError_t hipLaunchByPtr(char *);
 254     return CGM.CreateRuntimeFunction(
 255         llvm::FunctionType::get(IntTy, CharPtrTy, false), "hipLaunchByPtr");
 256   }
 257   // cudaError_t cudaLaunch(char *);
 258   return CGM.CreateRuntimeFunction(
 259       llvm::FunctionType::get(IntTy, CharPtrTy, false), "cudaLaunch");
 260 }
 261
 262 llvm::FunctionType *CGNVCUDARuntime::getRegisterGlobalsFnTy() const {
 263   return llvm::FunctionType::get(VoidTy, VoidPtrPtrTy, false);
 264 }
 265
 266 llvm::FunctionType *CGNVCUDARuntime::getCallbackFnTy() const {
 267   return llvm::FunctionType::get(VoidTy, VoidPtrTy, false);
 268 }
 269
 270 llvm::FunctionType *CGNVCUDARuntime::getRegisterLinkedBinaryFnTy() const {
 271   auto *CallbackFnTy = getCallbackFnTy();
 272   auto *RegisterGlobalsFnTy = getRegisterGlobalsFnTy();
 273   llvm::Type *Params[] = {RegisterGlobalsFnTy->getPointerTo(), VoidPtrTy,
 274                           VoidPtrTy, CallbackFnTy->getPointerTo()};
 275   return llvm::FunctionType::get(VoidTy, Params, false);
 276 }
 277
 278 std::string CGNVCUDARuntime::getDeviceSideName(const NamedDecl *ND) {
 279   GlobalDecl GD;
 280   // D could be either a kernel or a variable.
 281   if (auto *FD = dyn_cast<FunctionDecl>(ND))
 282     GD = GlobalDecl(FD, KernelReferenceKind::Kernel);
 283   else
 284     GD = GlobalDecl(ND);
 285   std::string DeviceSideName;
 286   MangleContext *MC;
 287   if (CGM.getLangOpts().CUDAIsDevice)
 288     MC = &CGM.getCXXABI().getMangleContext();
 289   else
 290     MC = DeviceMC.get();
 291   if (MC->shouldMangleDeclName(ND)) {
 292     SmallString<256> Buffer;
 293     llvm::raw_svector_ostream Out(Buffer);
 294     MC->mangleName(GD, Out);
 295     DeviceSideName = std::string(Out.str());
 296   } else
 297     DeviceSideName = std::string(ND->getIdentifier()->getName());
 298
 299   // Make unique name for device side static file-scope variable for HIP.
 300   if (CGM.getContext().shouldExternalize(ND) &&
 301       CGM.getLangOpts().GPURelocatableDeviceCode) {
 302     SmallString<256> Buffer;
 303     llvm::raw_svector_ostream Out(Buffer);
 304     Out << DeviceSideName;
 305     CGM.printPostfixForExternalizedDecl(Out, ND);
 306     DeviceSideName = std::string(Out.str());
 307   }
 308   return DeviceSideName;
 309 }
 310
 311 void CGNVCUDARuntime::emitDeviceStub(CodeGenFunction &CGF,
 312                                      FunctionArgList &Args) {
 313   EmittedKernels.push_back({CGF.CurFn, CGF.CurFuncDecl});
 314   if (auto *GV =
 315           dyn_cast<llvm::GlobalVariable>(KernelHandles[CGF.CurFn->getName()])) {
 316     GV->setLinkage(CGF.CurFn->getLinkage());
 317     GV->setInitializer(CGF.CurFn);
 318   }
 319   if (CudaFeatureEnabled(CGM.getTarget().getSDKVersion(),
 320                          CudaFeature::CUDA_USES_NEW_LAUNCH) ||
 321       (CGF.getLangOpts().HIP && CGF.getLangOpts().HIPUseNewLaunchAPI))
 322     emitDeviceStubBodyNew(CGF, Args);
 323   else
 324     emitDeviceStubBodyLegacy(CGF, Args);
 325 }
 326
 327 // CUDA 9.0+ uses new way to launch kernels. Parameters are packed in a local
 328 // array and kernels are launched using cudaLaunchKernel().
 329 void CGNVCUDARuntime::emitDeviceStubBodyNew(CodeGenFunction &CGF,
 330                                             FunctionArgList &Args) {
 331   // Build the shadow stack entry at the very start of the function.
 332
 333   // Calculate amount of space we will need for all arguments.  If we have no
 334   // args, allocate a single pointer so we still have a valid pointer to the
 335   // argument array that we can pass to runtime, even if it will be unused.
 336   Address KernelArgs = CGF.CreateTempAlloca(
 337       VoidPtrTy, CharUnits::fromQuantity(16), "kernel_args",
 338       llvm::ConstantInt::get(SizeTy, std::max<size_t>(1, Args.size())));
 339   // Store pointers to the arguments in a locally allocated launch_args.
 340   for (unsigned i = 0; i < Args.size(); ++i) {
 341     llvm::Value* VarPtr = CGF.GetAddrOfLocalVar(Args[i]).getPointer();
 342     llvm::Value *VoidVarPtr = CGF.Builder.CreatePointerCast(VarPtr, VoidPtrTy);
 343     CGF.Builder.CreateDefaultAlignedStore(
 344         VoidVarPtr,
 345         CGF.Builder.CreateConstGEP1_32(VoidPtrTy, KernelArgs.getPointer(), i));
 346   }
 347
 348   llvm::BasicBlock *EndBlock = CGF.createBasicBlock("setup.end");
 349
 350   // Lookup cudaLaunchKernel/hipLaunchKernel function.
 351   // HIP kernel launching API name depends on -fgpu-default-stream option. For
 352   // the default value 'legacy', it is hipLaunchKernel. For 'per-thread',
 353   // it is hipLaunchKernel_spt.
 354   // cudaError_t cudaLaunchKernel(const void *func, dim3 gridDim, dim3 blockDim,
 355   //                              void **args, size_t sharedMem,
 356   //                              cudaStream_t stream);
 357   // hipError_t hipLaunchKernel[_spt](const void *func, dim3 gridDim,
 358   //                                  dim3 blockDim, void **args,
 359   //                                  size_t sharedMem, hipStream_t stream);
 360   TranslationUnitDecl *TUDecl = CGM.getContext().getTranslationUnitDecl();
 361   DeclContext *DC = TranslationUnitDecl::castToDeclContext(TUDecl);
 362   std::string KernelLaunchAPI = "LaunchKernel";
 363   if (CGF.getLangOpts().HIP && CGF.getLangOpts().GPUDefaultStream ==
 364                                    LangOptions::GPUDefaultStreamKind::PerThread)
 365     KernelLaunchAPI = KernelLaunchAPI + "_spt";
 366   auto LaunchKernelName = addPrefixToName(KernelLaunchAPI);
 367   IdentifierInfo &cudaLaunchKernelII =
 368       CGM.getContext().Idents.get(LaunchKernelName);
 369   FunctionDecl *cudaLaunchKernelFD = nullptr;
 370   for (auto *Result : DC->lookup(&cudaLaunchKernelII)) {
 371     if (FunctionDecl *FD = dyn_cast<FunctionDecl>(Result))
 372       cudaLaunchKernelFD = FD;
 373   }
 374
 375   if (cudaLaunchKernelFD == nullptr) {
 376     CGM.Error(CGF.CurFuncDecl->getLocation(),
 377               "Can't find declaration for " + LaunchKernelName);
 378     return;
 379   }
 380   // Create temporary dim3 grid_dim, block_dim.
 381   ParmVarDecl *GridDimParam = cudaLaunchKernelFD->getParamDecl(1);
 382   QualType Dim3Ty = GridDimParam->getType();
 383   Address GridDim =
 384       CGF.CreateMemTemp(Dim3Ty, CharUnits::fromQuantity(8), "grid_dim");
 385   Address BlockDim =
 386       CGF.CreateMemTemp(Dim3Ty, CharUnits::fromQuantity(8), "block_dim");
 387   Address ShmemSize =
 388       CGF.CreateTempAlloca(SizeTy, CGM.getSizeAlign(), "shmem_size");
 389   Address Stream =
 390       CGF.CreateTempAlloca(VoidPtrTy, CGM.getPointerAlign(), "stream");
 391   llvm::FunctionCallee cudaPopConfigFn = CGM.CreateRuntimeFunction(
 392       llvm::FunctionType::get(IntTy,
 393                               {/*gridDim=*/GridDim.getType(),
 394                                /*blockDim=*/BlockDim.getType(),
 395                                /*ShmemSize=*/ShmemSize.getType(),
 396                                /*Stream=*/Stream.getType()},
 397                               /*isVarArg=*/false),
 398       addUnderscoredPrefixToName("PopCallConfiguration"));
 399
 400   CGF.EmitRuntimeCallOrInvoke(cudaPopConfigFn,
 401                               {GridDim.getPointer(), BlockDim.getPointer(),
 402                                ShmemSize.getPointer(), Stream.getPointer()});
 403
 404   // Emit the call to cudaLaunch
 405   llvm::Value *Kernel = CGF.Builder.CreatePointerCast(
 406       KernelHandles[CGF.CurFn->getName()], VoidPtrTy);
 407   CallArgList LaunchKernelArgs;
 408   LaunchKernelArgs.add(RValue::get(Kernel),
 409                        cudaLaunchKernelFD->getParamDecl(0)->getType());
 410   LaunchKernelArgs.add(RValue::getAggregate(GridDim), Dim3Ty);
 411   LaunchKernelArgs.add(RValue::getAggregate(BlockDim), Dim3Ty);
 412   LaunchKernelArgs.add(RValue::get(KernelArgs.getPointer()),
 413                        cudaLaunchKernelFD->getParamDecl(3)->getType());
 414   LaunchKernelArgs.add(RValue::get(CGF.Builder.CreateLoad(ShmemSize)),
 415                        cudaLaunchKernelFD->getParamDecl(4)->getType());
 416   LaunchKernelArgs.add(RValue::get(CGF.Builder.CreateLoad(Stream)),
 417                        cudaLaunchKernelFD->getParamDecl(5)->getType());
 418
 419   QualType QT = cudaLaunchKernelFD->getType();
 420   QualType CQT = QT.getCanonicalType();
 421   llvm::Type *Ty = CGM.getTypes().ConvertType(CQT);
 422   llvm::FunctionType *FTy = cast<llvm::FunctionType>(Ty);
 423
 424   const CGFunctionInfo &FI =
 425       CGM.getTypes().arrangeFunctionDeclaration(cudaLaunchKernelFD);
 426   llvm::FunctionCallee cudaLaunchKernelFn =
 427       CGM.CreateRuntimeFunction(FTy, LaunchKernelName);
 428   CGF.EmitCall(FI, CGCallee::forDirect(cudaLaunchKernelFn), ReturnValueSlot(),
 429                LaunchKernelArgs);
 430   CGF.EmitBranch(EndBlock);
 431
 432   CGF.EmitBlock(EndBlock);
 433 }
 434
 435 void CGNVCUDARuntime::emitDeviceStubBodyLegacy(CodeGenFunction &CGF,
 436                                                FunctionArgList &Args) {
 437   // Emit a call to cudaSetupArgument for each arg in Args.
 438   llvm::FunctionCallee cudaSetupArgFn = getSetupArgumentFn();
 439   llvm::BasicBlock *EndBlock = CGF.createBasicBlock("setup.end");
 440   CharUnits Offset = CharUnits::Zero();
 441   for (const VarDecl *A : Args) {
 442     auto TInfo = CGM.getContext().getTypeInfoInChars(A->getType());
 443     Offset = Offset.alignTo(TInfo.Align);
 444     llvm::Value *Args[] = {
 445         CGF.Builder.CreatePointerCast(CGF.GetAddrOfLocalVar(A).getPointer(),
 446                                       VoidPtrTy),
 447         llvm::ConstantInt::get(SizeTy, TInfo.Width.getQuantity()),
 448         llvm::ConstantInt::get(SizeTy, Offset.getQuantity()),
 449     };
 450     llvm::CallBase *CB = CGF.EmitRuntimeCallOrInvoke(cudaSetupArgFn, Args);
 451     llvm::Constant *Zero = llvm::ConstantInt::get(IntTy, 0);
 452     llvm::Value *CBZero = CGF.Builder.CreateICmpEQ(CB, Zero);
 453     llvm::BasicBlock *NextBlock = CGF.createBasicBlock("setup.next");
 454     CGF.Builder.CreateCondBr(CBZero, NextBlock, EndBlock);
 455     CGF.EmitBlock(NextBlock);
 456     Offset += TInfo.Width;
 457   }
 458
 459   // Emit the call to cudaLaunch
 460   llvm::FunctionCallee cudaLaunchFn = getLaunchFn();
 461   llvm::Value *Arg = CGF.Builder.CreatePointerCast(
 462       KernelHandles[CGF.CurFn->getName()], CharPtrTy);
 463   CGF.EmitRuntimeCallOrInvoke(cudaLaunchFn, Arg);
 464   CGF.EmitBranch(EndBlock);
 465
 466   CGF.EmitBlock(EndBlock);
 467 }
 468
 469 // Replace the original variable Var with the address loaded from variable
 470 // ManagedVar populated by HIP runtime.
 471 static void replaceManagedVar(llvm::GlobalVariable *Var,
 472                               llvm::GlobalVariable *ManagedVar) {
 473   SmallVector<SmallVector<llvm::User *, 8>, 8> WorkList;
 474   for (auto &&VarUse : Var->uses()) {
 475     WorkList.push_back({VarUse.getUser()});
 476   }
 477   while (!WorkList.empty()) {
 478     auto &&WorkItem = WorkList.pop_back_val();
 479     auto *U = WorkItem.back();
 480     if (isa<llvm::ConstantExpr>(U)) {
 481       for (auto &&UU : U->uses()) {
 482         WorkItem.push_back(UU.getUser());
 483         WorkList.push_back(WorkItem);
 484         WorkItem.pop_back();
 485       }
 486       continue;
 487     }
 488     if (auto *I = dyn_cast<llvm::Instruction>(U)) {
 489       llvm::Value *OldV = Var;
 490       llvm::Instruction *NewV =
 491           new llvm::LoadInst(Var->getType(), ManagedVar, "ld.managed", false,
 492                              llvm::Align(Var->getAlignment()), I);
 493       WorkItem.pop_back();
 494       // Replace constant expressions directly or indirectly using the managed
 495       // variable with instructions.
 496       for (auto &&Op : WorkItem) {
 497         auto *CE = cast<llvm::ConstantExpr>(Op);
 498         auto *NewInst = CE->getAsInstruction(I);
 499         NewInst->replaceUsesOfWith(OldV, NewV);
 500         OldV = CE;
 501         NewV = NewInst;
 502       }
 503       I->replaceUsesOfWith(OldV, NewV);
 504     } else {
 505       llvm_unreachable("Invalid use of managed variable");
 506     }
 507   }
 508 }
 509
 510 /// Creates a function that sets up state on the host side for CUDA objects that
 511 /// have a presence on both the host and device sides. Specifically, registers
 512 /// the host side of kernel functions and device global variables with the CUDA
 513 /// runtime.
 514 /// \code
 515 /// void __cuda_register_globals(void** GpuBinaryHandle) {
 516 ///    __cudaRegisterFunction(GpuBinaryHandle,Kernel0,...);
 517 ///    ...
 518 ///    __cudaRegisterFunction(GpuBinaryHandle,KernelM,...);
 519 ///    __cudaRegisterVar(GpuBinaryHandle, GlobalVar0, ...);
 520 ///    ...
 521 ///    __cudaRegisterVar(GpuBinaryHandle, GlobalVarN, ...);
 522 /// }
 523 /// \endcode
 524 llvm::Function *CGNVCUDARuntime::makeRegisterGlobalsFn() {
 525   // No need to register anything
 526   if (EmittedKernels.empty() && DeviceVars.empty())
 527     return nullptr;
 528
 529   llvm::Function *RegisterKernelsFunc = llvm::Function::Create(
 530       getRegisterGlobalsFnTy(), llvm::GlobalValue::InternalLinkage,
 531       addUnderscoredPrefixToName("_register_globals"), &TheModule);
 532   llvm::BasicBlock *EntryBB =
 533       llvm::BasicBlock::Create(Context, "entry", RegisterKernelsFunc);
 534   CGBuilderTy Builder(CGM, Context);
 535   Builder.SetInsertPoint(EntryBB);
 536
 537   // void __cudaRegisterFunction(void **, const char *, char *, const char *,
 538   //                             int, uint3*, uint3*, dim3*, dim3*, int*)
 539   llvm::Type *RegisterFuncParams[] = {
 540       VoidPtrPtrTy, CharPtrTy, CharPtrTy, CharPtrTy, IntTy,
 541       VoidPtrTy,    VoidPtrTy, VoidPtrTy, VoidPtrTy, IntTy->getPointerTo()};
 542   llvm::FunctionCallee RegisterFunc = CGM.CreateRuntimeFunction(
 543       llvm::FunctionType::get(IntTy, RegisterFuncParams, false),
 544       addUnderscoredPrefixToName("RegisterFunction"));
 545
 546   // Extract GpuBinaryHandle passed as the first argument passed to
 547   // __cuda_register_globals() and generate __cudaRegisterFunction() call for
 548   // each emitted kernel.
 549   llvm::Argument &GpuBinaryHandlePtr = *RegisterKernelsFunc->arg_begin();
 550   for (auto &&I : EmittedKernels) {
 551     llvm::Constant *KernelName =
 552         makeConstantString(getDeviceSideName(cast<NamedDecl>(I.D)));
 553     llvm::Constant *NullPtr = llvm::ConstantPointerNull::get(VoidPtrTy);
 554     llvm::Value *Args[] = {
 555         &GpuBinaryHandlePtr,
 556         Builder.CreateBitCast(KernelHandles[I.Kernel->getName()], VoidPtrTy),
 557         KernelName,
 558         KernelName,
 559         llvm::ConstantInt::get(IntTy, -1),
 560         NullPtr,
 561         NullPtr,
 562         NullPtr,
 563         NullPtr,
 564         llvm::ConstantPointerNull::get(IntTy->getPointerTo())};
 565     Builder.CreateCall(RegisterFunc, Args);
 566   }
 567
 568   llvm::Type *VarSizeTy = IntTy;
 569   // For HIP or CUDA 9.0+, device variable size is type of `size_t`.
 570   if (CGM.getLangOpts().HIP ||
 571       ToCudaVersion(CGM.getTarget().getSDKVersion()) >= CudaVersion::CUDA_90)
 572     VarSizeTy = SizeTy;
 573
 574   // void __cudaRegisterVar(void **, char *, char *, const char *,
 575   //                        int, int, int, int)
 576   llvm::Type *RegisterVarParams[] = {VoidPtrPtrTy, CharPtrTy, CharPtrTy,
 577                                      CharPtrTy,    IntTy,     VarSizeTy,
 578                                      IntTy,        IntTy};
 579   llvm::FunctionCallee RegisterVar = CGM.CreateRuntimeFunction(
 580       llvm::FunctionType::get(VoidTy, RegisterVarParams, false),
 581       addUnderscoredPrefixToName("RegisterVar"));
 582   // void __hipRegisterManagedVar(void **, char *, char *, const char *,
 583   //                              size_t, unsigned)
 584   llvm::Type *RegisterManagedVarParams[] = {VoidPtrPtrTy, CharPtrTy, CharPtrTy,
 585                                             CharPtrTy,    VarSizeTy, IntTy};
 586   llvm::FunctionCallee RegisterManagedVar = CGM.CreateRuntimeFunction(
 587       llvm::FunctionType::get(VoidTy, RegisterManagedVarParams, false),
 588       addUnderscoredPrefixToName("RegisterManagedVar"));
 589   // void __cudaRegisterSurface(void **, const struct surfaceReference *,
 590   //                            const void **, const char *, int, int);
 591   llvm::FunctionCallee RegisterSurf = CGM.CreateRuntimeFunction(
 592       llvm::FunctionType::get(
 593           VoidTy, {VoidPtrPtrTy, VoidPtrTy, CharPtrTy, CharPtrTy, IntTy, IntTy},
 594           false),
 595       addUnderscoredPrefixToName("RegisterSurface"));
 596   // void __cudaRegisterTexture(void **, const struct textureReference *,
 597   //                            const void **, const char *, int, int, int)
 598   llvm::FunctionCallee RegisterTex = CGM.CreateRuntimeFunction(
 599       llvm::FunctionType::get(
 600           VoidTy,
 601           {VoidPtrPtrTy, VoidPtrTy, CharPtrTy, CharPtrTy, IntTy, IntTy, IntTy},
 602           false),
 603       addUnderscoredPrefixToName("RegisterTexture"));
 604   for (auto &&Info : DeviceVars) {
 605     llvm::GlobalVariable *Var = Info.Var;
 606     assert((!Var->isDeclaration() || Info.Flags.isManaged()) &&
 607            "External variables should not show up here, except HIP managed "
 608            "variables");
 609     llvm::Constant *VarName = makeConstantString(getDeviceSideName(Info.D));
 610     switch (Info.Flags.getKind()) {
 611     case DeviceVarFlags::Variable: {
 612       uint64_t VarSize =
 613           CGM.getDataLayout().getTypeAllocSize(Var->getValueType());
 614       if (Info.Flags.isManaged()) {
 615         auto *ManagedVar = new llvm::GlobalVariable(
 616             CGM.getModule(), Var->getType(),
 617             /*isConstant=*/false, Var->getLinkage(),
 618             /*Init=*/Var->isDeclaration()
 619                 ? nullptr
 620                 : llvm::ConstantPointerNull::get(Var->getType()),
 621             /*Name=*/"", /*InsertBefore=*/nullptr,
 622             llvm::GlobalVariable::NotThreadLocal);
 623         ManagedVar->setDSOLocal(Var->isDSOLocal());
 624         ManagedVar->setVisibility(Var->getVisibility());
 625         ManagedVar->setExternallyInitialized(true);
 626         ManagedVar->takeName(Var);
 627         Var->setName(Twine(ManagedVar->getName() + ".managed"));
 628         replaceManagedVar(Var, ManagedVar);
 629         llvm::Value *Args[] = {
 630             &GpuBinaryHandlePtr,
 631             Builder.CreateBitCast(ManagedVar, VoidPtrTy),
 632             Builder.CreateBitCast(Var, VoidPtrTy),
 633             VarName,
 634             llvm::ConstantInt::get(VarSizeTy, VarSize),
 635             llvm::ConstantInt::get(IntTy, Var->getAlignment())};
 636         if (!Var->isDeclaration())
 637           Builder.CreateCall(RegisterManagedVar, Args);
 638       } else {
 639         llvm::Value *Args[] = {
 640             &GpuBinaryHandlePtr,
 641             Builder.CreateBitCast(Var, VoidPtrTy),
 642             VarName,
 643             VarName,
 644             llvm::ConstantInt::get(IntTy, Info.Flags.isExtern()),
 645             llvm::ConstantInt::get(VarSizeTy, VarSize),
 646             llvm::ConstantInt::get(IntTy, Info.Flags.isConstant()),
 647             llvm::ConstantInt::get(IntTy, 0)};
 648         Builder.CreateCall(RegisterVar, Args);
 649       }
 650       break;
 651     }
 652     case DeviceVarFlags::Surface:
 653       Builder.CreateCall(
 654           RegisterSurf,
 655           {&GpuBinaryHandlePtr, Builder.CreateBitCast(Var, VoidPtrTy), VarName,
 656            VarName, llvm::ConstantInt::get(IntTy, Info.Flags.getSurfTexType()),
 657            llvm::ConstantInt::get(IntTy, Info.Flags.isExtern())});
 658       break;
 659     case DeviceVarFlags::Texture:
 660       Builder.CreateCall(
 661           RegisterTex,
 662           {&GpuBinaryHandlePtr, Builder.CreateBitCast(Var, VoidPtrTy), VarName,
 663            VarName, llvm::ConstantInt::get(IntTy, Info.Flags.getSurfTexType()),
 664            llvm::ConstantInt::get(IntTy, Info.Flags.isNormalized()),
 665            llvm::ConstantInt::get(IntTy, Info.Flags.isExtern())});
 666       break;
 667     }
 668   }
 669
 670   Builder.CreateRetVoid();
 671   return RegisterKernelsFunc;
 672 }
 673
 674 /// Creates a global constructor function for the module:
 675 ///
 676 /// For CUDA:
 677 /// \code
 678 /// void __cuda_module_ctor() {
 679 ///     Handle = __cudaRegisterFatBinary(GpuBinaryBlob);
 680 ///     __cuda_register_globals(Handle);
 681 /// }
 682 /// \endcode
 683 ///
 684 /// For HIP:
 685 /// \code
 686 /// void __hip_module_ctor() {
 687 ///     if (__hip_gpubin_handle == 0) {
 688 ///         __hip_gpubin_handle  = __hipRegisterFatBinary(GpuBinaryBlob);
 689 ///         __hip_register_globals(__hip_gpubin_handle);
 690 ///     }
 691 /// }
 692 /// \endcode
 693 llvm::Function *CGNVCUDARuntime::makeModuleCtorFunction() {
 694   bool IsHIP = CGM.getLangOpts().HIP;
 695   bool IsCUDA = CGM.getLangOpts().CUDA;
 696   // No need to generate ctors/dtors if there is no GPU binary.
 697   StringRef CudaGpuBinaryFileName = CGM.getCodeGenOpts().CudaGpuBinaryFileName;
 698   if (CudaGpuBinaryFileName.empty() && !IsHIP)
 699     return nullptr;
 700   if ((IsHIP || (IsCUDA && !RelocatableDeviceCode)) && EmittedKernels.empty() &&
 701       DeviceVars.empty())
 702     return nullptr;
 703
 704   // void __{cuda|hip}_register_globals(void* handle);
 705   llvm::Function *RegisterGlobalsFunc = makeRegisterGlobalsFn();
 706   // We always need a function to pass in as callback. Create a dummy
 707   // implementation if we don't need to register anything.
 708   if (RelocatableDeviceCode && !RegisterGlobalsFunc)
 709     RegisterGlobalsFunc = makeDummyFunction(getRegisterGlobalsFnTy());
 710
 711   // void ** __{cuda|hip}RegisterFatBinary(void *);
 712   llvm::FunctionCallee RegisterFatbinFunc = CGM.CreateRuntimeFunction(
 713       llvm::FunctionType::get(VoidPtrPtrTy, VoidPtrTy, false),
 714       addUnderscoredPrefixToName("RegisterFatBinary"));
 715   // struct { int magic, int version, void * gpu_binary, void * dont_care };
 716   llvm::StructType *FatbinWrapperTy =
 717       llvm::StructType::get(IntTy, IntTy, VoidPtrTy, VoidPtrTy);
 718
 719   // Register GPU binary with the CUDA runtime, store returned handle in a
 720   // global variable and save a reference in GpuBinaryHandle to be cleaned up
 721   // in destructor on exit. Then associate all known kernels with the GPU binary
 722   // handle so CUDA runtime can figure out what to call on the GPU side.
 723   std::unique_ptr<llvm::MemoryBuffer> CudaGpuBinary = nullptr;
 724   if (!CudaGpuBinaryFileName.empty()) {
 725     auto VFS = CGM.getFileSystem();
 726     auto CudaGpuBinaryOrErr =
 727         VFS->getBufferForFile(CudaGpuBinaryFileName, -1, false);
 728     if (std::error_code EC = CudaGpuBinaryOrErr.getError()) {
 729       CGM.getDiags().Report(diag::err_cannot_open_file)
 730           << CudaGpuBinaryFileName << EC.message();
 731       return nullptr;
 732     }
 733     CudaGpuBinary = std::move(CudaGpuBinaryOrErr.get());
 734   }
 735
 736   llvm::Function *ModuleCtorFunc = llvm::Function::Create(
 737       llvm::FunctionType::get(VoidTy, false),
 738       llvm::GlobalValue::InternalLinkage,
 739       addUnderscoredPrefixToName("_module_ctor"), &TheModule);
 740   llvm::BasicBlock *CtorEntryBB =
 741       llvm::BasicBlock::Create(Context, "entry", ModuleCtorFunc);
 742   CGBuilderTy CtorBuilder(CGM, Context);
 743
 744   CtorBuilder.SetInsertPoint(CtorEntryBB);
 745
 746   const char *FatbinConstantName;
 747   const char *FatbinSectionName;
 748   const char *ModuleIDSectionName;
 749   StringRef ModuleIDPrefix;
 750   llvm::Constant *FatBinStr;
 751   unsigned FatMagic;
 752   if (IsHIP) {
 753     FatbinConstantName = ".hip_fatbin";
 754     FatbinSectionName = ".hipFatBinSegment";
 755
 756     ModuleIDSectionName = "__hip_module_id";
 757     ModuleIDPrefix = "__hip_";
 758
 759     if (CudaGpuBinary) {
 760       // If fatbin is available from early finalization, create a string
 761       // literal containing the fat binary loaded from the given file.
 762       const unsigned HIPCodeObjectAlign = 4096;
 763       FatBinStr = makeConstantArray(std::string(CudaGpuBinary->getBuffer()), "",
 764                                     FatbinConstantName, HIPCodeObjectAlign);
 765     } else {
 766       // If fatbin is not available, create an external symbol
 767       // __hip_fatbin in section .hip_fatbin. The external symbol is supposed
 768       // to contain the fat binary but will be populated somewhere else,
 769       // e.g. by lld through link script.
 770       FatBinStr = new llvm::GlobalVariable(
 771         CGM.getModule(), CGM.Int8Ty,
 772         /*isConstant=*/true, llvm::GlobalValue::ExternalLinkage, nullptr,
 773         "__hip_fatbin", nullptr,
 774         llvm::GlobalVariable::NotThreadLocal);
 775       cast<llvm::GlobalVariable>(FatBinStr)->setSection(FatbinConstantName);
 776     }
 777
 778     FatMagic = HIPFatMagic;
 779   } else {
 780     if (RelocatableDeviceCode)
 781       FatbinConstantName = CGM.getTriple().isMacOSX()
 782                                ? "__NV_CUDA,__nv_relfatbin"
 783                                : "__nv_relfatbin";
 784     else
 785       FatbinConstantName =
 786           CGM.getTriple().isMacOSX() ? "__NV_CUDA,__nv_fatbin" : ".nv_fatbin";
 787     // NVIDIA's cuobjdump looks for fatbins in this section.
 788     FatbinSectionName =
 789         CGM.getTriple().isMacOSX() ? "__NV_CUDA,__fatbin" : ".nvFatBinSegment";
 790
 791     ModuleIDSectionName = CGM.getTriple().isMacOSX()
 792                               ? "__NV_CUDA,__nv_module_id"
 793                               : "__nv_module_id";
 794     ModuleIDPrefix = "__nv_";
 795
 796     // For CUDA, create a string literal containing the fat binary loaded from
 797     // the given file.
 798     FatBinStr = makeConstantArray(std::string(CudaGpuBinary->getBuffer()), "",
 799                                   FatbinConstantName, 8);
 800     FatMagic = CudaFatMagic;
 801   }
 802
 803   // Create initialized wrapper structure that points to the loaded GPU binary
 804   ConstantInitBuilder Builder(CGM);
 805   auto Values = Builder.beginStruct(FatbinWrapperTy);
 806   // Fatbin wrapper magic.
 807   Values.addInt(IntTy, FatMagic);
 808   // Fatbin version.
 809   Values.addInt(IntTy, 1);
 810   // Data.
 811   Values.add(FatBinStr);
 812   // Unused in fatbin v1.
 813   Values.add(llvm::ConstantPointerNull::get(VoidPtrTy));
 814   llvm::GlobalVariable *FatbinWrapper = Values.finishAndCreateGlobal(
 815       addUnderscoredPrefixToName("_fatbin_wrapper"), CGM.getPointerAlign(),
 816       /*constant*/ true);
 817   FatbinWrapper->setSection(FatbinSectionName);
 818
 819   // There is only one HIP fat binary per linked module, however there are
 820   // multiple constructor functions. Make sure the fat binary is registered
 821   // only once. The constructor functions are executed by the dynamic loader
 822   // before the program gains control. The dynamic loader cannot execute the
 823   // constructor functions concurrently since doing that would not guarantee
 824   // thread safety of the loaded program. Therefore we can assume sequential
 825   // execution of constructor functions here.
 826   if (IsHIP) {
 827     auto Linkage = CudaGpuBinary ? llvm::GlobalValue::InternalLinkage :
 828         llvm::GlobalValue::LinkOnceAnyLinkage;
 829     llvm::BasicBlock *IfBlock =
 830         llvm::BasicBlock::Create(Context, "if", ModuleCtorFunc);
 831     llvm::BasicBlock *ExitBlock =
 832         llvm::BasicBlock::Create(Context, "exit", ModuleCtorFunc);
 833     // The name, size, and initialization pattern of this variable is part
 834     // of HIP ABI.
 835     GpuBinaryHandle = new llvm::GlobalVariable(
 836         TheModule, VoidPtrPtrTy, /*isConstant=*/false,
 837         Linkage,
 838         /*Initializer=*/llvm::ConstantPointerNull::get(VoidPtrPtrTy),
 839         "__hip_gpubin_handle");
 840     if (Linkage == llvm::GlobalValue::LinkOnceAnyLinkage)
 841       GpuBinaryHandle->setComdat(
 842           CGM.getModule().getOrInsertComdat(GpuBinaryHandle->getName()));
 843     GpuBinaryHandle->setAlignment(CGM.getPointerAlign().getAsAlign());
 844     // Prevent the weak symbol in different shared libraries being merged.
 845     if (Linkage != llvm::GlobalValue::InternalLinkage)
 846       GpuBinaryHandle->setVisibility(llvm::GlobalValue::HiddenVisibility);
 847     Address GpuBinaryAddr(
 848         GpuBinaryHandle, VoidPtrPtrTy,
 849         CharUnits::fromQuantity(GpuBinaryHandle->getAlignment()));
 850     {
 851       auto *HandleValue = CtorBuilder.CreateLoad(GpuBinaryAddr);
 852       llvm::Constant *Zero =
 853           llvm::Constant::getNullValue(HandleValue->getType());
 854       llvm::Value *EQZero = CtorBuilder.CreateICmpEQ(HandleValue, Zero);
 855       CtorBuilder.CreateCondBr(EQZero, IfBlock, ExitBlock);
 856     }
 857     {
 858       CtorBuilder.SetInsertPoint(IfBlock);
 859       // GpuBinaryHandle = __hipRegisterFatBinary(&FatbinWrapper);
 860       llvm::CallInst *RegisterFatbinCall = CtorBuilder.CreateCall(
 861           RegisterFatbinFunc,
 862           CtorBuilder.CreateBitCast(FatbinWrapper, VoidPtrTy));
 863       CtorBuilder.CreateStore(RegisterFatbinCall, GpuBinaryAddr);
 864       CtorBuilder.CreateBr(ExitBlock);
 865     }
 866     {
 867       CtorBuilder.SetInsertPoint(ExitBlock);
 868       // Call __hip_register_globals(GpuBinaryHandle);
 869       if (RegisterGlobalsFunc) {
 870         auto *HandleValue = CtorBuilder.CreateLoad(GpuBinaryAddr);
 871         CtorBuilder.CreateCall(RegisterGlobalsFunc, HandleValue);
 872       }
 873     }
 874   } else if (!RelocatableDeviceCode) {
 875     // Register binary with CUDA runtime. This is substantially different in
 876     // default mode vs. separate compilation!
 877     // GpuBinaryHandle = __cudaRegisterFatBinary(&FatbinWrapper);
 878     llvm::CallInst *RegisterFatbinCall = CtorBuilder.CreateCall(
 879         RegisterFatbinFunc,
 880         CtorBuilder.CreateBitCast(FatbinWrapper, VoidPtrTy));
 881     GpuBinaryHandle = new llvm::GlobalVariable(
 882         TheModule, VoidPtrPtrTy, false, llvm::GlobalValue::InternalLinkage,
 883         llvm::ConstantPointerNull::get(VoidPtrPtrTy), "__cuda_gpubin_handle");
 884     GpuBinaryHandle->setAlignment(CGM.getPointerAlign().getAsAlign());
 885     CtorBuilder.CreateAlignedStore(RegisterFatbinCall, GpuBinaryHandle,
 886                                    CGM.getPointerAlign());
 887
 888     // Call __cuda_register_globals(GpuBinaryHandle);
 889     if (RegisterGlobalsFunc)
 890       CtorBuilder.CreateCall(RegisterGlobalsFunc, RegisterFatbinCall);
 891
 892     // Call __cudaRegisterFatBinaryEnd(Handle) if this CUDA version needs it.
 893     if (CudaFeatureEnabled(CGM.getTarget().getSDKVersion(),
 894                            CudaFeature::CUDA_USES_FATBIN_REGISTER_END)) {
 895       // void __cudaRegisterFatBinaryEnd(void **);
 896       llvm::FunctionCallee RegisterFatbinEndFunc = CGM.CreateRuntimeFunction(
 897           llvm::FunctionType::get(VoidTy, VoidPtrPtrTy, false),
 898           "__cudaRegisterFatBinaryEnd");
 899       CtorBuilder.CreateCall(RegisterFatbinEndFunc, RegisterFatbinCall);
 900     }
 901   } else {
 902     // Generate a unique module ID.
 903     SmallString<64> ModuleID;
 904     llvm::raw_svector_ostream OS(ModuleID);
 905     OS << ModuleIDPrefix << llvm::format("%" PRIx64, FatbinWrapper->getGUID());
 906     llvm::Constant *ModuleIDConstant = makeConstantArray(
 907         std::string(ModuleID.str()), "", ModuleIDSectionName, 32, /*AddNull=*/true);
 908
 909     // Create an alias for the FatbinWrapper that nvcc will look for.
 910     llvm::GlobalAlias::create(llvm::GlobalValue::ExternalLinkage,
 911                               Twine("__fatbinwrap") + ModuleID, FatbinWrapper);
 912
 913     // void __cudaRegisterLinkedBinary%ModuleID%(void (*)(void *), void *,
 914     // void *, void (*)(void **))
 915     SmallString<128> RegisterLinkedBinaryName("__cudaRegisterLinkedBinary");
 916     RegisterLinkedBinaryName += ModuleID;
 917     llvm::FunctionCallee RegisterLinkedBinaryFunc = CGM.CreateRuntimeFunction(
 918         getRegisterLinkedBinaryFnTy(), RegisterLinkedBinaryName);
 919
 920     assert(RegisterGlobalsFunc && "Expecting at least dummy function!");
 921     llvm::Value *Args[] = {RegisterGlobalsFunc,
 922                            CtorBuilder.CreateBitCast(FatbinWrapper, VoidPtrTy),
 923                            ModuleIDConstant,
 924                            makeDummyFunction(getCallbackFnTy())};
 925     CtorBuilder.CreateCall(RegisterLinkedBinaryFunc, Args);
 926   }
 927
 928   // Create destructor and register it with atexit() the way NVCC does it. Doing
 929   // it during regular destructor phase worked in CUDA before 9.2 but results in
 930   // double-free in 9.2.
 931   if (llvm::Function *CleanupFn = makeModuleDtorFunction()) {
 932     // extern "C" int atexit(void (*f)(void));
 933     llvm::FunctionType *AtExitTy =
 934         llvm::FunctionType::get(IntTy, CleanupFn->getType(), false);
 935     llvm::FunctionCallee AtExitFunc =
 936         CGM.CreateRuntimeFunction(AtExitTy, "atexit", llvm::AttributeList(),
 937                                   /*Local=*/true);
 938     CtorBuilder.CreateCall(AtExitFunc, CleanupFn);
 939   }
 940
 941   CtorBuilder.CreateRetVoid();
 942   return ModuleCtorFunc;
 943 }
 944
 945 /// Creates a global destructor function that unregisters the GPU code blob
 946 /// registered by constructor.
 947 ///
 948 /// For CUDA:
 949 /// \code
 950 /// void __cuda_module_dtor() {
 951 ///     __cudaUnregisterFatBinary(Handle);
 952 /// }
 953 /// \endcode
 954 ///
 955 /// For HIP:
 956 /// \code
 957 /// void __hip_module_dtor() {
 958 ///     if (__hip_gpubin_handle) {
 959 ///         __hipUnregisterFatBinary(__hip_gpubin_handle);
 960 ///         __hip_gpubin_handle = 0;
 961 ///     }
 962 /// }
 963 /// \endcode
 964 llvm::Function *CGNVCUDARuntime::makeModuleDtorFunction() {
 965   // No need for destructor if we don't have a handle to unregister.
 966   if (!GpuBinaryHandle)
 967     return nullptr;
 968
 969   // void __cudaUnregisterFatBinary(void ** handle);
 970   llvm::FunctionCallee UnregisterFatbinFunc = CGM.CreateRuntimeFunction(
 971       llvm::FunctionType::get(VoidTy, VoidPtrPtrTy, false),
 972       addUnderscoredPrefixToName("UnregisterFatBinary"));
 973
 974   llvm::Function *ModuleDtorFunc = llvm::Function::Create(
 975       llvm::FunctionType::get(VoidTy, false),
 976       llvm::GlobalValue::InternalLinkage,
 977       addUnderscoredPrefixToName("_module_dtor"), &TheModule);
 978
 979   llvm::BasicBlock *DtorEntryBB =
 980       llvm::BasicBlock::Create(Context, "entry", ModuleDtorFunc);
 981   CGBuilderTy DtorBuilder(CGM, Context);
 982   DtorBuilder.SetInsertPoint(DtorEntryBB);
 983
 984   Address GpuBinaryAddr(
 985       GpuBinaryHandle, GpuBinaryHandle->getValueType(),
 986       CharUnits::fromQuantity(GpuBinaryHandle->getAlignment()));
 987   auto *HandleValue = DtorBuilder.CreateLoad(GpuBinaryAddr);
 988   // There is only one HIP fat binary per linked module, however there are
 989   // multiple destructor functions. Make sure the fat binary is unregistered
 990   // only once.
 991   if (CGM.getLangOpts().HIP) {
 992     llvm::BasicBlock *IfBlock =
 993         llvm::BasicBlock::Create(Context, "if", ModuleDtorFunc);
 994     llvm::BasicBlock *ExitBlock =
 995         llvm::BasicBlock::Create(Context, "exit", ModuleDtorFunc);
 996     llvm::Constant *Zero = llvm::Constant::getNullValue(HandleValue->getType());
 997     llvm::Value *NEZero = DtorBuilder.CreateICmpNE(HandleValue, Zero);
 998     DtorBuilder.CreateCondBr(NEZero, IfBlock, ExitBlock);
 999
1000     DtorBuilder.SetInsertPoint(IfBlock);
1001     DtorBuilder.CreateCall(UnregisterFatbinFunc, HandleValue);
1002     DtorBuilder.CreateStore(Zero, GpuBinaryAddr);
1003     DtorBuilder.CreateBr(ExitBlock);
1004
1005     DtorBuilder.SetInsertPoint(ExitBlock);
1006   } else {
1007     DtorBuilder.CreateCall(UnregisterFatbinFunc, HandleValue);
1008   }
1009   DtorBuilder.CreateRetVoid();
1010   return ModuleDtorFunc;
1011 }
1012
1013 CGCUDARuntime *CodeGen::CreateNVCUDARuntime(CodeGenModule &CGM) {
1014   return new CGNVCUDARuntime(CGM);
1015 }
1016
1017 void CGNVCUDARuntime::internalizeDeviceSideVar(
1018     const VarDecl *D, llvm::GlobalValue::LinkageTypes &Linkage) {
1019   // For -fno-gpu-rdc, host-side shadows of external declarations of device-side
1020   // global variables become internal definitions. These have to be internal in
1021   // order to prevent name conflicts with global host variables with the same
1022   // name in a different TUs.
1023   //
1024   // For -fgpu-rdc, the shadow variables should not be internalized because
1025   // they may be accessed by different TU.
1026   if (CGM.getLangOpts().GPURelocatableDeviceCode)
1027     return;
1028
1029   // __shared__ variables are odd. Shadows do get created, but
1030   // they are not registered with the CUDA runtime, so they
1031   // can't really be used to access their device-side
1032   // counterparts. It's not clear yet whether it's nvcc's bug or
1033   // a feature, but we've got to do the same for compatibility.
1034   if (D->hasAttr<CUDADeviceAttr>() || D->hasAttr<CUDAConstantAttr>() ||
1035       D->hasAttr<CUDASharedAttr>() ||
1036       D->getType()->isCUDADeviceBuiltinSurfaceType() ||
1037       D->getType()->isCUDADeviceBuiltinTextureType()) {
1038     Linkage = llvm::GlobalValue::InternalLinkage;
1039   }
1040 }
1041
1042 void CGNVCUDARuntime::handleVarRegistration(const VarDecl *D,
1043                                             llvm::GlobalVariable &GV) {
1044   if (D->hasAttr<CUDADeviceAttr>() || D->hasAttr<CUDAConstantAttr>()) {
1045     // Shadow variables and their properties must be registered with CUDA
1046     // runtime. Skip Extern global variables, which will be registered in
1047     // the TU where they are defined.
1048     //
1049     // Don't register a C++17 inline variable. The local symbol can be
1050     // discarded and referencing a discarded local symbol from outside the
1051     // comdat (__cuda_register_globals) is disallowed by the ELF spec.
1052     //
1053     // HIP managed variables need to be always recorded in device and host
1054     // compilations for transformation.
1055     //
1056     // HIP managed variables and variables in CUDADeviceVarODRUsedByHost are
1057     // added to llvm.compiler-used, therefore they are safe to be registered.
1058     if ((!D->hasExternalStorage() && !D->isInline()) ||
1059         CGM.getContext().CUDADeviceVarODRUsedByHost.contains(D) ||
1060         D->hasAttr<HIPManagedAttr>()) {
1061       registerDeviceVar(D, GV, !D->hasDefinition(),
1062                         D->hasAttr<CUDAConstantAttr>());
1063     }
1064   } else if (D->getType()->isCUDADeviceBuiltinSurfaceType() ||
1065              D->getType()->isCUDADeviceBuiltinTextureType()) {
1066     // Builtin surfaces and textures and their template arguments are
1067     // also registered with CUDA runtime.
1068     const auto *TD = cast<ClassTemplateSpecializationDecl>(
1069         D->getType()->castAs<RecordType>()->getDecl());
1070     const TemplateArgumentList &Args = TD->getTemplateArgs();
1071     if (TD->hasAttr<CUDADeviceBuiltinSurfaceTypeAttr>()) {
1072       assert(Args.size() == 2 &&
1073              "Unexpected number of template arguments of CUDA device "
1074              "builtin surface type.");
1075       auto SurfType = Args[1].getAsIntegral();
1076       if (!D->hasExternalStorage())
1077         registerDeviceSurf(D, GV, !D->hasDefinition(), SurfType.getSExtValue());
1078     } else {
1079       assert(Args.size() == 3 &&
1080              "Unexpected number of template arguments of CUDA device "
1081              "builtin texture type.");
1082       auto TexType = Args[1].getAsIntegral();
1083       auto Normalized = Args[2].getAsIntegral();
1084       if (!D->hasExternalStorage())
1085         registerDeviceTex(D, GV, !D->hasDefinition(), TexType.getSExtValue(),
1086                           Normalized.getZExtValue());
1087     }
1088   }
1089 }
1090
1091 // Transform managed variables to pointers to managed variables in device code.
1092 // Each use of the original managed variable is replaced by a load from the
1093 // transformed managed variable. The transformed managed variable contains
1094 // the address of managed memory which will be allocated by the runtime.
1095 void CGNVCUDARuntime::transformManagedVars() {
1096   for (auto &&Info : DeviceVars) {
1097     llvm::GlobalVariable *Var = Info.Var;
1098     if (Info.Flags.getKind() == DeviceVarFlags::Variable &&
1099         Info.Flags.isManaged()) {
1100       auto *ManagedVar = new llvm::GlobalVariable(
1101           CGM.getModule(), Var->getType(),
1102           /*isConstant=*/false, Var->getLinkage(),
1103           /*Init=*/Var->isDeclaration()
1104               ? nullptr
1105               : llvm::ConstantPointerNull::get(Var->getType()),
1106           /*Name=*/"", /*InsertBefore=*/nullptr,
1107           llvm::GlobalVariable::NotThreadLocal,
1108           CGM.getContext().getTargetAddressSpace(LangAS::cuda_device));
1109       ManagedVar->setDSOLocal(Var->isDSOLocal());
1110       ManagedVar->setVisibility(Var->getVisibility());
1111       ManagedVar->setExternallyInitialized(true);
1112       replaceManagedVar(Var, ManagedVar);
1113       ManagedVar->takeName(Var);
1114       Var->setName(Twine(ManagedVar->getName()) + ".managed");
1115       // Keep managed variables even if they are not used in device code since
1116       // they need to be allocated by the runtime.
1117       if (!Var->isDeclaration()) {
1118         assert(!ManagedVar->isDeclaration());
1119         CGM.addCompilerUsedGlobal(Var);
1120         CGM.addCompilerUsedGlobal(ManagedVar);
1121       }
1122     }
1123   }
1124 }
1125
1126 // Creates offloading entries for all the kernels and globals that must be
1127 // registered. The linker will provide a pointer to this section so we can
1128 // register the symbols with the linked device image.
1129 void CGNVCUDARuntime::createOffloadingEntries() {
1130   llvm::OpenMPIRBuilder OMPBuilder(CGM.getModule());
1131   OMPBuilder.initialize();
1132
1133   StringRef Section = CGM.getLangOpts().HIP ? "hip_offloading_entries"
1134                                             : "cuda_offloading_entries";
1135   for (KernelInfo &I : EmittedKernels)
1136     OMPBuilder.emitOffloadingEntry(KernelHandles[I.Kernel->getName()],
1137                                    getDeviceSideName(cast<NamedDecl>(I.D)), 0,
1138                                    DeviceVarFlags::OffloadGlobalEntry, Section);
1139
1140   for (VarInfo &I : DeviceVars) {
1141     uint64_t VarSize =
1142         CGM.getDataLayout().getTypeAllocSize(I.Var->getValueType());
1143     if (I.Flags.getKind() == DeviceVarFlags::Variable) {
1144       OMPBuilder.emitOffloadingEntry(
1145           I.Var, getDeviceSideName(I.D), VarSize,
1146           I.Flags.isManaged() ? DeviceVarFlags::OffloadGlobalManagedEntry
1147                               : DeviceVarFlags::OffloadGlobalEntry,
1148           Section);
1149     } else if (I.Flags.getKind() == DeviceVarFlags::Surface) {
1150       OMPBuilder.emitOffloadingEntry(I.Var, getDeviceSideName(I.D), VarSize,
1151                                      DeviceVarFlags::OffloadGlobalSurfaceEntry,
1152                                      Section);
1153     } else if (I.Flags.getKind() == DeviceVarFlags::Texture) {
1154       OMPBuilder.emitOffloadingEntry(I.Var, getDeviceSideName(I.D), VarSize,
1155                                      DeviceVarFlags::OffloadGlobalTextureEntry,
1156                                      Section);
1157     }
1158   }
1159 }
1160
1161 // Returns module constructor to be added.
1162 llvm::Function *CGNVCUDARuntime::finalizeModule() {
1163   if (CGM.getLangOpts().CUDAIsDevice) {
1164     transformManagedVars();
1165
1166     // Mark ODR-used device variables as compiler used to prevent it from being
1167     // eliminated by optimization. This is necessary for device variables
1168     // ODR-used by host functions. Sema correctly marks them as ODR-used no
1169     // matter whether they are ODR-used by device or host functions.
1170     //
1171     // We do not need to do this if the variable has used attribute since it
1172     // has already been added.
1173     //
1174     // Static device variables have been externalized at this point, therefore
1175     // variables with LLVM private or internal linkage need not be added.
1176     for (auto &&Info : DeviceVars) {
1177       auto Kind = Info.Flags.getKind();
1178       if (!Info.Var->isDeclaration() &&
1179           !llvm::GlobalValue::isLocalLinkage(Info.Var->getLinkage()) &&
1180           (Kind == DeviceVarFlags::Variable ||
1181            Kind == DeviceVarFlags::Surface ||
1182            Kind == DeviceVarFlags::Texture) &&
1183           Info.D->isUsed() && !Info.D->hasAttr<UsedAttr>()) {
1184         CGM.addCompilerUsedGlobal(Info.Var);
1185       }
1186     }
1187     return nullptr;
1188   }
1189   if (CGM.getLangOpts().OffloadingNewDriver && RelocatableDeviceCode)
1190     createOffloadingEntries();
1191   else
1192     return makeModuleCtorFunction();
1193
1194   return nullptr;
1195 }
1196
1197 llvm::GlobalValue *CGNVCUDARuntime::getKernelHandle(llvm::Function *F,
1198                                                     GlobalDecl GD) {
1199   auto Loc = KernelHandles.find(F->getName());
1200   if (Loc != KernelHandles.end()) {
1201     auto OldHandle = Loc->second;
1202     if (KernelStubs[OldHandle] == F)
1203       return OldHandle;
1204
1205     // We've found the function name, but F itself has changed, so we need to
1206     // update the references.
1207     if (CGM.getLangOpts().HIP) {
1208       // For HIP compilation the handle itself does not change, so we only need
1209       // to update the Stub value.
1210       KernelStubs[OldHandle] = F;
1211       return OldHandle;
1212     }
1213     // For non-HIP compilation, erase the old Stub and fall-through to creating
1214     // new entries.
1215     KernelStubs.erase(OldHandle);
1216   }
1217
1218   if (!CGM.getLangOpts().HIP) {
1219     KernelHandles[F->getName()] = F;
1220     KernelStubs[F] = F;
1221     return F;
1222   }
1223
1224   auto *Var = new llvm::GlobalVariable(
1225       TheModule, F->getType(), /*isConstant=*/true, F->getLinkage(),
1226       /*Initializer=*/nullptr,
1227       CGM.getMangledName(
1228           GD.getWithKernelReferenceKind(KernelReferenceKind::Kernel)));
1229   Var->setAlignment(CGM.getPointerAlign().getAsAlign());
1230   Var->setDSOLocal(F->isDSOLocal());
1231   Var->setVisibility(F->getVisibility());
1232   CGM.maybeSetTrivialComdat(*GD.getDecl(), *Var);
1233   KernelHandles[F->getName()] = Var;
1234   KernelStubs[Var] = F;
1235   return Var;
1236 }