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/IR/Constants.h"
14 #include "llvm/IR/Module.h"
15 #include "llvm/IRReader/IRReader.h"
16 #include "llvm/MC/StringTableBuilder.h"
17 #include "llvm/Object/Archive.h"
18 #include "llvm/Object/ArchiveWriter.h"
19 #include "llvm/Object/Binary.h"
20 #include "llvm/Object/COFF.h"
21 #include "llvm/Object/ELFObjectFile.h"
22 #include "llvm/Object/Error.h"
23 #include "llvm/Object/IRObjectFile.h"
24 #include "llvm/Object/ObjectFile.h"
25 #include "llvm/Support/Alignment.h"
26 #include "llvm/Support/FileOutputBuffer.h"
27 #include "llvm/Support/SourceMgr.h"
28
29 using namespace llvm;
30 using namespace llvm::object;
31
32 namespace {
33
34 /// Attempts to extract all the embedded device images contained inside the
35 /// buffer \p Contents. The buffer is expected to contain a valid offloading
36 /// binary format.
extractOffloadFiles(MemoryBufferRef Contents,SmallVectorImpl<OffloadFile> & Binaries)37 Error extractOffloadFiles(MemoryBufferRef Contents,
38 SmallVectorImpl<OffloadFile> &Binaries) {
39 uint64_t Offset = 0;
40 // There could be multiple offloading binaries stored at this section.
41 while (Offset < Contents.getBuffer().size()) {
42 std::unique_ptr<MemoryBuffer> Buffer =
43 MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "",
44 /*RequiresNullTerminator*/ false);
45 if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
46 Buffer->getBufferStart()))
47 Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(),
48 Buffer->getBufferIdentifier());
49 auto BinaryOrErr = OffloadBinary::create(*Buffer);
50 if (!BinaryOrErr)
51 return BinaryOrErr.takeError();
52 OffloadBinary &Binary = **BinaryOrErr;
53
54 // Create a new owned binary with a copy of the original memory.
55 std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
56 Binary.getData().take_front(Binary.getSize()),
57 Contents.getBufferIdentifier());
58 auto NewBinaryOrErr = OffloadBinary::create(*BufferCopy);
59 if (!NewBinaryOrErr)
60 return NewBinaryOrErr.takeError();
61 Binaries.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy));
62
63 Offset += Binary.getSize();
64 }
65
66 return Error::success();
67 }
68
69 // Extract offloading binaries from an Object file \p Obj.
extractFromObject(const ObjectFile & Obj,SmallVectorImpl<OffloadFile> & Binaries)70 Error extractFromObject(const ObjectFile &Obj,
71 SmallVectorImpl<OffloadFile> &Binaries) {
72 assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");
73
74 for (SectionRef Sec : Obj.sections()) {
75 // ELF files contain a section with the LLVM_OFFLOADING type.
76 if (Obj.isELF() &&
77 static_cast<ELFSectionRef>(Sec).getType() != ELF::SHT_LLVM_OFFLOADING)
78 continue;
79
80 // COFF has no section types so we rely on the name of the section.
81 if (Obj.isCOFF()) {
82 Expected<StringRef> NameOrErr = Sec.getName();
83 if (!NameOrErr)
84 return NameOrErr.takeError();
85
86 if (!NameOrErr->equals(".llvm.offloading"))
87 continue;
88 }
89
90 Expected<StringRef> Buffer = Sec.getContents();
91 if (!Buffer)
92 return Buffer.takeError();
93
94 MemoryBufferRef Contents(*Buffer, Obj.getFileName());
95 if (Error Err = extractOffloadFiles(Contents, Binaries))
96 return Err;
97 }
98
99 return Error::success();
100 }
101
extractFromBitcode(MemoryBufferRef Buffer,SmallVectorImpl<OffloadFile> & Binaries)102 Error extractFromBitcode(MemoryBufferRef Buffer,
103 SmallVectorImpl<OffloadFile> &Binaries) {
104 LLVMContext Context;
105 SMDiagnostic Err;
106 std::unique_ptr<Module> M = getLazyIRModule(
107 MemoryBuffer::getMemBuffer(Buffer, /*RequiresNullTerminator=*/false), Err,
108 Context);
109 if (!M)
110 return createStringError(inconvertibleErrorCode(),
111 "Failed to create module");
112
113 // Extract offloading data from globals referenced by the
114 // `llvm.embedded.object` metadata with the `.llvm.offloading` section.
115 auto *MD = M->getNamedMetadata("llvm.embedded.objects");
116 if (!MD)
117 return Error::success();
118
119 for (const MDNode *Op : MD->operands()) {
120 if (Op->getNumOperands() < 2)
121 continue;
122
123 MDString *SectionID = dyn_cast<MDString>(Op->getOperand(1));
124 if (!SectionID || SectionID->getString() != ".llvm.offloading")
125 continue;
126
127 GlobalVariable *GV =
128 mdconst::dyn_extract_or_null<GlobalVariable>(Op->getOperand(0));
129 if (!GV)
130 continue;
131
132 auto *CDS = dyn_cast<ConstantDataSequential>(GV->getInitializer());
133 if (!CDS)
134 continue;
135
136 MemoryBufferRef Contents(CDS->getAsString(), M->getName());
137 if (Error Err = extractOffloadFiles(Contents, Binaries))
138 return Err;
139 }
140
141 return Error::success();
142 }
143
extractFromArchive(const Archive & Library,SmallVectorImpl<OffloadFile> & Binaries)144 Error extractFromArchive(const Archive &Library,
145 SmallVectorImpl<OffloadFile> &Binaries) {
146 // Try to extract device code from each file stored in the static archive.
147 Error Err = Error::success();
148 for (auto Child : Library.children(Err)) {
149 auto ChildBufferOrErr = Child.getMemoryBufferRef();
150 if (!ChildBufferOrErr)
151 return ChildBufferOrErr.takeError();
152 std::unique_ptr<MemoryBuffer> ChildBuffer =
153 MemoryBuffer::getMemBuffer(*ChildBufferOrErr, false);
154
155 // Check if the buffer has the required alignment.
156 if (!isAddrAligned(Align(OffloadBinary::getAlignment()),
157 ChildBuffer->getBufferStart()))
158 ChildBuffer = MemoryBuffer::getMemBufferCopy(
159 ChildBufferOrErr->getBuffer(),
160 ChildBufferOrErr->getBufferIdentifier());
161
162 if (Error Err = extractOffloadBinaries(*ChildBuffer, Binaries))
163 return Err;
164 }
165
166 if (Err)
167 return Err;
168 return Error::success();
169 }
170
171 } // namespace
172
173 Expected<std::unique_ptr<OffloadBinary>>
create(MemoryBufferRef Buf)174 OffloadBinary::create(MemoryBufferRef Buf) {
175 if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
176 return errorCodeToError(object_error::parse_failed);
177
178 // Check for 0x10FF1OAD magic bytes.
179 if (identify_magic(Buf.getBuffer()) != file_magic::offload_binary)
180 return errorCodeToError(object_error::parse_failed);
181
182 // Make sure that the data has sufficient alignment.
183 if (!isAddrAligned(Align(getAlignment()), Buf.getBufferStart()))
184 return errorCodeToError(object_error::parse_failed);
185
186 const char *Start = Buf.getBufferStart();
187 const Header *TheHeader = reinterpret_cast<const Header *>(Start);
188 if (TheHeader->Version != OffloadBinary::Version)
189 return errorCodeToError(object_error::parse_failed);
190
191 if (TheHeader->Size > Buf.getBufferSize() ||
192 TheHeader->EntryOffset > TheHeader->Size - sizeof(Entry) ||
193 TheHeader->EntrySize > TheHeader->Size - sizeof(Header))
194 return errorCodeToError(object_error::unexpected_eof);
195
196 const Entry *TheEntry =
197 reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]);
198
199 if (TheEntry->ImageOffset > Buf.getBufferSize() ||
200 TheEntry->StringOffset > Buf.getBufferSize())
201 return errorCodeToError(object_error::unexpected_eof);
202
203 return std::unique_ptr<OffloadBinary>(
204 new OffloadBinary(Buf, TheHeader, TheEntry));
205 }
206
207 std::unique_ptr<MemoryBuffer>
write(const OffloadingImage & OffloadingData)208 OffloadBinary::write(const OffloadingImage &OffloadingData) {
209 // Create a null-terminated string table with all the used strings.
210 StringTableBuilder StrTab(StringTableBuilder::ELF);
211 for (auto &KeyAndValue : OffloadingData.StringData) {
212 StrTab.add(KeyAndValue.getKey());
213 StrTab.add(KeyAndValue.getValue());
214 }
215 StrTab.finalize();
216
217 uint64_t StringEntrySize =
218 sizeof(StringEntry) * OffloadingData.StringData.size();
219
220 // Make sure the image we're wrapping around is aligned as well.
221 uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
222 StringEntrySize + StrTab.getSize(),
223 getAlignment());
224
225 // Create the header and fill in the offsets. The entry will be directly
226 // placed after the header in memory. Align the size to the alignment of the
227 // header so this can be placed contiguously in a single section.
228 Header TheHeader;
229 TheHeader.Size = alignTo(
230 BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
231 TheHeader.EntryOffset = sizeof(Header);
232 TheHeader.EntrySize = sizeof(Entry);
233
234 // Create the entry using the string table offsets. The string table will be
235 // placed directly after the entry in memory, and the image after that.
236 Entry TheEntry;
237 TheEntry.TheImageKind = OffloadingData.TheImageKind;
238 TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
239 TheEntry.Flags = OffloadingData.Flags;
240 TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
241 TheEntry.NumStrings = OffloadingData.StringData.size();
242
243 TheEntry.ImageOffset = BinaryDataSize;
244 TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
245
246 SmallVector<char> Data;
247 Data.reserve(TheHeader.Size);
248 raw_svector_ostream OS(Data);
249 OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
250 OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
251 for (auto &KeyAndValue : OffloadingData.StringData) {
252 uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
253 StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.getKey()),
254 Offset + StrTab.getOffset(KeyAndValue.getValue())};
255 OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
256 }
257 StrTab.write(OS);
258 // Add padding to required image alignment.
259 OS.write_zeros(TheEntry.ImageOffset - OS.tell());
260 OS << OffloadingData.Image->getBuffer();
261
262 // Add final padding to required alignment.
263 assert(TheHeader.Size >= OS.tell() && "Too much data written?");
264 OS.write_zeros(TheHeader.Size - OS.tell());
265 assert(TheHeader.Size == OS.tell() && "Size mismatch");
266
267 return MemoryBuffer::getMemBufferCopy(OS.str());
268 }
269
extractOffloadBinaries(MemoryBufferRef Buffer,SmallVectorImpl<OffloadFile> & Binaries)270 Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
271 SmallVectorImpl<OffloadFile> &Binaries) {
272 file_magic Type = identify_magic(Buffer.getBuffer());
273 switch (Type) {
274 case file_magic::bitcode:
275 return extractFromBitcode(Buffer, Binaries);
276 case file_magic::elf_relocatable:
277 case file_magic::elf_executable:
278 case file_magic::elf_shared_object:
279 case file_magic::coff_object: {
280 Expected<std::unique_ptr<ObjectFile>> ObjFile =
281 ObjectFile::createObjectFile(Buffer, Type);
282 if (!ObjFile)
283 return ObjFile.takeError();
284 return extractFromObject(*ObjFile->get(), Binaries);
285 }
286 case file_magic::archive: {
287 Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
288 object::Archive::create(Buffer);
289 if (!LibFile)
290 return LibFile.takeError();
291 return extractFromArchive(*LibFile->get(), Binaries);
292 }
293 case file_magic::offload_binary:
294 return extractOffloadFiles(Buffer, Binaries);
295 default:
296 return Error::success();
297 }
298 }
299
getOffloadKind(StringRef Name)300 OffloadKind object::getOffloadKind(StringRef Name) {
301 return llvm::StringSwitch<OffloadKind>(Name)
302 .Case("openmp", OFK_OpenMP)
303 .Case("cuda", OFK_Cuda)
304 .Case("hip", OFK_HIP)
305 .Default(OFK_None);
306 }
307
getOffloadKindName(OffloadKind Kind)308 StringRef object::getOffloadKindName(OffloadKind Kind) {
309 switch (Kind) {
310 case OFK_OpenMP:
311 return "openmp";
312 case OFK_Cuda:
313 return "cuda";
314 case OFK_HIP:
315 return "hip";
316 default:
317 return "none";
318 }
319 }
320
getImageKind(StringRef Name)321 ImageKind object::getImageKind(StringRef Name) {
322 return llvm::StringSwitch<ImageKind>(Name)
323 .Case("o", IMG_Object)
324 .Case("bc", IMG_Bitcode)
325 .Case("cubin", IMG_Cubin)
326 .Case("fatbin", IMG_Fatbinary)
327 .Case("s", IMG_PTX)
328 .Default(IMG_None);
329 }
330
getImageKindName(ImageKind Kind)331 StringRef object::getImageKindName(ImageKind Kind) {
332 switch (Kind) {
333 case IMG_Object:
334 return "o";
335 case IMG_Bitcode:
336 return "bc";
337 case IMG_Cubin:
338 return "cubin";
339 case IMG_Fatbinary:
340 return "fatbin";
341 case IMG_PTX:
342 return "s";
343 default:
344 return "";
345 }
346 }
347