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