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