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