1 //===- OffloadWrapper.cpp ---------------------------------------*- C++ -*-===//
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 #include "OffloadWrapper.h"
10 #include "llvm/ADT/ArrayRef.h"
11 #include "llvm/ADT/Triple.h"
12 #include "llvm/IR/Constants.h"
13 #include "llvm/IR/GlobalVariable.h"
14 #include "llvm/IR/IRBuilder.h"
15 #include "llvm/IR/LLVMContext.h"
16 #include "llvm/IR/Module.h"
17 #include "llvm/Support/Error.h"
18 #include "llvm/Transforms/Utils/ModuleUtils.h"
19 
20 using namespace llvm;
21 
22 namespace {
23 /// Magic number that begins the section containing the CUDA fatbinary.
24 constexpr unsigned CudaFatMagic = 0x466243b1;
25 
26 /// Copied from clang/CGCudaRuntime.h.
27 enum OffloadEntryKindFlag : uint32_t {
28   /// Mark the entry as a global entry. This indicates the presense of a
29   /// kernel if the size size field is zero and a variable otherwise.
30   OffloadGlobalEntry = 0x0,
31   /// Mark the entry as a managed global variable.
32   OffloadGlobalManagedEntry = 0x1,
33   /// Mark the entry as a surface variable.
34   OffloadGlobalSurfaceEntry = 0x2,
35   /// Mark the entry as a texture variable.
36   OffloadGlobalTextureEntry = 0x3,
37 };
38 
39 IntegerType *getSizeTTy(Module &M) {
40   LLVMContext &C = M.getContext();
41   switch (M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))) {
42   case 4u:
43     return Type::getInt32Ty(C);
44   case 8u:
45     return Type::getInt64Ty(C);
46   }
47   llvm_unreachable("unsupported pointer type size");
48 }
49 
50 // struct __tgt_offload_entry {
51 //   void *addr;
52 //   char *name;
53 //   size_t size;
54 //   int32_t flags;
55 //   int32_t reserved;
56 // };
57 StructType *getEntryTy(Module &M) {
58   LLVMContext &C = M.getContext();
59   StructType *EntryTy = StructType::getTypeByName(C, "__tgt_offload_entry");
60   if (!EntryTy)
61     EntryTy = StructType::create("__tgt_offload_entry", Type::getInt8PtrTy(C),
62                                  Type::getInt8PtrTy(C), getSizeTTy(M),
63                                  Type::getInt32Ty(C), Type::getInt32Ty(C));
64   return EntryTy;
65 }
66 
67 PointerType *getEntryPtrTy(Module &M) {
68   return PointerType::getUnqual(getEntryTy(M));
69 }
70 
71 // struct __tgt_device_image {
72 //   void *ImageStart;
73 //   void *ImageEnd;
74 //   __tgt_offload_entry *EntriesBegin;
75 //   __tgt_offload_entry *EntriesEnd;
76 // };
77 StructType *getDeviceImageTy(Module &M) {
78   LLVMContext &C = M.getContext();
79   StructType *ImageTy = StructType::getTypeByName(C, "__tgt_device_image");
80   if (!ImageTy)
81     ImageTy = StructType::create("__tgt_device_image", Type::getInt8PtrTy(C),
82                                  Type::getInt8PtrTy(C), getEntryPtrTy(M),
83                                  getEntryPtrTy(M));
84   return ImageTy;
85 }
86 
87 PointerType *getDeviceImagePtrTy(Module &M) {
88   return PointerType::getUnqual(getDeviceImageTy(M));
89 }
90 
91 // struct __tgt_bin_desc {
92 //   int32_t NumDeviceImages;
93 //   __tgt_device_image *DeviceImages;
94 //   __tgt_offload_entry *HostEntriesBegin;
95 //   __tgt_offload_entry *HostEntriesEnd;
96 // };
97 StructType *getBinDescTy(Module &M) {
98   LLVMContext &C = M.getContext();
99   StructType *DescTy = StructType::getTypeByName(C, "__tgt_bin_desc");
100   if (!DescTy)
101     DescTy = StructType::create("__tgt_bin_desc", Type::getInt32Ty(C),
102                                 getDeviceImagePtrTy(M), getEntryPtrTy(M),
103                                 getEntryPtrTy(M));
104   return DescTy;
105 }
106 
107 PointerType *getBinDescPtrTy(Module &M) {
108   return PointerType::getUnqual(getBinDescTy(M));
109 }
110 
111 /// Creates binary descriptor for the given device images. Binary descriptor
112 /// is an object that is passed to the offloading runtime at program startup
113 /// and it describes all device images available in the executable or shared
114 /// library. It is defined as follows
115 ///
116 /// __attribute__((visibility("hidden")))
117 /// extern __tgt_offload_entry *__start_omp_offloading_entries;
118 /// __attribute__((visibility("hidden")))
119 /// extern __tgt_offload_entry *__stop_omp_offloading_entries;
120 ///
121 /// static const char Image0[] = { <Bufs.front() contents> };
122 ///  ...
123 /// static const char ImageN[] = { <Bufs.back() contents> };
124 ///
125 /// static const __tgt_device_image Images[] = {
126 ///   {
127 ///     Image0,                            /*ImageStart*/
128 ///     Image0 + sizeof(Image0),           /*ImageEnd*/
129 ///     __start_omp_offloading_entries,    /*EntriesBegin*/
130 ///     __stop_omp_offloading_entries      /*EntriesEnd*/
131 ///   },
132 ///   ...
133 ///   {
134 ///     ImageN,                            /*ImageStart*/
135 ///     ImageN + sizeof(ImageN),           /*ImageEnd*/
136 ///     __start_omp_offloading_entries,    /*EntriesBegin*/
137 ///     __stop_omp_offloading_entries      /*EntriesEnd*/
138 ///   }
139 /// };
140 ///
141 /// static const __tgt_bin_desc BinDesc = {
142 ///   sizeof(Images) / sizeof(Images[0]),  /*NumDeviceImages*/
143 ///   Images,                              /*DeviceImages*/
144 ///   __start_omp_offloading_entries,      /*HostEntriesBegin*/
145 ///   __stop_omp_offloading_entries        /*HostEntriesEnd*/
146 /// };
147 ///
148 /// Global variable that represents BinDesc is returned.
149 GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs) {
150   LLVMContext &C = M.getContext();
151   // Create external begin/end symbols for the offload entries table.
152   auto *EntriesB = new GlobalVariable(
153       M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
154       /*Initializer*/ nullptr, "__start_omp_offloading_entries");
155   EntriesB->setVisibility(GlobalValue::HiddenVisibility);
156   auto *EntriesE = new GlobalVariable(
157       M, getEntryTy(M), /*isConstant*/ true, GlobalValue::ExternalLinkage,
158       /*Initializer*/ nullptr, "__stop_omp_offloading_entries");
159   EntriesE->setVisibility(GlobalValue::HiddenVisibility);
160 
161   // We assume that external begin/end symbols that we have created above will
162   // be defined by the linker. But linker will do that only if linker inputs
163   // have section with "omp_offloading_entries" name which is not guaranteed.
164   // So, we just create dummy zero sized object in the offload entries section
165   // to force linker to define those symbols.
166   auto *DummyInit =
167       ConstantAggregateZero::get(ArrayType::get(getEntryTy(M), 0u));
168   auto *DummyEntry = new GlobalVariable(
169       M, DummyInit->getType(), true, GlobalVariable::ExternalLinkage, DummyInit,
170       "__dummy.omp_offloading.entry");
171   DummyEntry->setSection("omp_offloading_entries");
172   DummyEntry->setVisibility(GlobalValue::HiddenVisibility);
173 
174   auto *Zero = ConstantInt::get(getSizeTTy(M), 0u);
175   Constant *ZeroZero[] = {Zero, Zero};
176 
177   // Create initializer for the images array.
178   SmallVector<Constant *, 4u> ImagesInits;
179   ImagesInits.reserve(Bufs.size());
180   for (ArrayRef<char> Buf : Bufs) {
181     auto *Data = ConstantDataArray::get(C, Buf);
182     auto *Image = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
183                                      GlobalVariable::InternalLinkage, Data,
184                                      ".omp_offloading.device_image");
185     Image->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
186 
187     auto *Size = ConstantInt::get(getSizeTTy(M), Buf.size());
188     Constant *ZeroSize[] = {Zero, Size};
189 
190     auto *ImageB =
191         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroZero);
192     auto *ImageE =
193         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroSize);
194 
195     ImagesInits.push_back(ConstantStruct::get(getDeviceImageTy(M), ImageB,
196                                               ImageE, EntriesB, EntriesE));
197   }
198 
199   // Then create images array.
200   auto *ImagesData = ConstantArray::get(
201       ArrayType::get(getDeviceImageTy(M), ImagesInits.size()), ImagesInits);
202 
203   auto *Images =
204       new GlobalVariable(M, ImagesData->getType(), /*isConstant*/ true,
205                          GlobalValue::InternalLinkage, ImagesData,
206                          ".omp_offloading.device_images");
207   Images->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
208 
209   auto *ImagesB =
210       ConstantExpr::getGetElementPtr(Images->getValueType(), Images, ZeroZero);
211 
212   // And finally create the binary descriptor object.
213   auto *DescInit = ConstantStruct::get(
214       getBinDescTy(M),
215       ConstantInt::get(Type::getInt32Ty(C), ImagesInits.size()), ImagesB,
216       EntriesB, EntriesE);
217 
218   return new GlobalVariable(M, DescInit->getType(), /*isConstant*/ true,
219                             GlobalValue::InternalLinkage, DescInit,
220                             ".omp_offloading.descriptor");
221 }
222 
223 void createRegisterFunction(Module &M, GlobalVariable *BinDesc) {
224   LLVMContext &C = M.getContext();
225   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
226   auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
227                                 ".omp_offloading.descriptor_reg", &M);
228   Func->setSection(".text.startup");
229 
230   // Get __tgt_register_lib function declaration.
231   auto *RegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(M),
232                                       /*isVarArg*/ false);
233   FunctionCallee RegFuncC =
234       M.getOrInsertFunction("__tgt_register_lib", RegFuncTy);
235 
236   // Construct function body
237   IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
238   Builder.CreateCall(RegFuncC, BinDesc);
239   Builder.CreateRetVoid();
240 
241   // Add this function to constructors.
242   // Set priority to 1 so that __tgt_register_lib is executed AFTER
243   // __tgt_register_requires (we want to know what requirements have been
244   // asked for before we load a libomptarget plugin so that by the time the
245   // plugin is loaded it can report how many devices there are which can
246   // satisfy these requirements).
247   appendToGlobalCtors(M, Func, /*Priority*/ 1);
248 }
249 
250 void createUnregisterFunction(Module &M, GlobalVariable *BinDesc) {
251   LLVMContext &C = M.getContext();
252   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
253   auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
254                                 ".omp_offloading.descriptor_unreg", &M);
255   Func->setSection(".text.startup");
256 
257   // Get __tgt_unregister_lib function declaration.
258   auto *UnRegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(M),
259                                         /*isVarArg*/ false);
260   FunctionCallee UnRegFuncC =
261       M.getOrInsertFunction("__tgt_unregister_lib", UnRegFuncTy);
262 
263   // Construct function body
264   IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
265   Builder.CreateCall(UnRegFuncC, BinDesc);
266   Builder.CreateRetVoid();
267 
268   // Add this function to global destructors.
269   // Match priority of __tgt_register_lib
270   appendToGlobalDtors(M, Func, /*Priority*/ 1);
271 }
272 
273 // struct fatbin_wrapper {
274 //  int32_t magic;
275 //  int32_t version;
276 //  void *image;
277 //  void *reserved;
278 //};
279 StructType *getFatbinWrapperTy(Module &M) {
280   LLVMContext &C = M.getContext();
281   StructType *FatbinTy = StructType::getTypeByName(C, "fatbin_wrapper");
282   if (!FatbinTy)
283     FatbinTy = StructType::create("fatbin_wrapper", Type::getInt32Ty(C),
284                                   Type::getInt32Ty(C), Type::getInt8PtrTy(C),
285                                   Type::getInt8PtrTy(C));
286   return FatbinTy;
287 }
288 
289 /// Embed the image \p Image into the module \p M so it can be found by the
290 /// runtime.
291 GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image) {
292   LLVMContext &C = M.getContext();
293   llvm::Type *Int8PtrTy = Type::getInt8PtrTy(C);
294   llvm::Triple Triple = llvm::Triple(M.getTargetTriple());
295 
296   // Create the global string containing the fatbinary.
297   StringRef FatbinConstantSection =
298       Triple.isMacOSX() ? "__NV_CUDA,__nv_fatbin" : ".nv_fatbin";
299   auto *Data = ConstantDataArray::get(C, Image);
300   auto *Fatbin = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
301                                     GlobalVariable::InternalLinkage, Data,
302                                     ".fatbin_image");
303   Fatbin->setSection(FatbinConstantSection);
304 
305   // Create the fatbinary wrapper
306   StringRef FatbinWrapperSection =
307       Triple.isMacOSX() ? "__NV_CUDA,__fatbin" : ".nvFatBinSegment";
308   Constant *FatbinWrapper[] = {
309       ConstantInt::get(Type::getInt32Ty(C), CudaFatMagic),
310       ConstantInt::get(Type::getInt32Ty(C), 1),
311       ConstantExpr::getPointerBitCastOrAddrSpaceCast(Fatbin, Int8PtrTy),
312       ConstantPointerNull::get(Type::getInt8PtrTy(C))};
313 
314   Constant *FatbinInitializer =
315       ConstantStruct::get(getFatbinWrapperTy(M), FatbinWrapper);
316 
317   auto *FatbinDesc =
318       new GlobalVariable(M, getFatbinWrapperTy(M),
319                          /*isConstant*/ true, GlobalValue::InternalLinkage,
320                          FatbinInitializer, ".fatbin_wrapper");
321   FatbinDesc->setSection(FatbinWrapperSection);
322   FatbinDesc->setAlignment(Align(8));
323 
324   // We create a dummy entry to ensure the linker will define the begin / end
325   // symbols. The CUDA runtime should ignore the null address if we attempt to
326   // register it.
327   auto *DummyInit =
328       ConstantAggregateZero::get(ArrayType::get(getEntryTy(M), 0u));
329   auto *DummyEntry = new GlobalVariable(
330       M, DummyInit->getType(), true, GlobalVariable::ExternalLinkage, DummyInit,
331       "__dummy.cuda_offloading.entry");
332   DummyEntry->setSection("cuda_offloading_entries");
333   DummyEntry->setVisibility(GlobalValue::HiddenVisibility);
334 
335   return FatbinDesc;
336 }
337 
338 /// Create the register globals function. We will iterate all of the offloading
339 /// entries stored at the begin / end symbols and register them according to
340 /// their type. This creates the following function in IR:
341 ///
342 /// extern struct __tgt_offload_entry __start_cuda_offloading_entries;
343 /// extern struct __tgt_offload_entry __stop_cuda_offloading_entries;
344 ///
345 /// extern void __cudaRegisterFunction(void **, void *, void *, void *, int,
346 ///                                    void *, void *, void *, void *, int *);
347 /// extern void __cudaRegisterVar(void **, void *, void *, void *, int32_t,
348 ///                               int64_t, int32_t, int32_t);
349 ///
350 /// void __cudaRegisterTest(void **fatbinHandle) {
351 ///   for (struct __tgt_offload_entry *entry = &__start_cuda_offloading_entries;
352 ///        entry != &__stop_cuda_offloading_entries; ++entry) {
353 ///     if (!entry->size)
354 ///       __cudaRegisterFunction(fatbinHandle, entry->addr, entry->name,
355 ///                              entry->name, -1, 0, 0, 0, 0, 0);
356 ///     else
357 ///       __cudaRegisterVar(fatbinHandle, entry->addr, entry->name, entry->name,
358 ///                         0, entry->size, 0, 0);
359 ///   }
360 /// }
361 Function *createRegisterGlobalsFunction(Module &M) {
362   LLVMContext &C = M.getContext();
363   // Get the __cudaRegisterFunction function declaration.
364   auto *RegFuncTy = FunctionType::get(
365       Type::getInt32Ty(C),
366       {Type::getInt8PtrTy(C)->getPointerTo(), Type::getInt8PtrTy(C),
367        Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt32Ty(C),
368        Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt8PtrTy(C),
369        Type::getInt8PtrTy(C), Type::getInt32PtrTy(C)},
370       /*isVarArg*/ false);
371   FunctionCallee RegFunc =
372       M.getOrInsertFunction("__cudaRegisterFunction", RegFuncTy);
373 
374   // Get the __cudaRegisterVar function declaration.
375   auto *RegVarTy = FunctionType::get(
376       Type::getVoidTy(C),
377       {Type::getInt8PtrTy(C)->getPointerTo(), Type::getInt8PtrTy(C),
378        Type::getInt8PtrTy(C), Type::getInt8PtrTy(C), Type::getInt32Ty(C),
379        getSizeTTy(M), Type::getInt32Ty(C), Type::getInt32Ty(C)},
380       /*isVarArg*/ false);
381   FunctionCallee RegVar = M.getOrInsertFunction("__cudaRegisterVar", RegVarTy);
382 
383   // Create the references to the start / stop symbols defined by the linker.
384   auto *EntriesB = new GlobalVariable(
385       M, ArrayType::get(getEntryTy(M), 0), /*isConstant*/ true,
386       GlobalValue::ExternalLinkage,
387       /*Initializer*/ nullptr, "__start_cuda_offloading_entries");
388   EntriesB->setVisibility(GlobalValue::HiddenVisibility);
389   auto *EntriesE = new GlobalVariable(
390       M, ArrayType::get(getEntryTy(M), 0), /*isConstant*/ true,
391       GlobalValue::ExternalLinkage,
392       /*Initializer*/ nullptr, "__stop_cuda_offloading_entries");
393   EntriesE->setVisibility(GlobalValue::HiddenVisibility);
394 
395   auto *RegGlobalsTy = FunctionType::get(Type::getVoidTy(C),
396                                          Type::getInt8PtrTy(C)->getPointerTo(),
397                                          /*isVarArg*/ false);
398   auto *RegGlobalsFn = Function::Create(
399       RegGlobalsTy, GlobalValue::InternalLinkage, ".cuda.globals_reg", &M);
400   RegGlobalsFn->setSection(".text.startup");
401 
402   // Create the loop to register all the entries.
403   IRBuilder<> Builder(BasicBlock::Create(C, "entry", RegGlobalsFn));
404   auto *EntryBB = BasicBlock::Create(C, "while.entry", RegGlobalsFn);
405   auto *IfThenBB = BasicBlock::Create(C, "if.then", RegGlobalsFn);
406   auto *IfElseBB = BasicBlock::Create(C, "if.else", RegGlobalsFn);
407   auto *SwGlobalBB = BasicBlock::Create(C, "sw.global", RegGlobalsFn);
408   auto *SwManagedBB = BasicBlock::Create(C, "sw.managed", RegGlobalsFn);
409   auto *SwSurfaceBB = BasicBlock::Create(C, "sw.surface", RegGlobalsFn);
410   auto *SwTextureBB = BasicBlock::Create(C, "sw.texture", RegGlobalsFn);
411   auto *IfEndBB = BasicBlock::Create(C, "if.end", RegGlobalsFn);
412   auto *ExitBB = BasicBlock::Create(C, "while.end", RegGlobalsFn);
413 
414   auto *EntryCmp = Builder.CreateICmpNE(EntriesB, EntriesE);
415   Builder.CreateCondBr(EntryCmp, EntryBB, ExitBB);
416   Builder.SetInsertPoint(EntryBB);
417   auto *Entry = Builder.CreatePHI(getEntryPtrTy(M), 2, "entry");
418   auto *AddrPtr =
419       Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
420                                 {ConstantInt::get(getSizeTTy(M), 0),
421                                  ConstantInt::get(Type::getInt32Ty(C), 0)});
422   auto *Addr = Builder.CreateLoad(Type::getInt8PtrTy(C), AddrPtr, "addr");
423   auto *NamePtr =
424       Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
425                                 {ConstantInt::get(getSizeTTy(M), 0),
426                                  ConstantInt::get(Type::getInt32Ty(C), 1)});
427   auto *Name = Builder.CreateLoad(Type::getInt8PtrTy(C), NamePtr, "name");
428   auto *SizePtr =
429       Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
430                                 {ConstantInt::get(getSizeTTy(M), 0),
431                                  ConstantInt::get(Type::getInt32Ty(C), 2)});
432   auto *Size = Builder.CreateLoad(getSizeTTy(M), SizePtr, "size");
433   auto *FlagsPtr =
434       Builder.CreateInBoundsGEP(getEntryTy(M), Entry,
435                                 {ConstantInt::get(getSizeTTy(M), 0),
436                                  ConstantInt::get(Type::getInt32Ty(C), 3)});
437   auto *Flags = Builder.CreateLoad(Type::getInt32Ty(C), FlagsPtr, "flag");
438   auto *FnCond =
439       Builder.CreateICmpEQ(Size, ConstantInt::getNullValue(getSizeTTy(M)));
440   Builder.CreateCondBr(FnCond, IfThenBB, IfElseBB);
441 
442   // Create kernel registration code.
443   Builder.SetInsertPoint(IfThenBB);
444   Builder.CreateCall(RegFunc,
445                      {RegGlobalsFn->arg_begin(), Addr, Name, Name,
446                       ConstantInt::get(Type::getInt32Ty(C), -1),
447                       ConstantPointerNull::get(Type::getInt8PtrTy(C)),
448                       ConstantPointerNull::get(Type::getInt8PtrTy(C)),
449                       ConstantPointerNull::get(Type::getInt8PtrTy(C)),
450                       ConstantPointerNull::get(Type::getInt8PtrTy(C)),
451                       ConstantPointerNull::get(Type::getInt32PtrTy(C))});
452   Builder.CreateBr(IfEndBB);
453   Builder.SetInsertPoint(IfElseBB);
454 
455   auto *Switch = Builder.CreateSwitch(Flags, IfEndBB);
456   // Create global variable registration code.
457   Builder.SetInsertPoint(SwGlobalBB);
458   Builder.CreateCall(RegVar, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
459                               ConstantInt::get(Type::getInt32Ty(C), 0), Size,
460                               ConstantInt::get(Type::getInt32Ty(C), 0),
461                               ConstantInt::get(Type::getInt32Ty(C), 0)});
462   Builder.CreateBr(IfEndBB);
463   Switch->addCase(Builder.getInt32(OffloadGlobalEntry), SwGlobalBB);
464 
465   // Create managed variable registration code.
466   Builder.SetInsertPoint(SwManagedBB);
467   Builder.CreateBr(IfEndBB);
468   Switch->addCase(Builder.getInt32(OffloadGlobalManagedEntry), SwManagedBB);
469 
470   // Create surface variable registration code.
471   Builder.SetInsertPoint(SwSurfaceBB);
472   Builder.CreateBr(IfEndBB);
473   Switch->addCase(Builder.getInt32(OffloadGlobalSurfaceEntry), SwSurfaceBB);
474 
475   // Create texture variable registration code.
476   Builder.SetInsertPoint(SwTextureBB);
477   Builder.CreateBr(IfEndBB);
478   Switch->addCase(Builder.getInt32(OffloadGlobalTextureEntry), SwTextureBB);
479 
480   Builder.SetInsertPoint(IfEndBB);
481   auto *NewEntry = Builder.CreateInBoundsGEP(
482       getEntryTy(M), Entry, ConstantInt::get(getSizeTTy(M), 1));
483   auto *Cmp = Builder.CreateICmpEQ(
484       NewEntry,
485       ConstantExpr::getInBoundsGetElementPtr(
486           ArrayType::get(getEntryTy(M), 0), EntriesE,
487           ArrayRef<Constant *>({ConstantInt::get(getSizeTTy(M), 0),
488                                 ConstantInt::get(getSizeTTy(M), 0)})));
489   Entry->addIncoming(
490       ConstantExpr::getInBoundsGetElementPtr(
491           ArrayType::get(getEntryTy(M), 0), EntriesB,
492           ArrayRef<Constant *>({ConstantInt::get(getSizeTTy(M), 0),
493                                 ConstantInt::get(getSizeTTy(M), 0)})),
494       &RegGlobalsFn->getEntryBlock());
495   Entry->addIncoming(NewEntry, IfEndBB);
496   Builder.CreateCondBr(Cmp, ExitBB, EntryBB);
497   Builder.SetInsertPoint(ExitBB);
498   Builder.CreateRetVoid();
499 
500   return RegGlobalsFn;
501 }
502 
503 // Create the constructor and destructor to register the fatbinary with the CUDA
504 // runtime.
505 void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc) {
506   LLVMContext &C = M.getContext();
507   auto *CtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
508   auto *CtorFunc = Function::Create(CtorFuncTy, GlobalValue::InternalLinkage,
509                                     ".cuda.fatbin_reg", &M);
510   CtorFunc->setSection(".text.startup");
511 
512   auto *DtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
513   auto *DtorFunc = Function::Create(DtorFuncTy, GlobalValue::InternalLinkage,
514                                     ".cuda.fatbin_unreg", &M);
515   DtorFunc->setSection(".text.startup");
516 
517   // Get the __cudaRegisterFatBinary function declaration.
518   auto *RegFatTy = FunctionType::get(Type::getInt8PtrTy(C)->getPointerTo(),
519                                      Type::getInt8PtrTy(C),
520                                      /*isVarArg*/ false);
521   FunctionCallee RegFatbin =
522       M.getOrInsertFunction("__cudaRegisterFatBinary", RegFatTy);
523   // Get the __cudaRegisterFatBinaryEnd function declaration.
524   auto *RegFatEndTy = FunctionType::get(Type::getVoidTy(C),
525                                         Type::getInt8PtrTy(C)->getPointerTo(),
526                                         /*isVarArg*/ false);
527   FunctionCallee RegFatbinEnd =
528       M.getOrInsertFunction("__cudaRegisterFatBinaryEnd", RegFatEndTy);
529   // Get the __cudaUnregisterFatBinary function declaration.
530   auto *UnregFatTy = FunctionType::get(Type::getVoidTy(C),
531                                        Type::getInt8PtrTy(C)->getPointerTo(),
532                                        /*isVarArg*/ false);
533   FunctionCallee UnregFatbin =
534       M.getOrInsertFunction("__cudaUnregisterFatBinary", UnregFatTy);
535 
536   auto *AtExitTy =
537       FunctionType::get(Type::getInt32Ty(C), DtorFuncTy->getPointerTo(),
538                         /*isVarArg*/ false);
539   FunctionCallee AtExit = M.getOrInsertFunction("atexit", AtExitTy);
540 
541   auto *BinaryHandleGlobal = new llvm::GlobalVariable(
542       M, Type::getInt8PtrTy(C)->getPointerTo(), false,
543       llvm::GlobalValue::InternalLinkage,
544       llvm::ConstantPointerNull::get(Type::getInt8PtrTy(C)->getPointerTo()),
545       ".cuda.binary_handle");
546 
547   // Create the constructor to register this image with the runtime.
548   IRBuilder<> CtorBuilder(BasicBlock::Create(C, "entry", CtorFunc));
549   CallInst *Handle = CtorBuilder.CreateCall(
550       RegFatbin, ConstantExpr::getPointerBitCastOrAddrSpaceCast(
551                      FatbinDesc, Type::getInt8PtrTy(C)));
552   CtorBuilder.CreateAlignedStore(
553       Handle, BinaryHandleGlobal,
554       Align(M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))));
555   CtorBuilder.CreateCall(createRegisterGlobalsFunction(M), Handle);
556   CtorBuilder.CreateCall(RegFatbinEnd, Handle);
557   CtorBuilder.CreateCall(AtExit, DtorFunc);
558   CtorBuilder.CreateRetVoid();
559 
560   // Create the destructor to unregister the image with the runtime. We cannot
561   // use a standard global destructor after CUDA 9.2 so this must be called by
562   // `atexit()` intead.
563   IRBuilder<> DtorBuilder(BasicBlock::Create(C, "entry", DtorFunc));
564   LoadInst *BinaryHandle = DtorBuilder.CreateAlignedLoad(
565       Type::getInt8PtrTy(C)->getPointerTo(), BinaryHandleGlobal,
566       Align(M.getDataLayout().getPointerTypeSize(Type::getInt8PtrTy(C))));
567   DtorBuilder.CreateCall(UnregFatbin, BinaryHandle);
568   DtorBuilder.CreateRetVoid();
569 
570   // Add this function to constructors.
571   appendToGlobalCtors(M, CtorFunc, /*Priority*/ 1);
572 }
573 
574 } // namespace
575 
576 Error wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images) {
577   GlobalVariable *Desc = createBinDesc(M, Images);
578   if (!Desc)
579     return createStringError(inconvertibleErrorCode(),
580                              "No binary descriptors created.");
581   createRegisterFunction(M, Desc);
582   createUnregisterFunction(M, Desc);
583   return Error::success();
584 }
585 
586 Error wrapCudaBinary(Module &M, ArrayRef<char> Image) {
587   GlobalVariable *Desc = createFatbinDesc(M, Image);
588   if (!Desc)
589     return createStringError(inconvertibleErrorCode(),
590                              "No fatinbary section created.");
591 
592   createRegisterFatbinFunction(M, Desc);
593   return Error::success();
594 }
595