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.
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.
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 
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 
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>>
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>
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.first);
213     StrTab.add(KeyAndValue.second);
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.first),
254                     Offset + StrTab.getOffset(KeyAndValue.second)};
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 
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 
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 
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 
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 
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