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/MC/StringTableBuilder.h"
14 #include "llvm/Object/Error.h"
15 #include "llvm/Support/FileOutputBuffer.h"
16 
17 using namespace llvm;
18 using namespace llvm::object;
19 
20 Expected<std::unique_ptr<OffloadBinary>>
21 OffloadBinary::create(MemoryBufferRef Buf) {
22   if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry))
23     return errorCodeToError(object_error::parse_failed);
24 
25   // Check for 0x10FF1OAD magic bytes.
26   if (identify_magic(Buf.getBuffer()) != file_magic::offload_binary)
27     return errorCodeToError(object_error::parse_failed);
28 
29   const char *Start = Buf.getBufferStart();
30   const Header *TheHeader = reinterpret_cast<const Header *>(Start);
31   if (TheHeader->Version != OffloadBinary::Version)
32     return errorCodeToError(object_error::parse_failed);
33 
34   if (TheHeader->Size > Buf.getBufferSize() ||
35       TheHeader->EntryOffset > TheHeader->Size - sizeof(Entry) ||
36       TheHeader->EntrySize > TheHeader->Size - sizeof(Header))
37     return errorCodeToError(object_error::unexpected_eof);
38 
39   const Entry *TheEntry =
40       reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]);
41 
42   if (TheEntry->ImageOffset > Buf.getBufferSize() ||
43       TheEntry->StringOffset > Buf.getBufferSize())
44     return errorCodeToError(object_error::unexpected_eof);
45 
46   return std::unique_ptr<OffloadBinary>(
47       new OffloadBinary(Buf, TheHeader, TheEntry));
48 }
49 
50 std::unique_ptr<MemoryBuffer>
51 OffloadBinary::write(const OffloadingImage &OffloadingData) {
52   // Create a null-terminated string table with all the used strings.
53   StringTableBuilder StrTab(StringTableBuilder::ELF);
54   for (auto &KeyAndValue : OffloadingData.StringData) {
55     StrTab.add(KeyAndValue.getKey());
56     StrTab.add(KeyAndValue.getValue());
57   }
58   StrTab.finalize();
59 
60   uint64_t StringEntrySize =
61       sizeof(StringEntry) * OffloadingData.StringData.size();
62 
63   // Make sure the image we're wrapping around is aligned as well.
64   uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
65                                         StringEntrySize + StrTab.getSize(),
66                                     getAlignment());
67 
68   // Create the header and fill in the offsets. The entry will be directly
69   // placed after the header in memory. Align the size to the alignment of the
70   // header so this can be placed contiguously in a single section.
71   Header TheHeader;
72   TheHeader.Size = alignTo(
73       BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
74   TheHeader.EntryOffset = sizeof(Header);
75   TheHeader.EntrySize = sizeof(Entry);
76 
77   // Create the entry using the string table offsets. The string table will be
78   // placed directly after the entry in memory, and the image after that.
79   Entry TheEntry;
80   TheEntry.TheImageKind = OffloadingData.TheImageKind;
81   TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
82   TheEntry.Flags = OffloadingData.Flags;
83   TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
84   TheEntry.NumStrings = OffloadingData.StringData.size();
85 
86   TheEntry.ImageOffset = BinaryDataSize;
87   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
88 
89   SmallVector<char> Data;
90   Data.reserve(TheHeader.Size);
91   raw_svector_ostream OS(Data);
92   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
93   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
94   for (auto &KeyAndValue : OffloadingData.StringData) {
95     uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
96     StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.getKey()),
97                     Offset + StrTab.getOffset(KeyAndValue.getValue())};
98     OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
99   }
100   StrTab.write(OS);
101   // Add padding to required image alignment.
102   OS.write_zeros(TheEntry.ImageOffset - OS.tell());
103   OS << OffloadingData.Image->getBuffer();
104 
105   // Add final padding to required alignment.
106   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
107   OS.write_zeros(TheHeader.Size - OS.tell());
108   assert(TheHeader.Size == OS.tell() && "Size mismatch");
109 
110   return MemoryBuffer::getMemBufferCopy(OS.str());
111 }
112 
113 OffloadKind object::getOffloadKind(StringRef Name) {
114   return llvm::StringSwitch<OffloadKind>(Name)
115       .Case("openmp", OFK_OpenMP)
116       .Case("cuda", OFK_Cuda)
117       .Case("hip", OFK_HIP)
118       .Default(OFK_None);
119 }
120 
121 StringRef object::getOffloadKindName(OffloadKind Kind) {
122   switch (Kind) {
123   case OFK_OpenMP:
124     return "openmp";
125   case OFK_Cuda:
126     return "cuda";
127   case OFK_HIP:
128     return "hip";
129   default:
130     return "none";
131   }
132 }
133 
134 ImageKind object::getImageKind(StringRef Name) {
135   return llvm::StringSwitch<ImageKind>(Name)
136       .Case("o", IMG_Object)
137       .Case("bc", IMG_Bitcode)
138       .Case("cubin", IMG_Cubin)
139       .Case("fatbin", IMG_Fatbinary)
140       .Case("s", IMG_PTX)
141       .Default(IMG_None);
142 }
143 
144 StringRef object::getImageKindName(ImageKind Kind) {
145   switch (Kind) {
146   case IMG_Object:
147     return "o";
148   case IMG_Bitcode:
149     return "bc";
150   case IMG_Cubin:
151     return "cubin";
152   case IMG_Fatbinary:
153     return "fatbin";
154   case IMG_PTX:
155     return "s";
156   default:
157     return "";
158   }
159 }
160