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