diff --git a/clang/test/Driver/offload-packager.c b/clang/test/Driver/offload-packager.c --- a/clang/test/Driver/offload-packager.c +++ b/clang/test/Driver/offload-packager.c @@ -29,3 +29,25 @@ // RUN: diff *-amdgcn-amd-amdhsa-gfx908.2.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx908.2.o // RUN: diff *-amdgcn-amd-amdhsa-gfx90a.3.o %S/Inputs/dummy-elf.o; rm *-amdgcn-amd-amdhsa-gfx90a.3.o // RUN: not diff *-amdgcn-amd-amdhsa-gfx90c.4.o %S/Inputs/dummy-elf.o + +// Check that we can extract from an ELF object file +// RUN: clang-offload-packager -o %t.out \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 +// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-obj -o %t.o -fembed-offload-object=%t.out +// RUN: clang-offload-packager %t.out \ +// RUN: --image=file=%t-sm_70.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \ +// RUN: --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 +// RUN: diff %t-sm_70.o %S/Inputs/dummy-elf.o +// RUN: diff %t-gfx908.o %S/Inputs/dummy-elf.o + +// Check that we can extract from a bitcode file +// RUN: clang-offload-packager -o %t.out \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 \ +// RUN: --image=file=%S/Inputs/dummy-elf.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 +// RUN: %clang -cc1 %s -triple x86_64-unknown-linux-gnu -emit-llvm -o %t.bc -fembed-offload-object=%t.out +// RUN: clang-offload-packager %t.out \ +// RUN: --image=file=%t-sm_70.o,kind=openmp,triple=nvptx64-nvidia-cuda,arch=sm_70 \ +// RUN: --image=file=%t-gfx908.o,kind=openmp,triple=amdgcn-amd-amdhsa,arch=gfx908 +// RUN: diff %t-sm_70.o %S/Inputs/dummy-elf.o +// RUN: diff %t-gfx908.o %S/Inputs/dummy-elf.o diff --git a/clang/tools/clang-offload-packager/CMakeLists.txt b/clang/tools/clang-offload-packager/CMakeLists.txt --- a/clang/tools/clang-offload-packager/CMakeLists.txt +++ b/clang/tools/clang-offload-packager/CMakeLists.txt @@ -1,6 +1,9 @@ set(LLVM_LINK_COMPONENTS ${LLVM_TARGETS_TO_BUILD} BinaryFormat + BitWriter + Core + IRReader Object Support) diff --git a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp --- a/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp +++ b/clang/tools/clang-offload-packager/ClangOffloadPackager.cpp @@ -14,7 +14,14 @@ #include "clang/Basic/Version.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Module.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Object/Archive.h" +#include "llvm/Object/ArchiveWriter.h" #include "llvm/Object/Binary.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Object/IRObjectFile.h" #include "llvm/Object/ObjectFile.h" #include "llvm/Object/OffloadBinary.h" #include "llvm/Support/CommandLine.h" @@ -23,6 +30,7 @@ #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/Path.h" #include "llvm/Support/Signals.h" +#include "llvm/Support/SourceMgr.h" #include "llvm/Support/StringSaver.h" #include "llvm/Support/WithColor.h" @@ -54,6 +62,12 @@ OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n'; } +/// A single OffloadBinary and its associated memory. +using OffloadFile = OwningBinary; + +Error extractFromBuffer(std::unique_ptr Buffer, + SmallVectorImpl &DeviceFiles); + // Get a map containing all the arguments for the image. Repeated arguments will // be placed in a comma separated list. static DenseMap getImageArguments(StringRef Image, @@ -123,12 +137,11 @@ return Error::success(); } -static Expected>> -extractOffloadFiles(MemoryBufferRef Contents) { - if (identify_magic(Contents.getBuffer()) != file_magic::offload_binary) - return createStringError(inconvertibleErrorCode(), - "Input buffer not an offloading binary"); - SmallVector> Binaries; +/// Attempts to extract all the embedded device images contained inside the +/// buffer \p Contents. The buffer is expected to contain a valid offloading +/// binary format. +Error extractOffloadFiles(MemoryBufferRef Contents, + SmallVectorImpl &DeviceFiles) { uint64_t Offset = 0; // There could be multiple offloading binaries stored at this section. while (Offset < Contents.getBuffer().size()) { @@ -138,12 +151,136 @@ auto BinaryOrErr = OffloadBinary::create(*Buffer); if (!BinaryOrErr) return BinaryOrErr.takeError(); + OffloadBinary &Binary = **BinaryOrErr; + + // Create a new owned binary with a copy of the original memory. + std::unique_ptr BufferCopy = MemoryBuffer::getMemBufferCopy( + Binary.getData().take_front(Binary.getSize()), + Contents.getBufferIdentifier()); + auto NewBinaryOrErr = OffloadBinary::create(*BufferCopy); + if (!NewBinaryOrErr) + return NewBinaryOrErr.takeError(); + DeviceFiles.emplace_back(std::move(*NewBinaryOrErr), std::move(BufferCopy)); + + Offset += Binary.getSize(); + } + + return Error::success(); +} + +// Extract offloading binaries from an Object file \p Obj. +Error extractFromBinary(const ObjectFile &Obj, + SmallVectorImpl &DeviceFiles) { + for (ELFSectionRef Sec : Obj.sections()) { + if (Sec.getType() != ELF::SHT_LLVM_OFFLOADING) + continue; + + Expected Buffer = Sec.getContents(); + if (!Buffer) + return Buffer.takeError(); + + MemoryBufferRef Contents(*Buffer, Obj.getFileName()); + if (Error Err = extractOffloadFiles(Contents, DeviceFiles)) + return Err; + } + + return Error::success(); +} + +Error extractFromBitcode(std::unique_ptr Buffer, + SmallVectorImpl &DeviceFiles) { + LLVMContext Context; + SMDiagnostic Err; + std::unique_ptr M = getLazyIRModule(std::move(Buffer), Err, Context); + if (!M) + return createStringError(inconvertibleErrorCode(), + "Failed to create module"); + + // Extract offloading data from globals referenced by the + // `llvm.embedded.object` metadata with the `.llvm.offloading` section. + auto *MD = M->getNamedMetadata("llvm.embedded.objects"); + if (!MD) + return Error::success(); + + for (const MDNode *Op : MD->operands()) { + if (Op->getNumOperands() < 2) + continue; + + MDString *SectionID = dyn_cast(Op->getOperand(1)); + if (!SectionID || SectionID->getString() != ".llvm.offloading") + continue; + + GlobalVariable *GV = + mdconst::dyn_extract_or_null(Op->getOperand(0)); + if (!GV) + continue; + + auto *CDS = dyn_cast(GV->getInitializer()); + if (!CDS) + continue; + + MemoryBufferRef Contents(CDS->getAsString(), M->getName()); + if (Error Err = extractOffloadFiles(Contents, DeviceFiles)) + return Err; + } + + return Error::success(); +} - Offset += (*BinaryOrErr)->getSize(); - Binaries.emplace_back(std::move(*BinaryOrErr)); +Error extractFromArchive(const Archive &Library, + SmallVectorImpl &DeviceFiles) { + // Try to extract device code from each file stored in the static archive. + Error Err = Error::success(); + for (auto Child : Library.children(Err)) { + auto ChildBufferOrErr = Child.getMemoryBufferRef(); + if (!ChildBufferOrErr) + return ChildBufferOrErr.takeError(); + std::unique_ptr ChildBuffer = + MemoryBuffer::getMemBuffer(*ChildBufferOrErr, false); + + // Check if the buffer has the required alignment. + if (!isAddrAligned(Align(OffloadBinary::getAlignment()), + ChildBuffer->getBufferStart())) + ChildBuffer = MemoryBuffer::getMemBufferCopy( + ChildBufferOrErr->getBuffer(), + ChildBufferOrErr->getBufferIdentifier()); + + if (Error Err = extractFromBuffer(std::move(ChildBuffer), DeviceFiles)) + return Err; } - return std::move(Binaries); + if (Err) + return Err; + return Error::success(); +} + +/// Extracts embedded device offloading code from a memory \p Buffer to a list +/// of \p DeviceFiles. +Error extractFromBuffer(std::unique_ptr Buffer, + SmallVectorImpl &DeviceFiles) { + file_magic Type = identify_magic(Buffer->getBuffer()); + switch (Type) { + case file_magic::bitcode: + return extractFromBitcode(std::move(Buffer), DeviceFiles); + case file_magic::elf_relocatable: { + Expected> ObjFile = + ObjectFile::createObjectFile(*Buffer, Type); + if (!ObjFile) + return ObjFile.takeError(); + return extractFromBinary(*ObjFile->get(), DeviceFiles); + } + case file_magic::archive: { + Expected> LibFile = + object::Archive::create(*Buffer); + if (!LibFile) + return LibFile.takeError(); + return extractFromArchive(*LibFile->get(), DeviceFiles); + } + case file_magic::offload_binary: + return extractOffloadFiles(*Buffer, DeviceFiles); + default: + return Error::success(); + } } static Error unbundleImages() { @@ -159,9 +296,9 @@ Buffer = MemoryBuffer::getMemBufferCopy(Buffer->getBuffer(), Buffer->getBufferIdentifier()); - auto BinariesOrErr = extractOffloadFiles(*Buffer); - if (!BinariesOrErr) - return BinariesOrErr.takeError(); + SmallVector Binaries; + if (Error Err = extractFromBuffer(std::move(Buffer), Binaries)) + return Err; // Try to extract each device image specified by the user from the input file. for (StringRef Image : DeviceImages) { @@ -169,8 +306,8 @@ StringSaver Saver(Alloc); auto Args = getImageArguments(Image, Saver); - for (uint64_t I = 0, E = BinariesOrErr->size(); I != E; ++I) { - const auto &Binary = (*BinariesOrErr)[I]; + for (uint64_t I = 0, E = Binaries.size(); I != E; ++I) { + const auto *Binary = Binaries[I].getBinary(); // We handle the 'file' and 'kind' identifiers differently. bool Match = llvm::all_of(Args, [&](auto &Arg) { const auto [Key, Value] = Arg;