[CUDA] Add wrapper code generation for registering CUDA images

[jhuber6]

Summary

This patch adds the necessary code generation to create the wrapper code
that registers all the globals in CUDA. We create the necessary
functions and iterate through the list of
__start_cuda_offloading_entries to find which globals must be
registered. This is very similar to the code generation done currently
in Clang for non-rdc builds, but here we are registering a fully linked
fatbinary and finding the globals via the above sections.

With this we should be able to fully support basic RDC / LTO building of CUDA
code.

It's also worth noting that this does not include the necessary PTX to JIT the
image, so to use this support the offloading architecture must match the
system's architecture.

Depends on D123810

[jhuber6]

For reference, this is how I used this driver to offload a CUDA and OpenMP offload kernel that both called an external "hello world" function in device.cu.

$ clang++ hello.cu device.cu -foffload-new-driver --offload-arch=sm_70 -c
$ clang++ openmp.cpp -fopenmp-new-driver -fopenmp -fopenmp-targets=nvptx64 -Xopenmp-target=nvptx64 -march=sm_70 -fopenmp-offload-mandatory -c
$ clang++ hello.o device.o openmp.o -foffload-new-driver -fopenmp -fopenmp-targets=nvptx64

[tra]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

Do you think generation of the CUDA registration glue could be shared with the front-end?

[jhuber6]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

I was thinking about it, but ultimately decided to keep the noise outside of the new driver to a minimum. Maybe if we move to the offloading entries being a common format we can easily share this code. Keeping it in Clang would have the advantage that it's easier to test directly and ensures we don't de-sync if anything changes. The only downside is that in the future I may want to push this functionality to a linker plugin or similar, which would require pulling it out of Clang again to prevent us needing to link in Clang to build LLVM.

Also needing to do this all through the builder API isn't ideal, it would be nice if we had some kind of runtime to call to do this for us, but I didn't feel like adding yet another shared library for CUDA. I considered putting it inside the cuda header wrappers as well, but forcing every CUDA file to have some externally visible weak registration blob didn't sit well with me.

[tra]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

Perhaps front-end is not the right place for it, indeed. LLVM itself may be a better choice. We already have some things there for somewhat similar purposes (like lib/WindowsManifest) so adding a helper function to generate runtime glue for CUDA should not be unreasonable.

[jhuber6]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

I think it's fine here for this patch, but I definitely want to move it into LLVM in the future once I start generalizing more of this stuff.

[yaxunl]

need a CodeGenCUDA test for the registering. Also need a Driver test for the subcommands.

[yaxunl]

Also, I am wondering whether we should document the new embedding scheme: section names, symbol names, entries, etc, if it has not bee done.

[jhuber6]

In D123812#3459164, @yaxunl wrote:

need a CodeGenCUDA test for the registering. Also need a Driver test for the subcommands.

Testing things inside the linker wrapper is a little hairy. I may need to add a special option for doing dry runs and printing the wrapping code to we can test this more satisfactorily. Doing that will create some more noise for review unfortunately. For OpenMP is simply ran all of our unit tests using the new driver and considered that sufficient evidence that it was working. Also which driver sub-commands should be tested?

In D123812#3459172, @yaxunl wrote:

Also, I am wondering whether we should document the new embedding scheme: section names, symbol names, entries, etc, if it has not bee done.

There's some existing documentation I wrote for OpenMP offloading once we started using this scheme. I was planning on updating it with this scheme for CUDA once it's landed, no sense writing documentation before it's final.

[jhuber6]

Addings tests for wrapper codegen and fatbinary usage.

[yaxunl]

clang/test/Driver/linker-wrapper-image.c

32 ↗

