xref: /freebsd/contrib/llvm-project/llvm/lib/Frontend/Offloading/OffloadWrapper.cpp (revision 700637cbb5e582861067a11aaca4d053546871d2) 
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(M.getContext());
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(M.getContext());
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,bool Relocatable)114 GlobalVariable *createBinDesc(Module &M, ArrayRef<ArrayRef<char>> Bufs,
115                               EntryArrayTy EntryArray, StringRef Suffix,
116                               bool Relocatable) {
117   LLVMContext &C = M.getContext();
118   auto [EntriesB, EntriesE] = EntryArray;
119 
120   auto *Zero = ConstantInt::get(getSizeTTy(M), 0u);
121   Constant *ZeroZero[] = {Zero, Zero};
122 
123   // Create initializer for the images array.
124   SmallVector<Constant *, 4u> ImagesInits;
125   ImagesInits.reserve(Bufs.size());
126   for (ArrayRef<char> Buf : Bufs) {
127     // We embed the full offloading entry so the binary utilities can parse it.
128     auto *Data = ConstantDataArray::get(C, Buf);
129     auto *Image = new GlobalVariable(M, Data->getType(), /*isConstant=*/true,
130                                      GlobalVariable::InternalLinkage, Data,
131                                      ".omp_offloading.device_image" + Suffix);
132     Image->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
133     Image->setSection(Relocatable ? ".llvm.offloading.relocatable"
134                                   : ".llvm.offloading");
135     Image->setAlignment(Align(object::OffloadBinary::getAlignment()));
136 
137     StringRef Binary(Buf.data(), Buf.size());
138     assert(identify_magic(Binary) == file_magic::offload_binary &&
139            "Invalid binary format");
140 
141     // The device image struct contains the pointer to the beginning and end of
142     // the image stored inside of the offload binary. There should only be one
143     // of these for each buffer so we parse it out manually.
144     const auto *Header =
145         reinterpret_cast<const object::OffloadBinary::Header *>(
146             Binary.bytes_begin());
147     const auto *Entry = reinterpret_cast<const object::OffloadBinary::Entry *>(
148         Binary.bytes_begin() + Header->EntryOffset);
149 
150     auto *Begin = ConstantInt::get(getSizeTTy(M), Entry->ImageOffset);
151     auto *Size =
152         ConstantInt::get(getSizeTTy(M), Entry->ImageOffset + Entry->ImageSize);
153     Constant *ZeroBegin[] = {Zero, Begin};
154     Constant *ZeroSize[] = {Zero, Size};
155 
156     auto *ImageB =
157         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroBegin);
158     auto *ImageE =
159         ConstantExpr::getGetElementPtr(Image->getValueType(), Image, ZeroSize);
160 
161     ImagesInits.push_back(ConstantStruct::get(getDeviceImageTy(M), ImageB,
162                                               ImageE, EntriesB, EntriesE));
163   }
164 
165   // Then create images array.
166   auto *ImagesData = ConstantArray::get(
167       ArrayType::get(getDeviceImageTy(M), ImagesInits.size()), ImagesInits);
168 
169   auto *Images =
170       new GlobalVariable(M, ImagesData->getType(), /*isConstant*/ true,
171                          GlobalValue::InternalLinkage, ImagesData,
172                          ".omp_offloading.device_images" + Suffix);
173   Images->setUnnamedAddr(GlobalValue::UnnamedAddr::Global);
174 
175   auto *ImagesB =
176       ConstantExpr::getGetElementPtr(Images->getValueType(), Images, ZeroZero);
177 
178   // And finally create the binary descriptor object.
179   auto *DescInit = ConstantStruct::get(
180       getBinDescTy(M),
181       ConstantInt::get(Type::getInt32Ty(C), ImagesInits.size()), ImagesB,
182       EntriesB, EntriesE);
183 
184   return new GlobalVariable(M, DescInit->getType(), /*isConstant=*/true,
185                             GlobalValue::InternalLinkage, DescInit,
186                             ".omp_offloading.descriptor" + Suffix);
187 }
188 
createUnregisterFunction(Module & M,GlobalVariable * BinDesc,StringRef Suffix)189 Function *createUnregisterFunction(Module &M, GlobalVariable *BinDesc,
190                                    StringRef Suffix) {
191   LLVMContext &C = M.getContext();
192   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
193   auto *Func =
194       Function::Create(FuncTy, GlobalValue::InternalLinkage,
195                        ".omp_offloading.descriptor_unreg" + Suffix, &M);
196   Func->setSection(".text.startup");
197 
198   // Get __tgt_unregister_lib function declaration.
199   auto *UnRegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(M),
200                                         /*isVarArg*/ false);
201   FunctionCallee UnRegFuncC =
202       M.getOrInsertFunction("__tgt_unregister_lib", UnRegFuncTy);
203 
204   // Construct function body
205   IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
206   Builder.CreateCall(UnRegFuncC, BinDesc);
207   Builder.CreateRetVoid();
208 
209   return Func;
210 }
211 
createRegisterFunction(Module & M,GlobalVariable * BinDesc,StringRef Suffix)212 void createRegisterFunction(Module &M, GlobalVariable *BinDesc,
213                             StringRef Suffix) {
214   LLVMContext &C = M.getContext();
215   auto *FuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
216   auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage,
217                                 ".omp_offloading.descriptor_reg" + Suffix, &M);
218   Func->setSection(".text.startup");
219 
220   // Get __tgt_register_lib function declaration.
221   auto *RegFuncTy = FunctionType::get(Type::getVoidTy(C), getBinDescPtrTy(M),
222                                       /*isVarArg*/ false);
223   FunctionCallee RegFuncC =
224       M.getOrInsertFunction("__tgt_register_lib", RegFuncTy);
225 
226   auto *AtExitTy = FunctionType::get(
227       Type::getInt32Ty(C), PointerType::getUnqual(C), /*isVarArg=*/false);
228   FunctionCallee AtExit = M.getOrInsertFunction("atexit", AtExitTy);
229 
230   Function *UnregFunc = createUnregisterFunction(M, BinDesc, Suffix);
231 
232   // Construct function body
233   IRBuilder<> Builder(BasicBlock::Create(C, "entry", Func));
234 
235   Builder.CreateCall(RegFuncC, BinDesc);
236 
237   // Register the destructors with 'atexit'. This is expected by the CUDA
238   // runtime and ensures that we clean up before dynamic objects are destroyed.
239   // This needs to be done after plugin initialization to ensure that it is
240   // called before the plugin runtime is destroyed.
241   Builder.CreateCall(AtExit, UnregFunc);
242   Builder.CreateRetVoid();
243 
244   // Add this function to constructors.
245   appendToGlobalCtors(M, Func, /*Priority=*/101);
246 }
247 
248 // struct fatbin_wrapper {
249 //  int32_t magic;
250 //  int32_t version;
251 //  void *image;
252 //  void *reserved;
253 //};
getFatbinWrapperTy(Module & M)254 StructType *getFatbinWrapperTy(Module &M) {
255   LLVMContext &C = M.getContext();
256   StructType *FatbinTy = StructType::getTypeByName(C, "fatbin_wrapper");
257   if (!FatbinTy)
258     FatbinTy = StructType::create(
259         "fatbin_wrapper", Type::getInt32Ty(C), Type::getInt32Ty(C),
260         PointerType::getUnqual(C), PointerType::getUnqual(C));
261   return FatbinTy;
262 }
263 
264 /// Embed the image \p Image into the module \p M so it can be found by the
265 /// runtime.
createFatbinDesc(Module & M,ArrayRef<char> Image,bool IsHIP,StringRef Suffix)266 GlobalVariable *createFatbinDesc(Module &M, ArrayRef<char> Image, bool IsHIP,
267                                  StringRef Suffix) {
268   LLVMContext &C = M.getContext();
269   llvm::Type *Int8PtrTy = PointerType::getUnqual(C);
270   const llvm::Triple &Triple = M.getTargetTriple();
271 
272   // Create the global string containing the fatbinary.
273   StringRef FatbinConstantSection =
274       IsHIP ? ".hip_fatbin"
275             : (Triple.isMacOSX() ? "__NV_CUDA,__nv_fatbin" : ".nv_fatbin");
276   auto *Data = ConstantDataArray::get(C, Image);
277   auto *Fatbin = new GlobalVariable(M, Data->getType(), /*isConstant*/ true,
278                                     GlobalVariable::InternalLinkage, Data,
279                                     ".fatbin_image" + Suffix);
280   Fatbin->setSection(FatbinConstantSection);
281 
282   // Create the fatbinary wrapper
283   StringRef FatbinWrapperSection = IsHIP               ? ".hipFatBinSegment"
284                                    : Triple.isMacOSX() ? "__NV_CUDA,__fatbin"
285                                                        : ".nvFatBinSegment";
286   Constant *FatbinWrapper[] = {
287       ConstantInt::get(Type::getInt32Ty(C), IsHIP ? HIPFatMagic : CudaFatMagic),
288       ConstantInt::get(Type::getInt32Ty(C), 1),
289       ConstantExpr::getPointerBitCastOrAddrSpaceCast(Fatbin, Int8PtrTy),
290       ConstantPointerNull::get(PointerType::getUnqual(C))};
291 
292   Constant *FatbinInitializer =
293       ConstantStruct::get(getFatbinWrapperTy(M), FatbinWrapper);
294 
295   auto *FatbinDesc =
296       new GlobalVariable(M, getFatbinWrapperTy(M),
297                          /*isConstant*/ true, GlobalValue::InternalLinkage,
298                          FatbinInitializer, ".fatbin_wrapper" + Suffix);
299   FatbinDesc->setSection(FatbinWrapperSection);
300   FatbinDesc->setAlignment(Align(8));
301 
302   return FatbinDesc;
303 }
304 
305 /// Create the register globals function. We will iterate all of the offloading
306 /// entries stored at the begin / end symbols and register them according to
307 /// their type. This creates the following function in IR:
308 ///
309 /// extern struct __tgt_offload_entry __start_cuda_offloading_entries;
310 /// extern struct __tgt_offload_entry __stop_cuda_offloading_entries;
311 ///
312 /// extern void __cudaRegisterFunction(void **, void *, void *, void *, int,
313 ///                                    void *, void *, void *, void *, int *);
314 /// extern void __cudaRegisterVar(void **, void *, void *, void *, int32_t,
315 ///                               int64_t, int32_t, int32_t);
316 ///
317 /// void __cudaRegisterTest(void **fatbinHandle) {
318 ///   for (struct __tgt_offload_entry *entry = &__start_cuda_offloading_entries;
319 ///        entry != &__stop_cuda_offloading_entries; ++entry) {
320 ///     if (entry->Kind != OFK_CUDA)
321 ///       continue
322 ///
323 ///     if (!entry->Size)
324 ///       __cudaRegisterFunction(fatbinHandle, entry->addr, entry->name,
325 ///                              entry->name, -1, 0, 0, 0, 0, 0);
326 ///     else
327 ///       __cudaRegisterVar(fatbinHandle, entry->addr, entry->name, entry->name,
328 ///                         0, entry->size, 0, 0);
329 ///   }
330 /// }
createRegisterGlobalsFunction(Module & M,bool IsHIP,EntryArrayTy EntryArray,StringRef Suffix,bool EmitSurfacesAndTextures)331 Function *createRegisterGlobalsFunction(Module &M, bool IsHIP,
332                                         EntryArrayTy EntryArray,
333                                         StringRef Suffix,
334                                         bool EmitSurfacesAndTextures) {
335   LLVMContext &C = M.getContext();
336   auto [EntriesB, EntriesE] = EntryArray;
337 
338   // Get the __cudaRegisterFunction function declaration.
339   PointerType *Int8PtrTy = PointerType::get(C, 0);
340   PointerType *Int8PtrPtrTy = PointerType::get(C, 0);
341   PointerType *Int32PtrTy = PointerType::get(C, 0);
342   auto *RegFuncTy = FunctionType::get(
343       Type::getInt32Ty(C),
344       {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy, Type::getInt32Ty(C),
345        Int8PtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy, Int32PtrTy},
346       /*isVarArg*/ false);
347   FunctionCallee RegFunc = M.getOrInsertFunction(
348       IsHIP ? "__hipRegisterFunction" : "__cudaRegisterFunction", RegFuncTy);
349 
350   // Get the __cudaRegisterVar function declaration.
351   auto *RegVarTy = FunctionType::get(
352       Type::getVoidTy(C),
353       {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy, Type::getInt32Ty(C),
354        getSizeTTy(M), Type::getInt32Ty(C), Type::getInt32Ty(C)},
355       /*isVarArg*/ false);
356   FunctionCallee RegVar = M.getOrInsertFunction(
357       IsHIP ? "__hipRegisterVar" : "__cudaRegisterVar", RegVarTy);
358 
359   // Get the __cudaRegisterSurface function declaration.
360   FunctionType *RegManagedVarTy =
361       FunctionType::get(Type::getVoidTy(C),
362                         {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy,
363                          getSizeTTy(M), Type::getInt32Ty(C)},
364                         /*isVarArg=*/false);
365   FunctionCallee RegManagedVar = M.getOrInsertFunction(
366       IsHIP ? "__hipRegisterManagedVar" : "__cudaRegisterManagedVar",
367       RegManagedVarTy);
368 
369   // Get the __cudaRegisterSurface function declaration.
370   FunctionType *RegSurfaceTy =
371       FunctionType::get(Type::getVoidTy(C),
372                         {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy,
373                          Type::getInt32Ty(C), Type::getInt32Ty(C)},
374                         /*isVarArg=*/false);
375   FunctionCallee RegSurface = M.getOrInsertFunction(
376       IsHIP ? "__hipRegisterSurface" : "__cudaRegisterSurface", RegSurfaceTy);
377 
378   // Get the __cudaRegisterTexture function declaration.
379   FunctionType *RegTextureTy = FunctionType::get(
380       Type::getVoidTy(C),
381       {Int8PtrPtrTy, Int8PtrTy, Int8PtrTy, Int8PtrTy, Type::getInt32Ty(C),
382        Type::getInt32Ty(C), Type::getInt32Ty(C)},
383       /*isVarArg=*/false);
384   FunctionCallee RegTexture = M.getOrInsertFunction(
385       IsHIP ? "__hipRegisterTexture" : "__cudaRegisterTexture", RegTextureTy);
386 
387   auto *RegGlobalsTy = FunctionType::get(Type::getVoidTy(C), Int8PtrPtrTy,
388                                          /*isVarArg*/ false);
389   auto *RegGlobalsFn =
390       Function::Create(RegGlobalsTy, GlobalValue::InternalLinkage,
391                        IsHIP ? ".hip.globals_reg" : ".cuda.globals_reg", &M);
392   RegGlobalsFn->setSection(".text.startup");
393 
394   // Create the loop to register all the entries.
395   IRBuilder<> Builder(BasicBlock::Create(C, "entry", RegGlobalsFn));
396   auto *EntryBB = BasicBlock::Create(C, "while.entry", RegGlobalsFn);
397   auto *IfKindBB = BasicBlock::Create(C, "if.kind", RegGlobalsFn);
398   auto *IfThenBB = BasicBlock::Create(C, "if.then", RegGlobalsFn);
399   auto *IfElseBB = BasicBlock::Create(C, "if.else", RegGlobalsFn);
400   auto *SwGlobalBB = BasicBlock::Create(C, "sw.global", RegGlobalsFn);
401   auto *SwManagedBB = BasicBlock::Create(C, "sw.managed", RegGlobalsFn);
402   auto *SwSurfaceBB = BasicBlock::Create(C, "sw.surface", RegGlobalsFn);
403   auto *SwTextureBB = BasicBlock::Create(C, "sw.texture", RegGlobalsFn);
404   auto *IfEndBB = BasicBlock::Create(C, "if.end", RegGlobalsFn);
405   auto *ExitBB = BasicBlock::Create(C, "while.end", RegGlobalsFn);
406 
407   auto *EntryCmp = Builder.CreateICmpNE(EntriesB, EntriesE);
408   Builder.CreateCondBr(EntryCmp, EntryBB, ExitBB);
409   Builder.SetInsertPoint(EntryBB);
410   auto *Entry = Builder.CreatePHI(PointerType::getUnqual(C), 2, "entry");
411   auto *AddrPtr =
412       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
413                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
414                                  ConstantInt::get(Type::getInt32Ty(C), 4)});
415   auto *Addr = Builder.CreateLoad(Int8PtrTy, AddrPtr, "addr");
416   auto *AuxAddrPtr =
417       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
418                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
419                                  ConstantInt::get(Type::getInt32Ty(C), 8)});
420   auto *AuxAddr = Builder.CreateLoad(Int8PtrTy, AuxAddrPtr, "aux_addr");
421   auto *KindPtr =
422       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
423                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
424                                  ConstantInt::get(Type::getInt32Ty(C), 2)});
425   auto *Kind = Builder.CreateLoad(Type::getInt16Ty(C), KindPtr, "kind");
426   auto *NamePtr =
427       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
428                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
429                                  ConstantInt::get(Type::getInt32Ty(C), 5)});
430   auto *Name = Builder.CreateLoad(Int8PtrTy, NamePtr, "name");
431   auto *SizePtr =
432       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
433                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
434                                  ConstantInt::get(Type::getInt32Ty(C), 6)});
435   auto *Size = Builder.CreateLoad(Type::getInt64Ty(C), SizePtr, "size");
436   auto *FlagsPtr =
437       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
438                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
439                                  ConstantInt::get(Type::getInt32Ty(C), 3)});
440   auto *Flags = Builder.CreateLoad(Type::getInt32Ty(C), FlagsPtr, "flags");
441   auto *DataPtr =
442       Builder.CreateInBoundsGEP(offloading::getEntryTy(M), Entry,
443                                 {ConstantInt::get(Type::getInt32Ty(C), 0),
444                                  ConstantInt::get(Type::getInt32Ty(C), 7)});
445   auto *Data = Builder.CreateTrunc(
446       Builder.CreateLoad(Type::getInt64Ty(C), DataPtr, "data"),
447       Type::getInt32Ty(C));
448   auto *Type = Builder.CreateAnd(
449       Flags, ConstantInt::get(Type::getInt32Ty(C), 0x7), "type");
450 
451   // Extract the flags stored in the bit-field and convert them to C booleans.
452   auto *ExternBit = Builder.CreateAnd(
453       Flags, ConstantInt::get(Type::getInt32Ty(C),
454                               llvm::offloading::OffloadGlobalExtern));
455   auto *Extern = Builder.CreateLShr(
456       ExternBit, ConstantInt::get(Type::getInt32Ty(C), 3), "extern");
457   auto *ConstantBit = Builder.CreateAnd(
458       Flags, ConstantInt::get(Type::getInt32Ty(C),
459                               llvm::offloading::OffloadGlobalConstant));
460   auto *Const = Builder.CreateLShr(
461       ConstantBit, ConstantInt::get(Type::getInt32Ty(C), 4), "constant");
462   auto *NormalizedBit = Builder.CreateAnd(
463       Flags, ConstantInt::get(Type::getInt32Ty(C),
464                               llvm::offloading::OffloadGlobalNormalized));
465   auto *Normalized = Builder.CreateLShr(
466       NormalizedBit, ConstantInt::get(Type::getInt32Ty(C), 5), "normalized");
467   auto *KindCond = Builder.CreateICmpEQ(
468       Kind, ConstantInt::get(Type::getInt16Ty(C),
469                              IsHIP ? object::OffloadKind::OFK_HIP
470                                    : object::OffloadKind::OFK_Cuda));
471   Builder.CreateCondBr(KindCond, IfKindBB, IfEndBB);
472   Builder.SetInsertPoint(IfKindBB);
473   auto *FnCond = Builder.CreateICmpEQ(
474       Size, ConstantInt::getNullValue(Type::getInt64Ty(C)));
475   Builder.CreateCondBr(FnCond, IfThenBB, IfElseBB);
476 
477   // Create kernel registration code.
478   Builder.SetInsertPoint(IfThenBB);
479   Builder.CreateCall(RegFunc, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
480                                ConstantInt::get(Type::getInt32Ty(C), -1),
481                                ConstantPointerNull::get(Int8PtrTy),
482                                ConstantPointerNull::get(Int8PtrTy),
483                                ConstantPointerNull::get(Int8PtrTy),
484                                ConstantPointerNull::get(Int8PtrTy),
485                                ConstantPointerNull::get(Int32PtrTy)});
486   Builder.CreateBr(IfEndBB);
487   Builder.SetInsertPoint(IfElseBB);
488 
489   auto *Switch = Builder.CreateSwitch(Type, IfEndBB);
490   // Create global variable registration code.
491   Builder.SetInsertPoint(SwGlobalBB);
492   Builder.CreateCall(RegVar,
493                      {RegGlobalsFn->arg_begin(), Addr, Name, Name, Extern, Size,
494                       Const, ConstantInt::get(Type::getInt32Ty(C), 0)});
495   Builder.CreateBr(IfEndBB);
496   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalEntry),
497                   SwGlobalBB);
498 
499   // Create managed variable registration code.
500   Builder.SetInsertPoint(SwManagedBB);
501   Builder.CreateCall(RegManagedVar, {RegGlobalsFn->arg_begin(), AuxAddr, Addr,
502                                      Name, Size, Data});
503   Builder.CreateBr(IfEndBB);
504   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalManagedEntry),
505                   SwManagedBB);
506   // Create surface variable registration code.
507   Builder.SetInsertPoint(SwSurfaceBB);
508   if (EmitSurfacesAndTextures)
509     Builder.CreateCall(RegSurface, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
510                                     Data, Extern});
511   Builder.CreateBr(IfEndBB);
512   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalSurfaceEntry),
513                   SwSurfaceBB);
514 
515   // Create texture variable registration code.
516   Builder.SetInsertPoint(SwTextureBB);
517   if (EmitSurfacesAndTextures)
518     Builder.CreateCall(RegTexture, {RegGlobalsFn->arg_begin(), Addr, Name, Name,
519                                     Data, Normalized, Extern});
520   Builder.CreateBr(IfEndBB);
521   Switch->addCase(Builder.getInt32(llvm::offloading::OffloadGlobalTextureEntry),
522                   SwTextureBB);
523 
524   Builder.SetInsertPoint(IfEndBB);
525   auto *NewEntry = Builder.CreateInBoundsGEP(
526       offloading::getEntryTy(M), Entry, ConstantInt::get(getSizeTTy(M), 1));
527   auto *Cmp = Builder.CreateICmpEQ(
528       NewEntry,
529       ConstantExpr::getInBoundsGetElementPtr(
530           ArrayType::get(offloading::getEntryTy(M), 0), EntriesE,
531           ArrayRef<Constant *>({ConstantInt::get(getSizeTTy(M), 0),
532                                 ConstantInt::get(getSizeTTy(M), 0)})));
533   Entry->addIncoming(
534       ConstantExpr::getInBoundsGetElementPtr(
535           ArrayType::get(offloading::getEntryTy(M), 0), EntriesB,
536           ArrayRef<Constant *>({ConstantInt::get(getSizeTTy(M), 0),
537                                 ConstantInt::get(getSizeTTy(M), 0)})),
538       &RegGlobalsFn->getEntryBlock());
539   Entry->addIncoming(NewEntry, IfEndBB);
540   Builder.CreateCondBr(Cmp, ExitBB, EntryBB);
541   Builder.SetInsertPoint(ExitBB);
542   Builder.CreateRetVoid();
543 
544   return RegGlobalsFn;
545 }
546 
547 // Create the constructor and destructor to register the fatbinary with the CUDA
548 // runtime.
createRegisterFatbinFunction(Module & M,GlobalVariable * FatbinDesc,bool IsHIP,EntryArrayTy EntryArray,StringRef Suffix,bool EmitSurfacesAndTextures)549 void createRegisterFatbinFunction(Module &M, GlobalVariable *FatbinDesc,
550                                   bool IsHIP, EntryArrayTy EntryArray,
551                                   StringRef Suffix,
552                                   bool EmitSurfacesAndTextures) {
553   LLVMContext &C = M.getContext();
554   auto *CtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
555   auto *CtorFunc = Function::Create(
556       CtorFuncTy, GlobalValue::InternalLinkage,
557       (IsHIP ? ".hip.fatbin_reg" : ".cuda.fatbin_reg") + Suffix, &M);
558   CtorFunc->setSection(".text.startup");
559 
560   auto *DtorFuncTy = FunctionType::get(Type::getVoidTy(C), /*isVarArg*/ false);
561   auto *DtorFunc = Function::Create(
562       DtorFuncTy, GlobalValue::InternalLinkage,
563       (IsHIP ? ".hip.fatbin_unreg" : ".cuda.fatbin_unreg") + Suffix, &M);
564   DtorFunc->setSection(".text.startup");
565 
566   auto *PtrTy = PointerType::getUnqual(C);
567 
568   // Get the __cudaRegisterFatBinary function declaration.
569   auto *RegFatTy = FunctionType::get(PtrTy, PtrTy, /*isVarArg=*/false);
570   FunctionCallee RegFatbin = M.getOrInsertFunction(
571       IsHIP ? "__hipRegisterFatBinary" : "__cudaRegisterFatBinary", RegFatTy);
572   // Get the __cudaRegisterFatBinaryEnd function declaration.
573   auto *RegFatEndTy =
574       FunctionType::get(Type::getVoidTy(C), PtrTy, /*isVarArg=*/false);
575   FunctionCallee RegFatbinEnd =
576       M.getOrInsertFunction("__cudaRegisterFatBinaryEnd", RegFatEndTy);
577   // Get the __cudaUnregisterFatBinary function declaration.
578   auto *UnregFatTy =
579       FunctionType::get(Type::getVoidTy(C), PtrTy, /*isVarArg=*/false);
580   FunctionCallee UnregFatbin = M.getOrInsertFunction(
581       IsHIP ? "__hipUnregisterFatBinary" : "__cudaUnregisterFatBinary",
582       UnregFatTy);
583 
584   auto *AtExitTy =
585       FunctionType::get(Type::getInt32Ty(C), PtrTy, /*isVarArg=*/false);
586   FunctionCallee AtExit = M.getOrInsertFunction("atexit", AtExitTy);
587 
588   auto *BinaryHandleGlobal = new llvm::GlobalVariable(
589       M, PtrTy, false, llvm::GlobalValue::InternalLinkage,
590       llvm::ConstantPointerNull::get(PtrTy),
591       (IsHIP ? ".hip.binary_handle" : ".cuda.binary_handle") + Suffix);
592 
593   // Create the constructor to register this image with the runtime.
594   IRBuilder<> CtorBuilder(BasicBlock::Create(C, "entry", CtorFunc));
595   CallInst *Handle = CtorBuilder.CreateCall(
596       RegFatbin,
597       ConstantExpr::getPointerBitCastOrAddrSpaceCast(FatbinDesc, PtrTy));
598   CtorBuilder.CreateAlignedStore(
599       Handle, BinaryHandleGlobal,
600       Align(M.getDataLayout().getPointerTypeSize(PtrTy)));
601   CtorBuilder.CreateCall(createRegisterGlobalsFunction(M, IsHIP, EntryArray,
602                                                        Suffix,
603                                                        EmitSurfacesAndTextures),
604                          Handle);
605   if (!IsHIP)
606     CtorBuilder.CreateCall(RegFatbinEnd, Handle);
607   CtorBuilder.CreateCall(AtExit, DtorFunc);
608   CtorBuilder.CreateRetVoid();
609 
610   // Create the destructor to unregister the image with the runtime. We cannot
611   // use a standard global destructor after CUDA 9.2 so this must be called by
612   // `atexit()` instead.
613   IRBuilder<> DtorBuilder(BasicBlock::Create(C, "entry", DtorFunc));
614   LoadInst *BinaryHandle = DtorBuilder.CreateAlignedLoad(
615       PtrTy, BinaryHandleGlobal,
616       Align(M.getDataLayout().getPointerTypeSize(PtrTy)));
617   DtorBuilder.CreateCall(UnregFatbin, BinaryHandle);
618   DtorBuilder.CreateRetVoid();
619 
620   // Add this function to constructors.
621   appendToGlobalCtors(M, CtorFunc, /*Priority=*/101);
622 }
623 } // namespace
624 
wrapOpenMPBinaries(Module & M,ArrayRef<ArrayRef<char>> Images,EntryArrayTy EntryArray,llvm::StringRef Suffix,bool Relocatable)625 Error offloading::wrapOpenMPBinaries(Module &M, ArrayRef<ArrayRef<char>> Images,
626                                      EntryArrayTy EntryArray,
627                                      llvm::StringRef Suffix, bool Relocatable) {
628   GlobalVariable *Desc =
629       createBinDesc(M, Images, EntryArray, Suffix, Relocatable);
630   if (!Desc)
631     return createStringError(inconvertibleErrorCode(),
632                              "No binary descriptors created.");
633   createRegisterFunction(M, Desc, Suffix);
634   return Error::success();
635 }
636 
wrapCudaBinary(Module & M,ArrayRef<char> Image,EntryArrayTy EntryArray,llvm::StringRef Suffix,bool EmitSurfacesAndTextures)637 Error offloading::wrapCudaBinary(Module &M, ArrayRef<char> Image,
638                                  EntryArrayTy EntryArray,
639                                  llvm::StringRef Suffix,
640                                  bool EmitSurfacesAndTextures) {
641   GlobalVariable *Desc = createFatbinDesc(M, Image, /*IsHip=*/false, Suffix);
642   if (!Desc)
643     return createStringError(inconvertibleErrorCode(),
644                              "No fatbin section created.");
645 
646   createRegisterFatbinFunction(M, Desc, /*IsHip=*/false, EntryArray, Suffix,
647                                EmitSurfacesAndTextures);
648   return Error::success();
649 }
650 
wrapHIPBinary(Module & M,ArrayRef<char> Image,EntryArrayTy EntryArray,llvm::StringRef Suffix,bool EmitSurfacesAndTextures)651 Error offloading::wrapHIPBinary(Module &M, ArrayRef<char> Image,
652                                 EntryArrayTy EntryArray, llvm::StringRef Suffix,
653                                 bool EmitSurfacesAndTextures) {
654   GlobalVariable *Desc = createFatbinDesc(M, Image, /*IsHip=*/true, Suffix);
655   if (!Desc)
656     return createStringError(inconvertibleErrorCode(),
657                              "No fatbin section created.");
658 
659   createRegisterFatbinFunction(M, Desc, /*IsHip=*/true, EntryArray, Suffix,
660                                EmitSurfacesAndTextures);
661   return Error::success();
662 }
663