1 //===- Offloading.cpp - Utilities for handling offloading code  -*- 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/Object/OffloadBinary.h"
10 
11 #include "llvm/ADT/StringSwitch.h"
12 #include "llvm/BinaryFormat/Magic.h"
13 #include "llvm/MC/StringTableBuilder.h"
14 #include "llvm/Object/Error.h"
15 #include "llvm/Support/FileOutputBuffer.h"
16 
17 using namespace llvm;
18 using namespace llvm::object;
19 
20 Expected<std::unique_ptr<OffloadBinary>>
21 OffloadBinary::create(MemoryBufferRef Buf) {
22   if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
23     return errorCodeToError(object_error::parse_failed);
24 
25   // Check for 0x10FF1OAD magic bytes.
26   if (identify_magic(Buf.getBuffer()) != file_magic::offload_binary)
27     return errorCodeToError(object_error::parse_failed);
28 
29   const char *Start = Buf.getBufferStart();
30   const Header *TheHeader = reinterpret_cast<const Header *>(Start);
31   const Entry *TheEntry =
32       reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]);
33 
34   return std::unique_ptr<OffloadBinary>(
35       new OffloadBinary(Buf, TheHeader, TheEntry));
36 }
37 
38 std::unique_ptr<MemoryBuffer>
39 OffloadBinary::write(const OffloadingImage &OffloadingData) {
40   // Create a null-terminated string table with all the used strings.
41   StringTableBuilder StrTab(StringTableBuilder::ELF);
42   for (auto &KeyAndValue : OffloadingData.StringData) {
43     StrTab.add(KeyAndValue.getKey());
44     StrTab.add(KeyAndValue.getValue());
45   }
46   StrTab.finalize();
47 
48   uint64_t StringEntrySize =
49       sizeof(StringEntry) * OffloadingData.StringData.size();
50 
51   // Make sure the image we're wrapping around is aligned as well.
52   uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
53                                         StringEntrySize + StrTab.getSize(),
54                                     getAlignment());
55 
56   // Create the header and fill in the offsets. The entry will be directly
57   // placed after the header in memory. Align the size to the alignment of the
58   // header so this can be placed contiguously in a single section.
59   Header TheHeader;
60   TheHeader.Size = alignTo(
61       BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
62   TheHeader.EntryOffset = sizeof(Header);
63   TheHeader.EntrySize = sizeof(Entry);
64 
65   // Create the entry using the string table offsets. The string table will be
66   // placed directly after the entry in memory, and the image after that.
67   Entry TheEntry;
68   TheEntry.TheImageKind = OffloadingData.TheImageKind;
69   TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
70   TheEntry.Flags = OffloadingData.Flags;
71   TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
72   TheEntry.NumStrings = OffloadingData.StringData.size();
73 
74   TheEntry.ImageOffset = BinaryDataSize;
75   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
76 
77   SmallVector<char, 1024> Data;
78   raw_svector_ostream OS(Data);
79   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
80   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
81   for (auto &KeyAndValue : OffloadingData.StringData) {
82     uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
83     StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.getKey()),
84                     Offset + StrTab.getOffset(KeyAndValue.getValue())};
85     OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
86   }
87   StrTab.write(OS);
88   // Add padding to required image alignment.
89   OS.write_zeros(TheEntry.ImageOffset - OS.tell());
90   OS << OffloadingData.Image->getBuffer();
91 
92   // Add final padding to required alignment.
93   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
94   OS.write_zeros(TheHeader.Size - OS.tell());
95   assert(TheHeader.Size == OS.tell() && "Size mismatch");
96 
97   return MemoryBuffer::getMemBufferCopy(OS.str());
98 }
99 
100 OffloadKind object::getOffloadKind(StringRef Name) {
101   return llvm::StringSwitch<OffloadKind>(Name)
102       .Case("openmp", OFK_OpenMP)
103       .Case("cuda", OFK_Cuda)
104       .Case("hip", OFK_HIP)
105       .Default(OFK_None);
106 }
107 
108 StringRef object::getOffloadKindName(OffloadKind Kind) {
109   switch (Kind) {
110   case OFK_OpenMP:
111     return "openmp";
112   case OFK_Cuda:
113     return "cuda";
114   case OFK_HIP:
115     return "hip";
116   default:
117     return "none";
118   }
119 }
120 
121 ImageKind object::getImageKind(StringRef Name) {
122   return llvm::StringSwitch<ImageKind>(Name)
123       .Case("o", IMG_Object)
124       .Case("bc", IMG_Bitcode)
125       .Case("cubin", IMG_Cubin)
126       .Case("fatbin", IMG_Fatbinary)
127       .Case("s", IMG_PTX)
128       .Default(IMG_None);
129 }
130 
131 StringRef object::getImageKindName(ImageKind Kind) {
132   switch (Kind) {
133   case IMG_Object:
134     return "o";
135   case IMG_Bitcode:
136     return "bc";
137   case IMG_Cubin:
138     return "cubin";
139   case IMG_Fatbinary:
140     return "fatbin";
141   case IMG_PTX:
142     return "s";
143   default:
144     return "";
145   }
146 }
147