llvm/lib/Object/OffloadBinary.cpp

   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 SmallString<0> OffloadBinary::write(const OffloadingImage &OffloadingData) {
 208   // Create a null-terminated string table with all the used strings.
 209   StringTableBuilder StrTab(StringTableBuilder::ELF);
 210   for (auto &KeyAndValue : OffloadingData.StringData) {
 211     StrTab.add(KeyAndValue.first);
 212     StrTab.add(KeyAndValue.second);
 213   }
 214   StrTab.finalize();
 215
 216   uint64_t StringEntrySize =
 217       sizeof(StringEntry) * OffloadingData.StringData.size();
 218
 219   // Make sure the image we're wrapping around is aligned as well.
 220   uint64_t BinaryDataSize = alignTo(sizeof(Header) + sizeof(Entry) +
 221                                         StringEntrySize + StrTab.getSize(),
 222                                     getAlignment());
 223
 224   // Create the header and fill in the offsets. The entry will be directly
 225   // placed after the header in memory. Align the size to the alignment of the
 226   // header so this can be placed contiguously in a single section.
 227   Header TheHeader;
 228   TheHeader.Size = alignTo(
 229       BinaryDataSize + OffloadingData.Image->getBufferSize(), getAlignment());
 230   TheHeader.EntryOffset = sizeof(Header);
 231   TheHeader.EntrySize = sizeof(Entry);
 232
 233   // Create the entry using the string table offsets. The string table will be
 234   // placed directly after the entry in memory, and the image after that.
 235   Entry TheEntry;
 236   TheEntry.TheImageKind = OffloadingData.TheImageKind;
 237   TheEntry.TheOffloadKind = OffloadingData.TheOffloadKind;
 238   TheEntry.Flags = OffloadingData.Flags;
 239   TheEntry.StringOffset = sizeof(Header) + sizeof(Entry);
 240   TheEntry.NumStrings = OffloadingData.StringData.size();
 241
 242   TheEntry.ImageOffset = BinaryDataSize;
 243   TheEntry.ImageSize = OffloadingData.Image->getBufferSize();
 244
 245   SmallString<0> Data;
 246   Data.reserve(TheHeader.Size);
 247   raw_svector_ostream OS(Data);
 248   OS << StringRef(reinterpret_cast<char *>(&TheHeader), sizeof(Header));
 249   OS << StringRef(reinterpret_cast<char *>(&TheEntry), sizeof(Entry));
 250   for (auto &KeyAndValue : OffloadingData.StringData) {
 251     uint64_t Offset = sizeof(Header) + sizeof(Entry) + StringEntrySize;
 252     StringEntry Map{Offset + StrTab.getOffset(KeyAndValue.first),
 253                     Offset + StrTab.getOffset(KeyAndValue.second)};
 254     OS << StringRef(reinterpret_cast<char *>(&Map), sizeof(StringEntry));
 255   }
 256   StrTab.write(OS);
 257   // Add padding to required image alignment.
 258   OS.write_zeros(TheEntry.ImageOffset - OS.tell());
 259   OS << OffloadingData.Image->getBuffer();
 260
 261   // Add final padding to required alignment.
 262   assert(TheHeader.Size >= OS.tell() && "Too much data written?");
 263   OS.write_zeros(TheHeader.Size - OS.tell());
 264   assert(TheHeader.Size == OS.tell() && "Size mismatch");
 265
 266   return Data;
 267 }
 268
 269 Error object::extractOffloadBinaries(MemoryBufferRef Buffer,
 270                                      SmallVectorImpl<OffloadFile> &Binaries) {
 271   file_magic Type = identify_magic(Buffer.getBuffer());
 272   switch (Type) {
 273   case file_magic::bitcode:
 274     return extractFromBitcode(Buffer, Binaries);
 275   case file_magic::elf_relocatable:
 276   case file_magic::elf_executable:
 277   case file_magic::elf_shared_object:
 278   case file_magic::coff_object: {
 279     Expected<std::unique_ptr<ObjectFile>> ObjFile =
 280         ObjectFile::createObjectFile(Buffer, Type);
 281     if (!ObjFile)
 282       return ObjFile.takeError();
 283     return extractFromObject(*ObjFile->get(), Binaries);
 284   }
 285   case file_magic::archive: {
 286     Expected<std::unique_ptr<llvm::object::Archive>> LibFile =
 287         object::Archive::create(Buffer);
 288     if (!LibFile)
 289       return LibFile.takeError();
 290     return extractFromArchive(*LibFile->get(), Binaries);
 291   }
 292   case file_magic::offload_binary:
 293     return extractOffloadFiles(Buffer, Binaries);
 294   default:
 295     return Error::success();
 296   }
 297 }
 298
 299 OffloadKind object::getOffloadKind(StringRef Name) {
 300   return llvm::StringSwitch<OffloadKind>(Name)
 301       .Case("openmp", OFK_OpenMP)
 302       .Case("cuda", OFK_Cuda)
 303       .Case("hip", OFK_HIP)
 304       .Default(OFK_None);
 305 }
 306
 307 StringRef object::getOffloadKindName(OffloadKind Kind) {
 308   switch (Kind) {
 309   case OFK_OpenMP:
 310     return "openmp";
 311   case OFK_Cuda:
 312     return "cuda";
 313   case OFK_HIP:
 314     return "hip";
 315   default:
 316     return "none";
 317   }
 318 }
 319
 320 ImageKind object::getImageKind(StringRef Name) {
 321   return llvm::StringSwitch<ImageKind>(Name)
 322       .Case("o", IMG_Object)
 323       .Case("bc", IMG_Bitcode)
 324       .Case("cubin", IMG_Cubin)
 325       .Case("fatbin", IMG_Fatbinary)
 326       .Case("s", IMG_PTX)
 327       .Default(IMG_None);
 328 }
 329
 330 StringRef object::getImageKindName(ImageKind Kind) {
 331   switch (Kind) {
 332   case IMG_Object:
 333     return "o";
 334   case IMG_Bitcode:
 335     return "bc";
 336   case IMG_Cubin:
 337     return "cubin";
 338   case IMG_Fatbinary:
 339     return "fatbin";
 340   case IMG_PTX:
 341     return "s";
 342   default:
 343     return "";
 344   }
 345 }