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/IR/Constants.h" 14 #include "llvm/IR/Module.h" 15 #include "llvm/IRReader/IRReader.h" 16 #include "llvm/MC/StringTableBuilder.h" 17 #include "llvm/Object/Archive.h" 18 #include "llvm/Object/ArchiveWriter.h" 19 #include "llvm/Object/Binary.h" 20 #include "llvm/Object/COFF.h" 21 #include "llvm/Object/ELFObjectFile.h" 22 #include "llvm/Object/Error.h" 23 #include "llvm/Object/IRObjectFile.h" 24 #include "llvm/Object/ObjectFile.h" 25 #include "llvm/Support/Alignment.h" 26 #include "llvm/Support/FileOutputBuffer.h" 27 #include "llvm/Support/SourceMgr.h" 28 29 using namespace llvm; 30 using namespace llvm::object; 31 32 namespace { 33 34 /// Attempts to extract all the embedded device images contained inside the 35 /// buffer \p Contents. The buffer is expected to contain a valid offloading 36 /// binary format. 37 Error extractOffloadFiles(MemoryBufferRef Contents, 38 SmallVectorImpl<OffloadFile> &Binaries) { 39 uint64_t Offset = 0; 40 // There could be multiple offloading binaries stored at this section. 41 while (Offset < Contents.getBuffer().size()) { 42 std::unique_ptr<MemoryBuffer> Buffer = 43 MemoryBuffer::getMemBuffer(Contents.getBuffer().drop_front(Offset), "", 44 /*RequiresNullTerminator*/ false); 45 if (!isAddrAligned(Align(OffloadBinary::getAlignment()), 46 Buffer->getBufferStart())) 47 Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(), 48 Buffer->getBufferIdentifier()); 49 auto BinaryOrErr = OffloadBinary::create(*Buffer); 50 if (!BinaryOrErr) 51 return BinaryOrErr.takeError(); 52 OffloadBinary &Binary = **BinaryOrErr; 53 54 // Create a new owned binary with a copy of the original memory. 55 std::unique_ptr<MemoryBuffer> BufferCopy = MemoryBuffer::getMemBufferCopy( 56 Binary.getData().take_front(Binary.getSize()), 57 Contents.getBufferIdentifier()); 58 auto NewBinaryOrErr = OffloadBinary::create(*BufferCopy); 59 if (!NewBinaryOrErr) 60 return NewBinaryOrErr.takeError(); 61 Binaries.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy)); 62 63 Offset += Binary.getSize(); 64 } 65 66 return Error::success(); 67 } 68 69 // Extract offloading binaries from an Object file \p Obj. 70 Error extractFromObject(const ObjectFile &Obj, 71 SmallVectorImpl<OffloadFile> &Binaries) { 72 assert((Obj.isELF() || Obj.isCOFF()) && "Invalid file type"); 73 74 for (SectionRef Sec : Obj.sections()) { 75 // ELF files contain a section with the LLVM_OFFLOADING type. 76 if (Obj.isELF() && 77 static_cast<ELFSectionRef>(Sec).getType() != ELF::SHT_LLVM_OFFLOADING) 78 continue; 79 80 // COFF has no section types so we rely on the name of the section. 81 if (Obj.isCOFF()) { 82 Expected<StringRef> NameOrErr = Sec.getName(); 83 if (!NameOrErr) 84 return NameOrErr.takeError(); 85 86 if (!NameOrErr->equals(".llvm.offloading")) 87 continue; 88 } 89 90 Expected<StringRef> Buffer = Sec.getContents(); 91 if (!Buffer) 92 return Buffer.takeError(); 93 94 MemoryBufferRef Contents(*Buffer, Obj.getFileName()); 95 if (Error Err = extractOffloadFiles(Contents, Binaries)) 96 return Err; 97 } 98 99 return Error::success(); 100 } 101 102 Error extractFromBitcode(MemoryBufferRef Buffer, 103 SmallVectorImpl<OffloadFile> &Binaries) { 104 LLVMContext Context; 105 SMDiagnostic Err; 106 std::unique_ptr<Module> M = getLazyIRModule( 107 MemoryBuffer::getMemBuffer(Buffer, /*RequiresNullTerminator=*/false), Err, 108 Context); 109 if (!M) 110 return createStringError(inconvertibleErrorCode(), 111 "Failed to create module"); 112 113 // Extract offloading data from globals referenced by the 114 // `llvm.embedded.object` metadata with the `.llvm.offloading` section. 115 auto *MD = M->getNamedMetadata("llvm.embedded.objects"); 116 if (!MD) 117 return Error::success(); 118 119 for (const MDNode *Op : MD->operands()) { 120 if (Op->getNumOperands() < 2) 121 continue; 122 123 MDString *SectionID = dyn_cast<MDString>(Op->getOperand(1)); 124 if (!SectionID || SectionID->getString() != ".llvm.offloading") 125 continue; 126 127 GlobalVariable *GV = 128 mdconst::dyn_extract_or_null<GlobalVariable>(Op->getOperand(0)); 129 if (!GV) 130 continue; 131 132 auto *CDS = dyn_cast<ConstantDataSequential>(GV->getInitializer()); 133 if (!CDS) 134 continue; 135 136 MemoryBufferRef Contents(CDS->getAsString(), M->getName()); 137 if (Error Err = extractOffloadFiles(Contents, Binaries)) 138 return Err; 139 } 140 141 return Error::success(); 142 } 143 144 Error extractFromArchive(const Archive &Library, 145 SmallVectorImpl<OffloadFile> &Binaries) { 146 // Try to extract device code from each file stored in the static archive. 147 Error Err = Error::success(); 148 for (auto Child : Library.children(Err)) { 149 auto ChildBufferOrErr = Child.getMemoryBufferRef(); 150 if (!ChildBufferOrErr) 151 return ChildBufferOrErr.takeError(); 152 std::unique_ptr<MemoryBuffer> ChildBuffer = 153 MemoryBuffer::getMemBuffer(*ChildBufferOrErr, false); 154 155 // Check if the buffer has the required alignment. 156 if (!isAddrAligned(Align(OffloadBinary::getAlignment()), 157 ChildBuffer->getBufferStart())) 158 ChildBuffer = MemoryBuffer::getMemBufferCopy( 159 ChildBufferOrErr->getBuffer(), 160 ChildBufferOrErr->getBufferIdentifier()); 161 162 if (Error Err = extractOffloadBinaries(*ChildBuffer, Binaries)) 163 return Err; 164 } 165 166 if (Err) 167 return Err; 168 return Error::success(); 169 } 170 171 } // namespace 172 173 Expected<std::unique_ptr<OffloadBinary>> 174 OffloadBinary::create(MemoryBufferRef Buf) { 175 if (Buf.getBufferSize() < sizeof(Header) + sizeof(Entry)) 176 return errorCodeToError(object_error::parse_failed); 177 178 // Check for 0x10FF1OAD magic bytes. 179 if (identify_magic(Buf.getBuffer()) != file_magic::offload_binary) 180 return errorCodeToError(object_error::parse_failed); 181 182 // Make sure that the data has sufficient alignment. 183 if (!isAddrAligned(Align(getAlignment()), Buf.getBufferStart())) 184 return errorCodeToError(object_error::parse_failed); 185 186 const char *Start = Buf.getBufferStart(); 187 const Header *TheHeader = reinterpret_cast<const Header *>(Start); 188 if (TheHeader->Version != OffloadBinary::Version) 189 return errorCodeToError(object_error::parse_failed); 190 191 if (TheHeader->Size > Buf.getBufferSize() || 192 TheHeader->EntryOffset > TheHeader->Size - sizeof(Entry) || 193 TheHeader->EntrySize > TheHeader->Size - sizeof(Header)) 194 return errorCodeToError(object_error::unexpected_eof); 195 196 const Entry *TheEntry = 197 reinterpret_cast<const Entry *>(&Start[TheHeader->EntryOffset]); 198 199 if (TheEntry->ImageOffset > Buf.getBufferSize() || 200 TheEntry->StringOffset > Buf.getBufferSize()) 201 return errorCodeToError(object_error::unexpected_eof); 202 203 return std::unique_ptr<OffloadBinary>( 204 new OffloadBinary(Buf, TheHeader, TheEntry)); 205 } 206 207 std::unique_ptr<MemoryBuffer> 208 OffloadBinary::write(const OffloadingImage &OffloadingData) { 209 // Create a null-terminated string table with all the used strings. 210 StringTableBuilder StrTab(StringTableBuilder::ELF); 211 for (auto &KeyAndValue : OffloadingData.StringData) { 212 StrTab.add(KeyAndValue.first); 213 StrTab.add(KeyAndValue.second); 214 } 215 StrTab.finalize(); 216 217 uint64_t StringEntrySize = 218 sizeof(StringEntry) * OffloadingData.StringData.size(); 219 220 // Make sure the image we're wrapping around is aligned as well. 221 uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) + 222 StringEntrySize + StrTab.getSize(), 223 getAlignment()); 224 225 // Create the header and fill in the offsets. The entry will be directly 226 // placed after the header in memory. Align the size to the alignment of the 227 // header so this can be placed contiguously in a single section. 228 Header TheHeader; 229 TheHeader.Size = alignTo( 230 BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment()); 231 TheHeader.EntryOffset = sizeof(Header); 232 TheHeader.EntrySize = sizeof(Entry); 233 234 // Create the entry using the string table offsets. The string table will be 235 // placed directly after the entry in memory, and the image after that. 236 Entry TheEntry; 237 TheEntry.TheImageKind = OffloadingData.TheImageKind; 238 TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind; 239 TheEntry.Flags = OffloadingData.Flags; 240 TheEntry.StringOffset = sizeof(Header) + sizeof(Entry); 241 TheEntry.NumStrings = OffloadingData.StringData.size(); 242 243 TheEntry.ImageOffset = BinaryDataSize; 244 TheEntry.ImageSize = OffloadingData.Image->getBufferSize(); 245 246 SmallVector<char> Data; 247 Data.reserve(TheHeader.Size); 248 raw_svector_ostream OS(Data); 249 OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header)); 250 OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry)); 251 for (auto &KeyAndValue : OffloadingData.StringData) { 252 uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize; 253 StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.first), 254 Offset + StrTab.getOffset(KeyAndValue.second)}; 255 OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry)); 256 } 257 StrTab.write(OS); 258 // Add padding to required image alignment. 259 OS.write_zeros(TheEntry.ImageOffset - OS.tell()); 260 OS << OffloadingData.Image->getBuffer(); 261 262 // Add final padding to required alignment. 263 assert(TheHeader.Size >= OS.tell() && "Too much data written?"); 264 OS.write_zeros(TheHeader.Size - OS.tell()); 265 assert(TheHeader.Size == OS.tell() && "Size mismatch"); 266 267 return MemoryBuffer::getMemBufferCopy(OS.str()); 268 } 269 270 Error object::extractOffloadBinaries(MemoryBufferRef Buffer, 271 SmallVectorImpl<OffloadFile> &Binaries) { 272 file_magic Type = identify_magic(Buffer.getBuffer()); 273 switch (Type) { 274 case file_magic::bitcode: 275 return extractFromBitcode(Buffer, Binaries); 276 case file_magic::elf_relocatable: 277 case file_magic::elf_executable: 278 case file_magic::elf_shared_object: 279 case file_magic::coff_object: { 280 Expected<std::unique_ptr<ObjectFile>> ObjFile = 281 ObjectFile::createObjectFile(Buffer, Type); 282 if (!ObjFile) 283 return ObjFile.takeError(); 284 return extractFromObject(*ObjFile->get(), Binaries); 285 } 286 case file_magic::archive: { 287 Expected<std::unique_ptr<llvm::object::Archive>> LibFile = 288 object::Archive::create(Buffer); 289 if (!LibFile) 290 return LibFile.takeError(); 291 return extractFromArchive(*LibFile->get(), Binaries); 292 } 293 case file_magic::offload_binary: 294 return extractOffloadFiles(Buffer, Binaries); 295 default: 296 return Error::success(); 297 } 298 } 299 300 OffloadKind object::getOffloadKind(StringRef Name) { 301 return llvm::StringSwitch<OffloadKind>(Name) 302 .Case("openmp", OFK_OpenMP) 303 .Case("cuda", OFK_Cuda) 304 .Case("hip", OFK_HIP) 305 .Default(OFK_None); 306 } 307 308 StringRef object::getOffloadKindName(OffloadKind Kind) { 309 switch (Kind) { 310 case OFK_OpenMP: 311 return "openmp"; 312 case OFK_Cuda: 313 return "cuda"; 314 case OFK_HIP: 315 return "hip"; 316 default: 317 return "none"; 318 } 319 } 320 321 ImageKind object::getImageKind(StringRef Name) { 322 return llvm::StringSwitch<ImageKind>(Name) 323 .Case("o", IMG_Object) 324 .Case("bc", IMG_Bitcode) 325 .Case("cubin", IMG_Cubin) 326 .Case("fatbin", IMG_Fatbinary) 327 .Case("s", IMG_PTX) 328 .Default(IMG_None); 329 } 330 331 StringRef object::getImageKindName(ImageKind Kind) { 332 switch (Kind) { 333 case IMG_Object: 334 return "o"; 335 case IMG_Bitcode: 336 return "bc"; 337 case IMG_Cubin: 338 return "cubin"; 339 case IMG_Fatbinary: 340 return "fatbin"; 341 case IMG_PTX: 342 return "s"; 343 default: 344 return ""; 345 } 346 } 347