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