diff --git a/clang/test/Driver/clang-offload-bundler.c b/clang/test/Driver/clang-offload-bundler.c --- a/clang/test/Driver/clang-offload-bundler.c +++ b/clang/test/Driver/clang-offload-bundler.c @@ -23,6 +23,7 @@ // // RUN: echo 'Content of device file 1' > %t.tgt1 // RUN: echo 'Content of device file 2' > %t.tgt2 +// RUN: echo 'Content of device file 3' > %t.tgt3 // // Check help message. @@ -91,6 +92,10 @@ // CK-ERR9A: error: expecting exactly one host target but got 0 // CK-ERR9B: error: Duplicate targets are not allowed +// RUN: not clang-offload-bundler -type=a -targets=hxst-powerpcxxle-ibm-linux-gnu,openxp-pxxerpc64le-ibm-linux-gnu,xpenmp-x86_xx-pc-linux-gnu -inputs=%t.i,%t.tgt1,%t.tgt2 -outputs=%t.bundle.i 2>&1 | FileCheck %s --check-prefix CK-ERR10A + +// CK-ERR10A: error: Archive files are only supported for unbundling + // // Check text bundle. This is a readable format, so we check for the format we expect to find. // @@ -319,6 +324,42 @@ // RUN: 2>&1 | FileCheck -check-prefix=DUP %s // DUP: error: Duplicate targets are not allowed +// COM: Archive unbundling test #1: +// COM: The input is an archive of bundled object file and the output should be +// COM: an archive of device-specific files extracted from the bundled object +// COM: files in the archive. + +// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.o,%t.tgt1 -outputs=%t.abundle1.o +// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.o,%t.tgt2 -outputs=%t.abundle2.o +// RUN: llvm-ar cr %t.lib.a %t.abundle1.o %t.abundle2.o +// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx900 -inputs=%t.lib.a -outputs=%t.devicelib.a +// RUN: llvm-ar t %t.devicelib.a | FileCheck %s +// CHECK: abundle1.bc +// CHECK: abundle2.bc + +// COM: Create an archive (*-archive.lib.a) containing two bundles of +// COM: device 1 (gfx903) and one bundle of device 2 (gfx906) + +// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t.o,%t.tgt1 -outputs=%t.bundle-1-gfx903.o +// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx906 -inputs=%t.o,%t.tgt1 -outputs=%t.bundle-1-gfx906.o +// RUN: clang-offload-bundler -type=o -targets=host-%itanium_abi_triple,openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t.o,%t.tgt2 -outputs=%t.bundle-2-gfx903.o +// RUN: llvm-ar cr %t-archive.lib.a %t.bundle-1-gfx903.o %t.bundle-1-gfx906.o %t.bundle-2-gfx903.o + +// COM: Unbundle archive.lib.a to create a device sepcific archive for device 1 +// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx903 -inputs=%t-archive.lib.a -outputs=%t-archive.lib-gfx903.a +// RUN: llvm-ar t %t-archive.lib-gfx903.a | FileCheck %s -check-prefix=GFX903 +// GFX903: bundle-1-gfx903.bc +// GFX903: bundle-2-gfx903.bc +// GFX903-NOT: bundle-1-gfx906.bc + +// COM: Multi target mode of archive unbundling +// RUN: clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx903,openmp-amdgcn-amd-amdhsa-gfx906 -inputs=%t-archive.lib.a -outputs=%t-marchive.lib-gfx903.a,%t-marchive.lib-gfx906.a +// RUN: llvm-ar t %t-marchive.lib-gfx906.a | FileCheck %s -check-prefix=GFX906a +// RUN: llvm-ar t %t-marchive.lib-gfx903.a | FileCheck %s -check-prefix=GFX903a +// GFX906a: bundle-1-gfx906.bc +// GFX903a: bundle-1-gfx903.bc +// GFX903a: bundle-2-gfx903.bc + // Some code so that we can create a binary out of this file. int A = 0; void test_func(void) { diff --git a/clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp b/clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp --- a/clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp +++ b/clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp @@ -22,6 +22,8 @@ #include "llvm/ADT/StringRef.h" #include "llvm/ADT/StringSwitch.h" #include "llvm/ADT/Triple.h" +#include "llvm/Object/Archive.h" +#include "llvm/Object/ArchiveWriter.h" #include "llvm/Object/Binary.h" #include "llvm/Object/ObjectFile.h" #include "llvm/Support/Casting.h" @@ -30,6 +32,7 @@ #include "llvm/Support/Error.h" #include "llvm/Support/ErrorOr.h" #include "llvm/Support/FileSystem.h" +#include "llvm/Support/Host.h" #include "llvm/Support/MemoryBuffer.h" #include "llvm/Support/Path.h" #include "llvm/Support/Program.h" @@ -82,6 +85,7 @@ " bc - llvm-bc\n" " s - assembler\n" " o - object\n" + " a - archive of objects\n" " gch - precompiled-header\n" " ast - clang AST file"), cl::cat(ClangOffloadBundlerCategory)); @@ -152,7 +156,7 @@ virtual Error ReadBundleEnd(MemoryBuffer &Input) = 0; /// Read the current bundle and write the result into the stream \a OS. - virtual Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) = 0; + virtual Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) = 0; /// Write the header of the bundled file to \a OS based on the information /// gathered from \a Inputs. @@ -324,7 +328,7 @@ return Error::success(); } - Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final { + Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final { assert(CurBundleInfo != BundlesInfo.end() && "Invalid reader info!"); StringRef FC = Input.getBuffer(); OS.write(FC.data() + CurBundleInfo->second.Offset, @@ -487,7 +491,7 @@ Error ReadBundleEnd(MemoryBuffer &Input) final { return Error::success(); } - Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final { + Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final { Expected ContentOrErr = CurrentSection->getContents(); if (!ContentOrErr) return ContentOrErr.takeError(); @@ -681,7 +685,7 @@ return Error::success(); } - Error ReadBundle(raw_fd_ostream &OS, MemoryBuffer &Input) final { + Error ReadBundle(raw_ostream &OS, MemoryBuffer &Input) final { StringRef FC = Input.getBuffer(); size_t BundleStart = ReadChars; @@ -764,6 +768,8 @@ return std::make_unique(/*Comment=*/"#"); if (FilesType == "o") return CreateObjectFileHandler(FirstInput); + if (FilesType == "a") + return CreateObjectFileHandler(FirstInput); if (FilesType == "gch") return std::make_unique(); if (FilesType == "ast") @@ -942,6 +948,171 @@ return Error::success(); } +static Archive::Kind getDefaultArchiveKindForHost() { + return Triple(sys::getDefaultTargetTriple()).isOSDarwin() ? Archive::K_DARWIN + : Archive::K_GNU; +} + +/// Returns arch specific extension for device object files in the generated +/// device specific archive in archive unbundling mode +static StringRef getDeviceFileExtension(StringRef Device) { + if (Device.contains("gfx")) // for amdgcn + return "bc"; + else if (Device.contains("sm_")) // for nvptx + return "cubin"; + return "o"; +} + +static bool isValidOffloadTarget(StringRef ArchiveDevice) { + for (auto &Target : TargetNames) + if (Target == ArchiveDevice) + return true; + return false; +} + +/// UnbundleArchive takes an archive file (".a") as input containing bundled +/// object files, and a list of offload targets (not host), and extracts the +/// device code into a new archive file for each offload target. Each resulting +/// archive file contains all device object files corresponding to that +/// particular offload target device. Names of object files in the archive are +/// preserved but their extension is changed according to the target +/// architecture (either ".cubin" or ".bc"). The created archive file does not +/// contain an index of the symbols. +static Error UnbundleArchive() { + std::vector> ArchiveBuffers; + std::vector> LibraryNames; + + /// Map of target names with list of object files that will form the device + /// specific archive for that target + StringMap> ArchiveMembersMap; + + // Map of target names and output archive filenames + StringMap TargetOutputFileNameMap; + auto Output = OutputFileNames.begin(); + for (auto &Target : TargetNames) { + TargetOutputFileNameMap[Target] = *Output; + ++Output; + } + + StringRef IFName = InputFileNames.front(); + ErrorOr> BufOrErr = + MemoryBuffer::getFileOrSTDIN(IFName, -1, false); + if (std::error_code EC = BufOrErr.getError()) + return createFileError(InputFileNames.front(), EC); + + ArchiveBuffers.push_back(std::move(*BufOrErr)); + Expected> LibOrErr = + Archive::create(ArchiveBuffers.back()->getMemBufferRef()); + if (!LibOrErr) + return LibOrErr.takeError(); + + auto Archive = std::move(*LibOrErr); + + Error ArchiveErr = Error::success(); + auto ChildEnd = Archive->child_end(); + for (auto ChildIter = Archive->child_begin(ArchiveErr); ChildIter != ChildEnd; + ++ChildIter) { + if (ArchiveErr) + return ArchiveErr; + auto ChildNameOrErr = (*ChildIter).getName(); + if (!ChildNameOrErr) + return ChildNameOrErr.takeError(); + + StringRef ChildName = sys::path::filename(*ChildNameOrErr); + + auto ChildBufferRefOrErr = (*ChildIter).getMemoryBufferRef(); + if (!ChildBufferRefOrErr) + return ChildBufferRefOrErr.takeError(); + + auto ChildBuffer = MemoryBuffer::getMemBuffer(*ChildBufferRefOrErr, false); + + Expected> FileHandlerOrErr = + CreateFileHandler(*ChildBuffer); + if (!FileHandlerOrErr) + return FileHandlerOrErr.takeError(); + + std::unique_ptr &FileHandler = *FileHandlerOrErr; + assert(FileHandler); + + if (Error ReadErr = FileHandler.get()->ReadHeader(*ChildBuffer)) + return ReadErr; + + Expected> CurTripleOrErr = + FileHandler->ReadBundleStart(*ChildBuffer); + if (!CurTripleOrErr) + return CurTripleOrErr.takeError(); + + Optional OptionalCurKindTriple = *CurTripleOrErr; + // No device code in this child, skip + if (!OptionalCurKindTriple.hasValue()) + continue; + StringRef CurOffloadTargetDevice = *OptionalCurKindTriple; + assert(!CurOffloadTargetDevice.empty()); + + while (!CurOffloadTargetDevice.empty()) { + if (hasHostKind(CurOffloadTargetDevice)) { + // Do nothing, we don't extract host code yet + } else if (isValidOffloadTarget(CurOffloadTargetDevice)) { + std::string BundleData; + raw_string_ostream DataStream(BundleData); + if (Error Err = FileHandler.get()->ReadBundle(DataStream, *ChildBuffer)) + return Err; + + SmallString<128> StatsFile; + StatsFile.assign(ChildName); + llvm::sys::path::replace_extension( + StatsFile, getDeviceFileExtension(CurOffloadTargetDevice)); + LibraryNames.push_back( + std::unique_ptr(new std::string(StatsFile.c_str()))); + + auto MemBuf = MemoryBuffer::getMemBufferCopy( + DataStream.str(), *(LibraryNames.back().get())); + ArchiveBuffers.push_back(std::move(MemBuf)); + auto MemBufRef = MemoryBufferRef(*(ArchiveBuffers.back())); + + // Add this archive into ArchiveMembers vector of CurOffloadTargetDevice + if (ArchiveMembersMap.find(CurOffloadTargetDevice) == + ArchiveMembersMap.end()) { + std::vector ArchiveMembers; + ArchiveMembers.push_back(NewArchiveMember(MemBufRef)); + ArchiveMembersMap.insert_or_assign(CurOffloadTargetDevice, + std::move(ArchiveMembers)); + } else { + ArchiveMembersMap[CurOffloadTargetDevice].push_back( + NewArchiveMember(MemBufRef)); + } + } + + if (Error Err = FileHandler.get()->ReadBundleEnd(*ChildBuffer)) + return Err; + + Expected> NextTripleOrErr = + FileHandler->ReadBundleStart(*ChildBuffer); + if (!NextTripleOrErr) + return NextTripleOrErr.takeError(); + + CurOffloadTargetDevice = + ((*NextTripleOrErr).hasValue()) ? **NextTripleOrErr : ""; + } + } + assert(!ArchiveErr); + + /// Write out an archive for each target + for (auto &Target : TargetNames) { + StringRef FileName = TargetOutputFileNameMap[Target]; + StringMapIterator> CurArchiveMembers = + ArchiveMembersMap.find(Target); + if (CurArchiveMembers != ArchiveMembersMap.end()) { + if (Error WriteErr = writeArchive(FileName, CurArchiveMembers->getValue(), + true, getDefaultArchiveKindForHost(), + true, false, nullptr)) + return WriteErr; + } + } + + return Error::success(); +} + static void PrintVersion(raw_ostream &OS) { OS << clang::getClangToolFullVersion("clang-offload-bundler") << '\n'; } @@ -982,6 +1153,12 @@ "match in unbundling mode")); } } else { + if (FilesType == "a") { + reportError(createStringError(errc::invalid_argument, + "Archive files are only supported " + "for unbundling")); + Error = true; + } if (OutputFileNames.size() != 1) { Error = true; reportError(createStringError( @@ -1064,7 +1241,11 @@ if (!llvm::sys::fs::exists(BundlerExecutable)) BundlerExecutable = sys::fs::getMainExecutable(argv[0], &BundlerExecutable); - if (llvm::Error Err = Unbundle ? UnbundleFiles() : BundleFiles()) { + llvm::Error Err = (Unbundle && FilesType == "a") ? UnbundleArchive() + : (Unbundle) ? UnbundleFiles() + : BundleFiles(); + + if (Err) { reportError(std::move(Err)); return 1; }