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/Binary.h"
19#include "llvm/Object/ELFObjectFile.h"
20#include "llvm/Object/Error.h"
21#include "llvm/Object/IRObjectFile.h"
22#include "llvm/Object/ObjectFile.h"
23#include "llvm/Support/Alignment.h"
24#include "llvm/Support/SourceMgr.h"
25
26using namespace llvm;
27using namespace llvm::object;
28
29namespace {
30
31/// A MemoryBuffer that shares ownership of the underlying memory.
32/// This allows multiple OffloadBinary instances to share the same buffer.
33class SharedMemoryBuffer : public MemoryBuffer {
34public:
35 SharedMemoryBuffer(std::shared_ptr<MemoryBuffer> Buf)
36 : SharedBuf(std::move(Buf)) {
37 init(BufStart: SharedBuf->getBufferStart(), BufEnd: SharedBuf->getBufferEnd(),
38 /*RequiresNullTerminator=*/false);
39 }
40
41 BufferKind getBufferKind() const override { return MemoryBuffer_Malloc; }
42
43 StringRef getBufferIdentifier() const override {
44 return SharedBuf->getBufferIdentifier();
45 }
46
47private:
48 const std::shared_ptr<MemoryBuffer> SharedBuf;
49};
50
51/// Attempts to extract all the embedded device images contained inside the
52/// buffer \p Contents. The buffer is expected to contain a valid offloading
53/// binary format.
54Error extractOffloadFiles(MemoryBufferRef Contents,
55 SmallVectorImpl<OffloadFile> &Binaries) {
56 uint64_t Offset = 0;
57 // There could be multiple offloading binaries stored at this section.
58 while (Offset < Contents.getBufferSize()) {
59 std::unique_ptr<MemoryBuffer> Buffer =
60 MemoryBuffer::getMemBuffer(InputData: Contents.getBuffer().drop_front(N: Offset), BufferName: "",
61 /*RequiresNullTerminator*/ false);
62 if (!isAddrAligned(Lhs: Align(OffloadBinary::getAlignment()),
63 Addr: Buffer->getBufferStart()))
64 Buffer = MemoryBuffer::getMemBufferCopy(InputData: Buffer->getBuffer(),
65 BufferName: Buffer->getBufferIdentifier());
66
67 auto HeaderOrErr = OffloadBinary::extractHeader(Buf: *Buffer);
68 if (!HeaderOrErr)
69 return HeaderOrErr.takeError();
70 const OffloadBinary::Header *Header = *HeaderOrErr;
71
72 // Create a copy of original memory containing only the current binary.
73 std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy(
74 InputData: Buffer->getBuffer().take_front(N: Header->Size),
75 BufferName: Contents.getBufferIdentifier());
76
77 auto BinariesOrErr = OffloadBinary::create(Buf: *BufferCopy);
78 if (!BinariesOrErr)
79 return BinariesOrErr.takeError();
80
81 // Share ownership among multiple OffloadFiles.
82 std::shared_ptr<MemoryBuffer> SharedBuffer =
83 std::shared_ptr<MemoryBuffer>(std::move(BufferCopy));
84
85 for (auto &Binary : *BinariesOrErr) {
86 std::unique_ptr<SharedMemoryBuffer> SharedBufferPtr =
87 std::make_unique<SharedMemoryBuffer>(args&: SharedBuffer);
88 Binaries.emplace_back(Args: std::move(Binary), Args: std::move(SharedBufferPtr));
89 }
90
91 Offset += Header->Size;
92 }
93
94 return Error::success();
95}
96
97// Extract offloading binaries from an Object file \p Obj.
98Error extractFromObject(const ObjectFile &Obj,
99 SmallVectorImpl<OffloadFile> &Binaries) {
100 assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type");
101
102 for (SectionRef Sec : Obj.sections()) {
103 // ELF files contain a section with the LLVM_OFFLOADING type.
104 if (Obj.isELF() &&
105 static_cast<ELFSectionRef>(Sec).getType() != ELF::SHT_LLVM_OFFLOADING)
106 continue;
107
108 // COFF has no section types so we rely on the name of the section.
109 if (Obj.isCOFF()) {
110 Expected<StringRef> NameOrErr = Sec.getName();
111 if (!NameOrErr)
112 return NameOrErr.takeError();
113
114 if (!NameOrErr->starts_with(Prefix: ".llvm.offloading"))
115 continue;
116 }
117
118 Expected<StringRef> Buffer = Sec.getContents();
119 if (!Buffer)
120 return Buffer.takeError();
121
122 MemoryBufferRef Contents(*Buffer, Obj.getFileName());
123 if (Error Err = extractOffloadFiles(Contents, Binaries))
124 return Err;
125 }
126
127 return Error::success();
128}
129
130Error extractFromBitcode(MemoryBufferRef Buffer,
131 SmallVectorImpl<OffloadFile> &Binaries) {
132 LLVMContext Context;
133 SMDiagnostic Err;
134 std::unique_ptr<Module> M = getLazyIRModule(
135 Buffer: MemoryBuffer::getMemBuffer(Ref: Buffer, /*RequiresNullTerminator=*/false), Err,
136 Context);
137 if (!M)
138 return createStringError(EC: inconvertibleErrorCode(),
139 S: "Failed to create module");
140
141 // Extract offloading data from globals referenced by the
142 // `llvm.embedded.object` metadata with the `.llvm.offloading` section.
143 auto *MD = M->getNamedMetadata(Name: "llvm.embedded.objects");
144 if (!MD)
145 return Error::success();
146
147 for (const MDNode *Op : MD->operands()) {
148 if (Op->getNumOperands() < 2)
149 continue;
150
151 MDString *SectionID = dyn_cast<MDString>(Val: Op->getOperand(I: 1));
152 if (!SectionID || SectionID->getString() != ".llvm.offloading")
153 continue;
154
155 GlobalVariable *GV =
156 mdconst::dyn_extract_or_null<GlobalVariable>(MD: Op->getOperand(I: 0));
157 if (!GV)
158 continue;
159
160 auto *CDS = dyn_cast<ConstantDataSequential>(Val: GV->getInitializer());
161 if (!CDS)
162 continue;
163
164 MemoryBufferRef Contents(CDS->getAsString(), M->getName());
165 if (Error Err = extractOffloadFiles(Contents, Binaries))
166 return Err;
167 }
168
169 return Error::success();
170}
171
172Error extractFromArchive(const Archive &Library,
173 SmallVectorImpl<OffloadFile> &Binaries) {
174 // Try to extract device code from each file stored in the static archive.
175 Error Err = Error::success();
176 for (auto Child : Library.children(Err)) {
177 auto ChildBufferOrErr = Child.getMemoryBufferRef();
178 if (!ChildBufferOrErr)
179 return ChildBufferOrErr.takeError();
180 std::unique_ptr<MemoryBuffer> ChildBuffer =
181 MemoryBuffer::getMemBuffer(Ref: *ChildBufferOrErr, RequiresNullTerminator: false);
182
183 // Check if the buffer has the required alignment.
184 if (!isAddrAligned(Lhs: Align(OffloadBinary::getAlignment()),
185 Addr: ChildBuffer->getBufferStart()))
186 ChildBuffer = MemoryBuffer::getMemBufferCopy(
187 InputData: ChildBufferOrErr->getBuffer(),
188 BufferName: ChildBufferOrErr->getBufferIdentifier());
189
190 if (Error Err = extractOffloadBinaries(Buffer: *ChildBuffer, Binaries))
191 return Err;
192 }
193
194 if (Err)
195 return Err;
196 return Error::success();
197}
198
199} // namespace
200
201Expected<const OffloadBinary::Header *>
202OffloadBinary::extractHeader(MemoryBufferRef Buf) {
203 if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
204 return errorCodeToError(EC: object_error::parse_failed);
205
206 // Check for 0x10FF1OAD magic bytes.
207 if (identify_magic(magic: Buf.getBuffer()) != file_magic::offload_binary)
208 return errorCodeToError(EC: object_error::parse_failed);
209
210 // Make sure that the data has sufficient alignment.
211 if (!isAddrAligned(Lhs: Align(getAlignment()), Addr: Buf.getBufferStart()))
212 return errorCodeToError(EC: object_error::parse_failed);
213
214 const char *Start = Buf.getBufferStart();
215 const Header *TheHeader = reinterpret_cast<const Header *>(Start);
216 if (TheHeader->Version == 0 || TheHeader->Version > OffloadBinary::Version)
217 return errorCodeToError(EC: object_error::parse_failed);
218
219 if (TheHeader->Size > Buf.getBufferSize() ||
220 TheHeader->Size < sizeof(Entry) || TheHeader->Size < sizeof(Header))
221 return errorCodeToError(EC: object_error::unexpected_eof);
222
223 uint64_t EntriesCount =
224 (TheHeader->Version == 1) ? 1 : TheHeader->EntriesCount;
225 uint64_t EntriesSize = sizeof(Entry) * EntriesCount;
226 if (TheHeader->EntriesOffset > TheHeader->Size - EntriesSize ||
227 EntriesSize > TheHeader->Size - sizeof(Header))
228 return errorCodeToError(EC: object_error::unexpected_eof);
229
230 return TheHeader;
231}
232
233Expected<SmallVector<std::unique_ptr<OffloadBinary>>>
234OffloadBinary::create(MemoryBufferRef Buf, std::optional<uint64_t> Index) {
235 auto HeaderOrErr = OffloadBinary::extractHeader(Buf);
236 if (!HeaderOrErr)
237 return HeaderOrErr.takeError();
238 const Header *TheHeader = *HeaderOrErr;
239
240 const char *Start = Buf.getBufferStart();
241 const Entry *Entries =
242 reinterpret_cast<const Entry *>(&Start[TheHeader->EntriesOffset]);
243
244 auto validateEntry = [&](const Entry *TheEntry) -> Error {
245 if (TheEntry->ImageOffset > Buf.getBufferSize() ||
246 TheEntry->StringOffset > Buf.getBufferSize() ||
247 TheEntry->StringOffset + TheEntry->NumStrings * sizeof(StringEntry) >
248 Buf.getBufferSize())
249 return errorCodeToError(EC: object_error::unexpected_eof);
250 return Error::success();
251 };
252
253 SmallVector<std::unique_ptr<OffloadBinary>> Binaries;
254 if (TheHeader->Version > 1 && Index.has_value()) {
255 if (*Index >= TheHeader->EntriesCount)
256 return errorCodeToError(EC: object_error::parse_failed);
257 const Entry *TheEntry = &Entries[*Index];
258 if (auto Err = validateEntry(TheEntry))
259 return std::move(Err);
260
261 Binaries.emplace_back(Args: new OffloadBinary(Buf, TheHeader, TheEntry, *Index));
262 return std::move(Binaries);
263 }
264
265 uint64_t EntriesCount = TheHeader->Version == 1 ? 1 : TheHeader->EntriesCount;
266 for (uint64_t I = 0; I < EntriesCount; ++I) {
267 const Entry *TheEntry = &Entries[I];
268 if (auto Err = validateEntry(TheEntry))
269 return std::move(Err);
270
271 Binaries.emplace_back(Args: new OffloadBinary(Buf, TheHeader, TheEntry, I));
272 }
273
274 return std::move(Binaries);
275}
276
277SmallString<0> OffloadBinary::write(ArrayRef<OffloadingImage> OffloadingData) {
278 uint64_t EntriesCount = OffloadingData.size();
279 assert(EntriesCount > 0 && "At least one offloading image is required");
280
281 // Create a null-terminated string table with all the used strings.
282 // Also calculate total size of images.
283 StringTableBuilder StrTab(StringTableBuilder::ELF);
284 uint64_t TotalStringEntries = 0;
285 uint64_t TotalImagesSize = 0;
286 for (const OffloadingImage &Img : OffloadingData) {
287 for (auto &KeyAndValue : Img.StringData) {
288 StrTab.add(S: KeyAndValue.first);
289 StrTab.add(S: KeyAndValue.second);
290 }
291 TotalStringEntries += Img.StringData.size();
292 TotalImagesSize += Img.Image->getBufferSize();
293 }
294 StrTab.finalize();
295
296 uint64_t StringEntrySize = sizeof(StringEntry) * TotalStringEntries;
297 uint64_t EntriesSize = sizeof(Entry) * EntriesCount;
298 uint64_t StrTabOffset = sizeof(Header) + EntriesSize + StringEntrySize;
299
300 // Make sure the image we're wrapping around is aligned as well.
301 uint64_t BinaryDataSize =
302 alignTo(Value: StrTabOffset + StrTab.getSize(), Align: getAlignment());
303
304 // Create the header and fill in the offsets. The entries will be directly
305 // placed after the header in memory. Align the size to the alignment of the
306 // header so this can be placed contiguously in a single section.
307 Header TheHeader;
308 TheHeader.Size = alignTo(Value: BinaryDataSize + TotalImagesSize, Align: getAlignment());
309 TheHeader.EntriesOffset = sizeof(Header);
310 TheHeader.EntriesCount = EntriesCount;
311
312 SmallString<0> Data;
313 Data.reserve(N: TheHeader.Size);
314 raw_svector_ostream OS(Data);
315 OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
316
317 // Create the entries using the string table offsets. The string table will be
318 // placed directly after the set of entries in memory, and all the images are
319 // after that.
320 uint64_t StringEntryOffset = sizeof(Header) + EntriesSize;
321 uint64_t ImageOffset = BinaryDataSize;
322 for (const OffloadingImage &Img : OffloadingData) {
323 Entry TheEntry;
324
325 TheEntry.TheImageKind = Img.TheImageKind;
326 TheEntry.TheOffloadKind = Img.TheOffloadKind;
327 TheEntry.Flags = Img.Flags;
328
329 TheEntry.StringOffset = StringEntryOffset;
330 StringEntryOffset += sizeof(StringEntry) * Img.StringData.size();
331 TheEntry.NumStrings = Img.StringData.size();
332
333 TheEntry.ImageOffset = ImageOffset;
334 ImageOffset += Img.Image->getBufferSize();
335 TheEntry.ImageSize = Img.Image->getBufferSize();
336
337 OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
338 }
339
340 // Create the string map entries.
341 for (const OffloadingImage &Img : OffloadingData) {
342 for (auto &KeyAndValue : Img.StringData) {
343 StringEntry Map{.KeyOffset: StrTabOffset + StrTab.getOffset(S: KeyAndValue.first),
344 .ValueOffset: StrTabOffset + StrTab.getOffset(S: KeyAndValue.second),
345 .ValueSize: KeyAndValue.second.size()};
346 OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
347 }
348 }
349
350 StrTab.write(OS);
351 // Add padding to required image alignment.
352 OS.write_zeros(NumZeros: BinaryDataSize - OS.tell());
353
354 for (const OffloadingImage &Img : OffloadingData)
355 OS << Img.Image->getBuffer();
356
357 // Add final padding to required alignment.
358 assert(TheHeader.Size >= OS.tell() && "Too much data written?");
359 OS.write_zeros(NumZeros: TheHeader.Size - OS.tell());
360 assert(TheHeader.Size == OS.tell() && "Size mismatch");
361
362 return Data;
363}
364
365Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
366 SmallVectorImpl<OffloadFile> &Binaries) {
367 file_magic Type = identify_magic(magic: Buffer.getBuffer());
368 switch (Type) {
369 case file_magic::bitcode:
370 return extractFromBitcode(Buffer, Binaries);
371 case file_magic::elf_relocatable:
372 case file_magic::elf_executable:
373 case file_magic::elf_shared_object:
374 case file_magic::coff_object: {
375 Expected<std::unique_ptr<ObjectFile>> ObjFile =
376 ObjectFile::createObjectFile(Object: Buffer, Type);
377 if (!ObjFile)
378 return ObjFile.takeError();
379 return extractFromObject(Obj: *ObjFile->get(), Binaries);
380 }
381 case file_magic::archive: {
382 Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
383 object::Archive::create(Source: Buffer);
384 if (!LibFile)
385 return LibFile.takeError();
386 return extractFromArchive(Library: *LibFile->get(), Binaries);
387 }
388 case file_magic::offload_binary:
389 return extractOffloadFiles(Contents: Buffer, Binaries);
390 default:
391 return Error::success();
392 }
393}
394
395OffloadKind object::getOffloadKind(StringRef Name) {
396 return llvm::StringSwitch<OffloadKind>(Name)
397 .Case(S: "openmp", Value: OFK_OpenMP)
398 .Case(S: "cuda", Value: OFK_Cuda)
399 .Case(S: "hip", Value: OFK_HIP)
400 .Case(S: "sycl", Value: OFK_SYCL)
401 .Default(Value: OFK_None);
402}
403
404StringRef object::getOffloadKindName(OffloadKind Kind) {
405 switch (Kind) {
406 case OFK_OpenMP:
407 return "openmp";
408 case OFK_Cuda:
409 return "cuda";
410 case OFK_HIP:
411 return "hip";
412 case OFK_SYCL:
413 return "sycl";
414 default:
415 return "none";
416 }
417}
418
419ImageKind object::getImageKind(StringRef Name) {
420 return llvm::StringSwitch<ImageKind>(Name)
421 .Case(S: "o", Value: IMG_Object)
422 .Case(S: "bc", Value: IMG_Bitcode)
423 .Case(S: "cubin", Value: IMG_Cubin)
424 .Case(S: "fatbin", Value: IMG_Fatbinary)
425 .Case(S: "s", Value: IMG_PTX)
426 .Case(S: "spv", Value: IMG_SPIRV)
427 .Default(Value: IMG_None);
428}
429
430StringRef object::getImageKindName(ImageKind Kind) {
431 switch (Kind) {
432 case IMG_Object:
433 return "o";
434 case IMG_Bitcode:
435 return "bc";
436 case IMG_Cubin:
437 return "cubin";
438 case IMG_Fatbinary:
439 return "fatbin";
440 case IMG_PTX:
441 return "s";
442 case IMG_SPIRV:
443 return "spv";
444 default:
445 return "";
446 }
447}
448
449bool object::areTargetsCompatible(const OffloadFile::TargetID &LHS,
450 const OffloadFile::TargetID &RHS) {
451 // Exact matches are not considered compatible because they are the same
452 // target. We are interested in different targets that are compatible.
453 if (LHS == RHS)
454 return false;
455
456 // The triples must match at all times.
457 if (LHS.first != RHS.first)
458 return false;
459
460 // If the architecture is "all" we assume it is always compatible.
461 if (LHS.second == "generic" || RHS.second == "generic")
462 return true;
463
464 // Only The AMDGPU target requires additional checks.
465 llvm::Triple T(LHS.first);
466 if (!T.isAMDGPU())
467 return false;
468
469 // The base processor must always match.
470 if (LHS.second.split(Separator: ":").first != RHS.second.split(Separator: ":").first)
471 return false;
472
473 // Check combintions of on / off features that must match.
474 if (LHS.second.contains(Other: "xnack+") && RHS.second.contains(Other: "xnack-"))
475 return false;
476 if (LHS.second.contains(Other: "xnack-") && RHS.second.contains(Other: "xnack+"))
477 return false;
478 if (LHS.second.contains(Other: "sramecc-") && RHS.second.contains(Other: "sramecc+"))
479 return false;
480 if (LHS.second.contains(Other: "sramecc+") && RHS.second.contains(Other: "sramecc-"))
481 return false;
482 return true;
483}
484