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