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