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