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 
write(const OffloadingImage & OffloadingData)207 SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
208   // Create a null-terminated string table with all the used strings.
209   StringTableBuilder StrTab(StringTableBuilder::ELF);
210   for (auto &KeyAndValue : OffloadingData.StringData) {
211     StrTab.add(KeyAndValue.first);
212     StrTab.add(KeyAndValue.second);
213   }
214   StrTab.finalize();
215 
216   uint64_t StringEntrySize =
217       sizeof(StringEntry) * OffloadingData.StringData.size();
218 
219   // Make sure the image we're wrapping around is aligned as well.
220   uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
221                                         StringEntrySize + StrTab.getSize(),
222                                     getAlignment());
223 
224   // Create the header and fill in the offsets. The entry will be directly
225   // placed after the header in memory. Align the size to the alignment of the
226   // header so this can be placed contiguously in a single section.
227   Header TheHeader;
228   TheHeader.Size = alignTo(
229       BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
230   TheHeader.EntryOffset = sizeof(Header);
231   TheHeader.EntrySize = sizeof(Entry);
232 
233   // Create the entry using the string table offsets. The string table will be
234   // placed directly after the entry in memory, and the image after that.
235   Entry TheEntry;
236   TheEntry.TheImageKind = OffloadingData.TheImageKind;
237   TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
238   TheEntry.Flags = OffloadingData.Flags;
239   TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
240   TheEntry.NumStrings = OffloadingData.StringData.size();
241 
242   TheEntry.ImageOffset = BinaryDataSize;
243   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
244 
245   SmallString<0> Data;
246   Data.reserve(TheHeader.Size);
247   raw_svector_ostream OS(Data);
248   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
249   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
250   for (auto &KeyAndValue : OffloadingData.StringData) {
251     uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
252     StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.first),
253                     Offset + StrTab.getOffset(KeyAndValue.second)};
254     OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
255   }
256   StrTab.write(OS);
257   // Add padding to required image alignment.
258   OS.write_zeros(TheEntry.ImageOffset - OS.tell());
259   OS << OffloadingData.Image->getBuffer();
260 
261   // Add final padding to required alignment.
262   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
263   OS.write_zeros(TheHeader.Size - OS.tell());
264   assert(TheHeader.Size == OS.tell() && "Size mismatch");
265 
266   return Data;
267 }
268 
extractOffloadBinaries(MemoryBufferRef Buffer,SmallVectorImpl<OffloadFile> & Binaries)269 Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
270                                      SmallVectorImpl<OffloadFile> &Binaries) {
271   file_magic Type = identify_magic(Buffer.getBuffer());
272   switch (Type) {
273   case file_magic::bitcode:
274     return extractFromBitcode(Buffer, Binaries);
275   case file_magic::elf_relocatable:
276   case file_magic::elf_executable:
277   case file_magic::elf_shared_object:
278   case file_magic::coff_object: {
279     Expected<std::unique_ptr<ObjectFile>> ObjFile =
280         ObjectFile::createObjectFile(Buffer, Type);
281     if (!ObjFile)
282       return ObjFile.takeError();
283     return extractFromObject(*ObjFile->get(), Binaries);
284   }
285   case file_magic::archive: {
286     Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
287         object::Archive::create(Buffer);
288     if (!LibFile)
289       return LibFile.takeError();
290     return extractFromArchive(*LibFile->get(), Binaries);
291   }
292   case file_magic::offload_binary:
293     return extractOffloadFiles(Buffer, Binaries);
294   default:
295     return Error::success();
296   }
297 }
298 
getOffloadKind(StringRef Name)299 OffloadKind object::getOffloadKind(StringRef Name) {
300   return llvm::StringSwitch<OffloadKind>(Name)
301       .Case("openmp", OFK_OpenMP)
302       .Case("cuda", OFK_Cuda)
303       .Case("hip", OFK_HIP)
304       .Default(OFK_None);
305 }
306 
getOffloadKindName(OffloadKind Kind)307 StringRef object::getOffloadKindName(OffloadKind Kind) {
308   switch (Kind) {
309   case OFK_OpenMP:
310     return "openmp";
311   case OFK_Cuda:
312     return "cuda";
313   case OFK_HIP:
314     return "hip";
315   default:
316     return "none";
317   }
318 }
319 
getImageKind(StringRef Name)320 ImageKind object::getImageKind(StringRef Name) {
321   return llvm::StringSwitch<ImageKind>(Name)
322       .Case("o", IMG_Object)
323       .Case("bc", IMG_Bitcode)
324       .Case("cubin", IMG_Cubin)
325       .Case("fatbin", IMG_Fatbinary)
326       .Case("s", IMG_PTX)
327       .Default(IMG_None);
328 }
329 
getImageKindName(ImageKind Kind)330 StringRef object::getImageKindName(ImageKind Kind) {
331   switch (Kind) {
332   case IMG_Object:
333     return "o";
334   case IMG_Bitcode:
335     return "bc";
336   case IMG_Cubin:
337     return "cubin";
338   case IMG_Fatbinary:
339     return "fatbin";
340   case IMG_PTX:
341     return "s";
342   default:
343     return "";
344   }
345 }
346 
areTargetsCompatible(const OffloadFile::TargetID & LHS,const OffloadFile::TargetID & RHS)347 bool object::areTargetsCompatible(const OffloadFile::TargetID &LHS,
348                                   const OffloadFile::TargetID &RHS) {
349   // Exact matches are not considered compatible because they are the same
350   // target. We are interested in different targets that are compatible.
351   if (LHS == RHS)
352     return false;
353 
354   // The triples must match at all times.
355   if (LHS.first != RHS.first)
356     return false;
357 
358   // Only The AMDGPU target requires additional checks.
359   llvm::Triple T(LHS.first);
360   if (!T.isAMDGPU())
361     return false;
362 
363   // The base processor must always match.
364   if (LHS.second.split(":").first != RHS.second.split(":").first)
365     return false;
366 
367   // Check combintions of on / off features that must match.
368   if (LHS.second.contains("xnack+") && RHS.second.contains("xnack-"))
369     return false;
370   if (LHS.second.contains("xnack-") && RHS.second.contains("xnack+"))
371     return false;
372   if (LHS.second.contains("sramecc-") && RHS.second.contains("sramecc+"))
373     return false;
374   if (LHS.second.contains("sramecc+") && RHS.second.contains("sramecc-"))
375     return false;
376   return true;
377 }
378