(On Diff #423796)

what happens if there are multiple binaries for different GPUs? will the linker-wrapper generates one fatbinary containing both elfs and embed the fatbinary as one image?

clang/test/Driver/linker-wrapper.c

46 ↗

(On Diff #423796)

This option is the same as the preceding option. Is this intentional? Can we have a test that embeds multiple binaries for different GPUs?

[jhuber6]

clang/test/Driver/linker-wrapper-image.c

32 ↗

(On Diff #423796)

Yes, I'll add it to the other test.

clang/test/Driver/linker-wrapper.c

46 ↗

(On Diff #423796)

It's intentional to show that we can pull out two objects embedded in a single file (Like if someone did ld -r or something). I'll add binaries for different GPUs to show that works.

[jhuber6]

Adjusting tests.

[jhuber6]

This comment was removed by jhuber6.

[jhuber6]

Applied changes to wrong commit, whoops.

[jhuber6]

Updating code generation. Previously we would seg-fault in the case that no offloading entries were created. To solve this we simply check that Begin != End before trying to register anything.

[jhuber6]

ping

[jdoerfert]

Are there unresolved concerns we should address in this review?

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

I'm OK with it being here but the place to consider (IMHO) is llvm/lib/Frontend, maybe /CUDA/Register.cpp.

[yaxunl]

LGTM.

[jhuber6]

In D123812#3496914, @yaxunl wrote:

LGTM.

Did you forget to accept the revision? D123810 and D123471 still need to be looked at, but these are mostly non-intrusive so I don't think they'll break anything.

[tra]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

OK. I'm fine keeping it all here for now.
Please add a comment pointing towards the origin of this code. and, maybe, a TODO item to consolidate and move it into a better place.

[jhuber6]

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

351

Will do, thanks for the reviews. I'll land these tomorrow morning and see if anything breaks.

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

int main(int argc, const char **argv) {
   if (Error Err = linkDeviceFiles(DeviceFiles, LinkedImages))
     return reportError(std::move(Err));
 
   // Wrap each linked device image into a linkable host binary and add it to the
   // link job's inputs.
   auto FileOrErr = wrapDeviceImages(LinkedImages);
   if (!FileOrErr)
     return reportError(FileOrErr.takeError());
-  LinkerArgs.append(*FileOrErr);
+
+  // We need to insert the new files next to the old ones to make sure they're
+  // linked with the same libraries / arguments.
+  auto FirstInput = std::next(llvm::find_if(LinkerArgs, [](StringRef Str) {
+    return sys::fs::exists(Str) && !sys::fs::is_directory(Str) &&
+           Str != ExecutableName;
+  }));
+  LinkerArgs.insert(FirstInput, FileOrErr->begin(), FileOrErr->end());
 
   // Run the host linking job.
   if (Error Err = runLinker(LinkerUserPath, LinkerArgs))
     return reportError(std::move(Err));
 
   // Remove the temporary files created.
   for (const auto &TempFile : TempFiles)
     if (std::error_code EC = sys::fs::remove(TempFile))
       reportError(createFileError(TempFile, EC));
 
   return EXIT_SUCCESS;
 }

clang/tools/clang-linker-wrapper/OffloadWrapper.cpp

 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
 
 namespace {
+/// Magic number that begins the section containing the CUDA fatbinary.
+constexpr unsigned CudaFatMagic = 0x466243b1;
 
 IntegerType *getSizeTTy(Module &M) {
   LLVMContext &C = M.getContext();
   switch (M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))) {
   case 4u:
     return Type::getInt32Ty(C);
   case 8u:
     return Type::getInt64Ty(C);
void createUnregisterFunction(Module &M, GlobalVariable *BinDesc) {
   Builder.CreateCall(UnRegFuncC, BinDesc);
   Builder.CreateRetVoid();
 
   // Add this function to global destructors.
   // Match priority of __tgt_register_lib
   appendToGlobalDtors(M, Func, /*Priority*/ 1);
 }
 
+// struct fatbin_wrapper {
+//  int32_t magic;
+//  int32_t version;
+//  void *image;
+//  void *reserved;
+//};
+StructType *getFatbinWrapperTy(Module &M) {
+  LLVMContext &C = M.getContext();
+  StructType *FatbinTy = StructType::getTypeByName(C, "fatbin_wrapper");
+  if (!FatbinTy)
+    FatbinTy = StructType::create("fatbin_wrapper", Type::getInt32Ty(C),
+                                  Type::getInt32Ty(C), Type::getInt8PtrTy(C),
+                                  Type::getInt8PtrTy(C));
+  return FatbinTy;
+}
+
+/// Embed the image \p Image into the module \p M so it can be found by the
+/// runtime.
+GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image) {
+  LLVMContext &C = M.getContext();
+  llvm::Type *Int8PtrTy = Type::getInt8PtrTy(C);
+  llvm::Triple Triple = llvm::Triple(M.getTargetTriple());
+
+  // Create the global string containing the fatbinary.
+  StringRef FatbinConstantSection =
+      Triple.isMacOSX() ? "__NV_CUDA,__nv_fatbin" : ".nv_fatbin";
+  auto *Data = ConstantDataArray::get(C, Image);
+  auto *Fatbin = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
+                                    GlobalVariable::InternalLinkage, Data,
+                                    ".fatbin_image");
+  Fatbin->setSection(FatbinConstantSection);
+
+  // Create the fatbinary wrapper
+  StringRef FatbinWrapperSection =
+      Triple.isMacOSX() ? "__NV_CUDA,__fatbin" : ".nvFatBinSegment";
+  Constant *FatbinWrapper[] = {
+      ConstantInt::get(Type::getInt32Ty(C), CudaFatMagic),
+      ConstantInt::get(Type::getInt32Ty(C), 1),
+      ConstantExpr::getPointerBitCastOrAddrSpaceCast(Fatbin, Int8PtrTy),
+      ConstantPointerNull::get(Type::getInt8PtrTy(C))};
+
+  Constant *FatbinInitializer =
+      ConstantStruct::get(getFatbinWrapperTy(M), FatbinWrapper);
+
+  auto *FatbinDesc =
+      new GlobalVariable(M, getFatbinWrapperTy(M),
+                         /*isConstant*/ true, GlobalValue::InternalLinkage,
+                         FatbinInitializer, ".fatbin_wrapper");
+  FatbinDesc->setSection(FatbinWrapperSection);
+  FatbinDesc->setAlignment(Align(8));
+
+  // We create a dummy entry to ensure the linker will define the begin / end
+  // symbols. The CUDA runtime should ignore the null address if we attempt to
+  // register it.
+  auto *DummyInit =
+      ConstantAggregateZero::get(ArrayType::get(getEntryTy(M), 0u));
+  auto *DummyEntry = new GlobalVariable(
+      M, DummyInit->getType(), true, GlobalVariable::ExternalLinkage, DummyInit,
+      "__dummy.cuda_offloading.entry");
+  DummyEntry->setSection("cuda_offloading_entries");
+  DummyEntry->setVisibility(GlobalValue::HiddenVisibility);
+
+  return FatbinDesc;
+}
+
+/// Create the register globals function. We will iterate all of the offloading
+/// entries stored at the begin / end symbols and register them according to
+/// their type. This creates the following function in IR:
+///
+/// extern struct __tgt_offload_entry __start_cuda_offloading_entries;
+/// extern struct __tgt_offload_entry __stop_cuda_offloading_entries;
+///
+/// extern void __cudaRegisterFunction(void **, void *, void *, void *, int,
+///                                    void *, void *, void *, void *, int *);
+/// extern void __cudaRegisterVar(void **, void *, void *, void *, int32_t,
+///                               int64_t, int32_t, int32_t);
+///
+/// void __cudaRegisterTest(void **fatbinHandle) {
+///   for (struct __tgt_offload_entry *entry = &__start_cuda_offloading_entries;
+///        entry != &__stop_cuda_offloading_entries; ++entry) {
+///     if (!entry->size)
+///       __cudaRegisterFunction(fatbinHandle, entry->addr, entry->name,
+///                              entry->name, -1, 0, 0, 0, 0, 0);
+///     else
+///       __cudaRegisterVar(fatbinHandle, entry->addr, entry->name, entry->name,
+///                         0, entry->size, 0, 0);
+///   }
+/// }
+///
+/// TODO: This only registers functions are variables. Additional support is
+///       required for texture / surface / managed variables.
+Function *createRegisterGlobalsFunction(Module &M) {
+  LLVMContext &C = M.getContext();
+  // Get the __cudaRegisterFunction function declaration.
+  auto *RegFuncTy = FunctionType::get(
+      Type::getInt32Ty(C),
+      {Type::getInt8PtrTy(C)->getPointerTo(), Type::getInt8PtrTy(C),
+       Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt32Ty(C),
+       Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt8PtrTy(C),
+       Type::getInt8PtrTy(C), Type::getInt32PtrTy(C)},
+      /*isVarArg*/ false);
+  FunctionCallee RegFunc =
+      M.getOrInsertFunction("__cudaRegisterFunction", RegFuncTy);
+
+  // Get the __cudaRegisterVar function declaration.
+  auto *RegVarTy = FunctionType::get(
+      Type::getInt32Ty(C),
+      {Type::getInt8PtrTy(C)->getPointerTo(), Type::getInt8PtrTy(C),
+       Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt32Ty(C),
+       getSizeTTy(M), Type::getInt32Ty(C), Type::getInt32Ty(C)},
+      /*isVarArg*/ false);
+  FunctionCallee RegVar = M.getOrInsertFunction("__cudaRegisterVar", RegVarTy);
+
+  // Create the references to the start / stop symbols defined by the linker.
+  auto *EntriesB = new GlobalVariable(
+      M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
+      /*Initializer*/ nullptr, "__start_cuda_offloading_entries");
+  EntriesB->setVisibility(GlobalValue::HiddenVisibility);
+  auto *EntriesE = new GlobalVariable(
+      M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
+      /*Initializer*/ nullptr, "__stop_cuda_offloading_entries");
+  EntriesE->setVisibility(GlobalValue::HiddenVisibility);
+
+  auto *RegGlobalsTy = FunctionType::get(Type::getVoidTy(C),
+                                         Type::getInt8PtrTy(C)->getPointerTo(),
+                                         /*isVarArg*/ false);
+  auto *RegGlobalsFn = Function::Create(
+      RegGlobalsTy, GlobalValue::InternalLinkage, ".cuda.globals_reg", &M);
+  RegGlobalsFn->setSection(".text.startup");
+
+  // Create the loop to register all the entries.
+  IRBuilder<> Builder(BasicBlock::Create(C, "entry", RegGlobalsFn));
+  auto *EntryBB = BasicBlock::Create(C, "while.entry", RegGlobalsFn);
+  auto *IfThenBB = BasicBlock::Create(C, "if.then", RegGlobalsFn);
+  auto *IfElseBB = BasicBlock::Create(C, "if.else", RegGlobalsFn);
+  auto *IfEndBB = BasicBlock::Create(C, "if.end", RegGlobalsFn);
+  auto *ExitBB = BasicBlock::Create(C, "while.end", RegGlobalsFn);
+
+  Builder.CreateBr(EntryBB);
+  Builder.SetInsertPoint(EntryBB);
+  auto *Entry = Builder.CreatePHI(getEntryPtrTy(M), 2, "entry");
+  auto *AddrPtr =
+      Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
+                                {ConstantInt::get(getSizeTTy(M), 0),
+                                 ConstantInt::get(Type::getInt32Ty(C), 0)});
+  auto *Addr = Builder.CreateLoad(Type::getInt8PtrTy(C), AddrPtr, "addr");
+  auto *NamePtr =
+      Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
+                                {ConstantInt::get(getSizeTTy(M), 0),
+                                 ConstantInt::get(Type::getInt32Ty(C), 1)});
+  auto *Name = Builder.CreateLoad(Type::getInt8PtrTy(C), NamePtr, "name");
+  auto *SizePtr =
+      Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
+                                {ConstantInt::get(getSizeTTy(M), 0),
+                                 ConstantInt::get(Type::getInt32Ty(C), 2)});
+  auto *Size = Builder.CreateLoad(getSizeTTy(M), SizePtr, "size");
+  auto *FnCond =
+      Builder.CreateICmpEQ(Size, ConstantInt::getNullValue(getSizeTTy(M)));
+  Builder.CreateCondBr(FnCond, IfThenBB, IfElseBB);
+  Builder.SetInsertPoint(IfThenBB);
+  Builder.CreateCall(RegFunc,
+                     {RegGlobalsFn->arg_begin(), Addr, Name, Name,
+                      ConstantInt::get(Type::getInt32Ty(C), -1),
+                      ConstantPointerNull::get(Type::getInt8PtrTy(C)),
+                      ConstantPointerNull::get(Type::getInt8PtrTy(C)),
+                      ConstantPointerNull::get(Type::getInt8PtrTy(C)),
+                      ConstantPointerNull::get(Type::getInt8PtrTy(C)),
+                      ConstantPointerNull::get(Type::getInt32PtrTy(C))});
+  Builder.CreateBr(IfEndBB);
+  Builder.SetInsertPoint(IfElseBB);
+  Builder.CreateCall(RegVar, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
+                              ConstantInt::get(Type::getInt32Ty(C), 0), Size,
+                              ConstantInt::get(Type::getInt32Ty(C), 0),
+                              ConstantInt::get(Type::getInt32Ty(C), 0)});
+  Builder.CreateBr(IfEndBB);
+  Builder.SetInsertPoint(IfEndBB);
+  auto *NewEntry = Builder.CreateInBoundsGEP(
+      getEntryTy(M), Entry, ConstantInt::get(getSizeTTy(M), 1));
+  auto *Cmp = Builder.CreateICmpEQ(NewEntry, EntriesE);
+  Entry->addIncoming(EntriesB, &RegGlobalsFn->getEntryBlock());
+  Entry->addIncoming(NewEntry, IfEndBB);
+  Builder.CreateCondBr(Cmp, ExitBB, EntryBB);
+  Builder.SetInsertPoint(ExitBB);
+  Builder.CreateRetVoid();
+
+  return RegGlobalsFn;
+}
+
+// Create the constructor and destructor to register the fatbinary with the CUDA
+// runtime.
+void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc) {
+  LLVMContext &C = M.getContext();
+  auto *CtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
+  auto *CtorFunc = Function::Create(CtorFuncTy, GlobalValue::InternalLinkage,
+                                    ".cuda.fatbin_reg", &M);
+  CtorFunc->setSection(".text.startup");
+
+  auto *DtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
+  auto *DtorFunc = Function::Create(DtorFuncTy, GlobalValue::InternalLinkage,
+                                    ".cuda.fatbin_unreg", &M);
+  DtorFunc->setSection(".text.startup");
+
+  // Get the __cudaRegisterFatBinary function declaration.
+  auto *RegFatTy = FunctionType::get(Type::getInt8PtrTy(C)->getPointerTo(),
+                                     Type::getInt8PtrTy(C),
+                                     /*isVarArg*/ false);
+  FunctionCallee RegFatbin =
+      M.getOrInsertFunction("__cudaRegisterFatBinary", RegFatTy);
+  // Get the __cudaRegisterFatBinaryEnd function declaration.
+  auto *RegFatEndTy = FunctionType::get(Type::getVoidTy(C),
+                                        Type::getInt8PtrTy(C)->getPointerTo(),
+                                        /*isVarArg*/ false);
+  FunctionCallee RegFatbinEnd =
+      M.getOrInsertFunction("__cudaRegisterFatBinaryEnd", RegFatEndTy);
+  // Get the __cudaUnregisterFatBinary function declaration.
+  auto *UnregFatTy = FunctionType::get(Type::getVoidTy(C),
+                                       Type::getInt8PtrTy(C)->getPointerTo(),
+                                       /*isVarArg*/ false);
+  FunctionCallee UnregFatbin =
+      M.getOrInsertFunction("__cudaUnregisterFatBinary", UnregFatTy);
+
+  auto *AtExitTy =
+      FunctionType::get(Type::getInt32Ty(C), DtorFuncTy->getPointerTo(),
+                        /*isVarArg*/ false);
+  FunctionCallee AtExit = M.getOrInsertFunction("atexit", AtExitTy);
+
+  auto *BinaryHandleGlobal = new llvm::GlobalVariable(
+      M, Type::getInt8PtrTy(C)->getPointerTo(), false,
+      llvm::GlobalValue::InternalLinkage,
+      llvm::ConstantPointerNull::get(Type::getInt8PtrTy(C)->getPointerTo()),
+      ".cuda.binary_handle");
+
+  // Create the constructor to register this image with the runtime.
+  IRBuilder<> CtorBuilder(BasicBlock::Create(C, "entry", CtorFunc));
+  CallInst *Handle = CtorBuilder.CreateCall(
+      RegFatbin, ConstantExpr::getPointerBitCastOrAddrSpaceCast(
+                     FatbinDesc, Type::getInt8PtrTy(C)));
+  CtorBuilder.CreateAlignedStore(
+      Handle, BinaryHandleGlobal,
+      Align(M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))));
+  CtorBuilder.CreateCall(createRegisterGlobalsFunction(M), Handle);
+  CtorBuilder.CreateCall(RegFatbinEnd, Handle);
+  CtorBuilder.CreateCall(AtExit, DtorFunc);
+  CtorBuilder.CreateRetVoid();
+
+  // Create the destructor to unregister the image with the runtime. We cannot
+  // use a standard global destructor after CUDA 9.2 so this must be called by
+  // `atexit()` intead.
+  IRBuilder<> DtorBuilder(BasicBlock::Create(C, "entry", DtorFunc));
+  LoadInst *BinaryHandle = DtorBuilder.CreateAlignedLoad(
+      Type::getInt8PtrTy(C)->getPointerTo(), BinaryHandleGlobal,
+      Align(M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))));
+  DtorBuilder.CreateCall(UnregFatbin, BinaryHandle);
+  DtorBuilder.CreateRetVoid();
+
+  // Add this function to constructors.
+  appendToGlobalCtors(M, CtorFunc, /*Priority*/ 1);
+}
+
 } // namespace
 
 Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
   GlobalVariable *Desc = createBinDesc(M, Images);
   if (!Desc)
     return createStringError(inconvertibleErrorCode(),
                              "No binary descriptors created.");
   createRegisterFunction(M, Desc);
   createUnregisterFunction(M, Desc);
   return Error::success();
 }
 
-llvm::Error wrapCudaBinary(llvm::Module &M, llvm::ArrayRef<char> Images) {
+Error wrapCudaBinary(Module &M, ArrayRef<char> Image) {
+  GlobalVariable *Desc = createFatbinDesc(M, Image);
+  if (!Desc)
     return createStringError(inconvertibleErrorCode(),
-                           "Cuda wrapping is not yet supported.");
+                             "No fatinbary section created.");
+
+  createRegisterFatbinFunction(M, Desc);
+  return Error::success();
 }