1 //===--- Offloading.h - 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 // This file contains the binary format used for budingling device metadata with
10 // an associated device image. The data can then be stored inside a host object
11 // file to create a fat binary and read by the linker. This is intended to be a
12 // thin wrapper around the image itself. If this format becomes sufficiently
13 // complex it should be moved to a standard binary format like msgpack or ELF.
14 //
15 //===----------------------------------------------------------------------===//
16 
17 #ifndef LLVM_OBJECT_OFFLOADBINARY_H
18 #define LLVM_OBJECT_OFFLOADBINARY_H
19 
20 #include "llvm/ADT/StringMap.h"
21 #include "llvm/ADT/StringRef.h"
22 #include "llvm/Object/Binary.h"
23 #include "llvm/Support/Error.h"
24 #include "llvm/Support/MemoryBuffer.h"
25 #include <memory>
26 
27 namespace llvm {
28 
29 namespace object {
30 
31 /// The producer of the associated offloading image.
32 enum OffloadKind : uint16_t {
33   OFK_None = 0,
34   OFK_OpenMP,
35   OFK_Cuda,
36   OFK_HIP,
37   OFK_LAST,
38 };
39 
40 /// The type of contents the offloading image contains.
41 enum ImageKind : uint16_t {
42   IMG_None = 0,
43   IMG_Object,
44   IMG_Bitcode,
45   IMG_Cubin,
46   IMG_Fatbinary,
47   IMG_PTX,
48   IMG_LAST,
49 };
50 
51 /// A simple binary serialization of an offloading file. We use this format to
52 /// embed the offloading image into the host executable so it can be extracted
53 /// and used by the linker.
54 ///
55 /// Many of these could be stored in the same section by the time the linker
56 /// sees it so we mark this information with a header. The version is used to
57 /// detect ABI stability and the size is used to find other offloading entries
58 /// that may exist in the same section. All offsets are given as absolute byte
59 /// offsets from the beginning of the file.
60 class OffloadBinary : public Binary {
61 public:
62   using string_iterator = StringMap<StringRef>::const_iterator;
63   using string_iterator_range = iterator_range<string_iterator>;
64 
65   /// The current version of the binary used for backwards compatibility.
66   static const uint32_t Version = 1;
67 
68   /// The offloading metadata that will be serialized to a memory buffer.
69   struct OffloadingImage {
70     ImageKind TheImageKind;
71     OffloadKind TheOffloadKind;
72     uint32_t Flags;
73     StringMap<StringRef> StringData;
74     std::unique_ptr<MemoryBuffer> Image;
75   };
76 
77   /// Attempt to parse the offloading binary stored in \p Data.
78   static Expected<std::unique_ptr<OffloadBinary>> create(MemoryBufferRef);
79 
80   /// Serialize the contents of \p File to a binary buffer to be read later.
81   static std::unique_ptr<MemoryBuffer> write(const OffloadingImage &);
82 
getAlignment()83   static uint64_t getAlignment() { return 8; }
84 
getImageKind()85   ImageKind getImageKind() const { return TheEntry->TheImageKind; }
getOffloadKind()86   OffloadKind getOffloadKind() const { return TheEntry->TheOffloadKind; }
getVersion()87   uint32_t getVersion() const { return TheHeader->Version; }
getFlags()88   uint32_t getFlags() const { return TheEntry->Flags; }
getSize()89   uint64_t getSize() const { return TheHeader->Size; }
90 
getTriple()91   StringRef getTriple() const { return getString("triple"); }
getArch()92   StringRef getArch() const { return getString("arch"); }
getImage()93   StringRef getImage() const {
94     return StringRef(&Buffer[TheEntry->ImageOffset], TheEntry->ImageSize);
95   }
96 
97   // Iterator over all the key and value pairs in the binary.
strings()98   string_iterator_range strings() const {
99     return string_iterator_range(StringData.begin(), StringData.end());
100   }
101 
getString(StringRef Key)102   StringRef getString(StringRef Key) const { return StringData.lookup(Key); }
103 
classof(const Binary * V)104   static bool classof(const Binary *V) { return V->isOffloadFile(); }
105 
106   struct Header {
107     uint8_t Magic[4] = {0x10, 0xFF, 0x10, 0xAD}; // 0x10FF10AD magic bytes.
108     uint32_t Version = OffloadBinary::Version;   // Version identifier.
109     uint64_t Size;        // Size in bytes of this entire binary.
110     uint64_t EntryOffset; // Offset of the metadata entry in bytes.
111     uint64_t EntrySize;   // Size of the metadata entry in bytes.
112   };
113 
114   struct Entry {
115     ImageKind TheImageKind;     // The kind of the image stored.
116     OffloadKind TheOffloadKind; // The producer of this image.
117     uint32_t Flags;             // Additional flags associated with the image.
118     uint64_t StringOffset;      // Offset in bytes to the string map.
119     uint64_t NumStrings;        // Number of entries in the string map.
120     uint64_t ImageOffset;       // Offset in bytes of the actual binary image.
121     uint64_t ImageSize;         // Size in bytes of the binary image.
122   };
123 
124   struct StringEntry {
125     uint64_t KeyOffset;
126     uint64_t ValueOffset;
127   };
128 
129 private:
OffloadBinary(MemoryBufferRef Source,const Header * TheHeader,const Entry * TheEntry)130   OffloadBinary(MemoryBufferRef Source, const Header *TheHeader,
131                 const Entry *TheEntry)
132       : Binary(Binary::ID_Offload, Source), Buffer(Source.getBufferStart()),
133         TheHeader(TheHeader), TheEntry(TheEntry) {
134     const StringEntry *StringMapBegin =
135         reinterpret_cast<const StringEntry *>(&Buffer[TheEntry->StringOffset]);
136     for (uint64_t I = 0, E = TheEntry->NumStrings; I != E; ++I) {
137       StringRef Key = &Buffer[StringMapBegin[I].KeyOffset];
138       StringData[Key] = &Buffer[StringMapBegin[I].ValueOffset];
139     }
140   }
141 
142   OffloadBinary(const OffloadBinary &Other) = delete;
143 
144   /// Map from keys to offsets in the binary.
145   StringMap<StringRef> StringData;
146   /// Raw pointer to the MemoryBufferRef for convenience.
147   const char *Buffer;
148   /// Location of the header within the binary.
149   const Header *TheHeader;
150   /// Location of the metadata entries within the binary.
151   const Entry *TheEntry;
152 };
153 
154 /// A class to contain the binary information for a single OffloadBinary that
155 /// owns its memory.
156 class OffloadFile : public OwningBinary<OffloadBinary> {
157 public:
158   using TargetID = std::pair<StringRef, StringRef>;
159 
OffloadFile(std::unique_ptr<OffloadBinary> Binary,std::unique_ptr<MemoryBuffer> Buffer)160   OffloadFile(std::unique_ptr<OffloadBinary> Binary,
161               std::unique_ptr<MemoryBuffer> Buffer)
162       : OwningBinary<OffloadBinary>(std::move(Binary), std::move(Buffer)) {}
163 
164   /// We use the Triple and Architecture pair to group linker inputs together.
165   /// This conversion function lets us use these inputs in a hash-map.
TargetID()166   operator TargetID() const {
167     return std::make_pair(getBinary()->getTriple(), getBinary()->getArch());
168   }
169 };
170 
171 /// Extracts embedded device offloading code from a memory \p Buffer to a list
172 /// of \p Binaries.
173 Error extractOffloadBinaries(MemoryBufferRef Buffer,
174                              SmallVectorImpl<OffloadFile> &Binaries);
175 
176 /// Convert a string \p Name to an image kind.
177 ImageKind getImageKind(StringRef Name);
178 
179 /// Convert an image kind to its string representation.
180 StringRef getImageKindName(ImageKind Name);
181 
182 /// Convert a string \p Name to an offload kind.
183 OffloadKind getOffloadKind(StringRef Name);
184 
185 /// Convert an offload kind to its string representation.
186 StringRef getOffloadKindName(OffloadKind Name);
187 
188 } // namespace object
189 
190 } // namespace llvm
191 #endif
192