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>>
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>
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 
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 
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 
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 
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