Skip to content

Fix offload extract #10257

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Jul 28, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions clang/test/Driver/clang-offload-extract.c
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,8 @@
// Extract target images (deprecated use model)
//
// RUN: clang-offload-extract --output=%t.deprecated %t.fat.bin | FileCheck %s --check-prefix CHECK-DEPRECATE
// CHECK-DEPRECATE: Saving target image to
// CHECK-DEPRECATE: Saving target image to
// CHECK-DEPRECATE: Section {{.*}}: Image 1'-> File
// CHECK-DEPRECATE: Section {{.*}}: Image 2'-> File

//
// Check that extracted contents match the original images.
Expand All @@ -47,8 +47,8 @@
// Extract target images (new use model)
//
// RUN: clang-offload-extract --stem=%t.extracted %t.fat.bin | FileCheck %s --check-prefix CHECK-EXTRACT
// CHECK-EXTRACT: Saving target image to
// CHECK-EXTRACT: Saving target image to
// CHECK-EXTRACT: Section {{.*}}: Image 1'-> File
// CHECK-EXTRACT: Section {{.*}}: Image 2'-> File

//
// Check that extracted contents match the original images.
Expand Down
240 changes: 141 additions & 99 deletions clang/tools/clang-offload-extract/ClangOffloadExtract.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,18 +28,24 @@
#include "llvm/Support/raw_ostream.h"

#define IMAGE_INFO_SECTION_NAME ".tgtimg"
#define IMAGE_SECTION_NAME_PREFIX "__CLANG_OFFLOAD_BUNDLE__"
// Windows truncates the names of sections to 8 bytes
#define IMAGE_SECTION_NAME_PREFIX_COFF "__CLANG_"

#define DEBUG_TYPE "clang-offload-extract"

using namespace llvm;
using namespace llvm::object;

// Command-line parsing
// Create a category to label utility-specific options; This will allow
// us to distinguish those specific options from generic options and
// irrelevant options
static cl::OptionCategory
ClangOffloadExtractCategory("Utility-specific options");

// Create options for the input and output files, each with an
// appropriate default when not specified
// Create options for the input (fat binary) and output (target images)
// files, each with an appropriate default when not specified
static cl::opt<std::string> Input(cl::Positional, cl::init("a.out"),
cl::desc("<input file>"),
cl::cat(ClangOffloadExtractCategory));
Expand All @@ -61,10 +67,8 @@ static cl::alias FileNameStemAlias(
"output", cl::desc("Deprecated option, replaced by option '--stem'"),
cl::aliasopt(FileNameStem), cl::cat(ClangOffloadExtractCategory));

// Path to the current binary
static std::string ToolPath;

// Report error (and handle any deferred errors)
// This is the main diagnostic used by the utility
static void reportError(Error E, Twine message = "\n") {
std::string S;
raw_string_ostream OSS(S);
Expand All @@ -80,7 +84,6 @@ static void reportError(Error E, Twine message = "\n") {

int main(int argc, const char **argv) {
sys::PrintStackTraceOnErrorSignal(argv[0]);
ToolPath = argv[0];

// Hide non-generic options that are not in this utility's explicit
// category;
Expand Down Expand Up @@ -138,105 +141,144 @@ linked fat binary, and store them in separate files.
}

// We are dealing with an appropriate fat binary;
// Locate the section IMAGE_INFO_SECTION_NAME (which contains the
// metadata on the embedded binaries)
unsigned FileNum = 0;

// * Create an array all the sections that have
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
// * Create an array all the sections that have
// * Create an array of all the sections that have

// IMAGE_SECTION_NAME in the section name:
auto OffloadSections = SmallVector<SectionRef>();
// * Locate the section that starts with IMAGE_INFO_SECTION_NAME_PREFIX
// and extract the index for all the embedded binaries:
auto OffloadIndex = SmallVector<SectionRef>();
for (const auto &Section : Binary->sections()) {
Expected<StringRef> NameOrErr = Section.getName();
if (auto E = NameOrErr.takeError()) {
Expected<StringRef> InfoSecNameOrErr = Section.getName();
if (auto E = InfoSecNameOrErr.takeError()) {
reportError(std::move(E), "Input File: '" + Input + "'\n");
}
if (*NameOrErr != IMAGE_INFO_SECTION_NAME)
continue;
LLVM_DEBUG(dbgs() << "Section: " << *InfoSecNameOrErr << "\n");

// This is the section we are looking for;
// Extract the meta information:
// The IMAGE_INFO_SECTION_NAME section contains packed <address,
// size> pairs describing target images that are stored in the fat
// binary.
Expected<StringRef> DataOrErr = Section.getContents();
if (auto E = DataOrErr.takeError()) {
reportError(std::move(E), "Input File: '" + Input + "'\n");
// We have a valid section name
std::string SectionNameToCompare = isa<COFFObjectFile>(Binary)
? IMAGE_SECTION_NAME_PREFIX_COFF
: IMAGE_SECTION_NAME_PREFIX;
if (InfoSecNameOrErr->find(SectionNameToCompare) != std::string::npos) {
// This section contains embedded binaries
OffloadSections.push_back(Section);
} else if (*InfoSecNameOrErr == IMAGE_INFO_SECTION_NAME) {
// This section is the index
OffloadIndex.push_back(Section);
}
// Data type to store the metadata for an individual target image
struct ImgInfoType {
uintptr_t Addr;
uintptr_t Size;
};

// Store the metadata for all target images in an array of target
// image information descriptors
auto ImgInfo =
ArrayRef(reinterpret_cast<const ImgInfoType *>(DataOrErr->data()),
DataOrErr->size() / sizeof(ImgInfoType));

// Loop over the image information descriptors to extract each
// target image.
for (const auto &Img : ImgInfo) {
// Ignore zero padding that can be inserted by the linker.
if (!Img.Addr)
continue;

// Find section which contains this image.
// TODO: can use more efficient algorithm than linear search. For
// example sections and images could be sorted by address then one pass
// performed through both at the same time.
// std::find_if
// * searches for a true predicate in [first,last] =~ [first,end)
// * returns end if no predicate is true
// It is probably faster to track success through a bool (ImgFound)
bool ImgFound = false;
auto ImgSec =
find_if(Binary->sections(), [&Img, &ImgFound](SectionRef Sec) {
bool pred = ( //
Sec.isData() //
&& (Img.Addr == Sec.getAddress()) //
&& (Img.Size == Sec.getSize()) //
);
ImgFound = ImgFound || pred;
return pred;
});
if (!ImgFound) {

reportError(
createStringError(
inconvertibleErrorCode(),
"cannot find section containing <0x%lx, 0x%lx> target image",
Img.Addr, Img.Size),
"Input File: '" + Input + "'\n");
}

Expected<StringRef> SecDataOrErr = ImgSec->getContents();
if (auto E = SecDataOrErr.takeError()) {
reportError(std::move(E), "Input File: '" + Input + "'\n");
}

// Output file name is composed from the name prefix provided by the
// user and the image number which is appended to the prefix
std::string FileName = FileNameStem + "." + std::to_string(FileNum++);

// Tell user that we are saving an image.
outs() << "Saving target image to '" << FileName << "'\n";

// Write image data to the output
std::error_code EC;
raw_fd_ostream OS(FileName, EC);
if (EC) {
reportError(createFileError(FileName, EC),
"Specify a different Output File ('--stem' option)\n");
}

OS << SecDataOrErr->substr(Img.Addr - ImgSec->getAddress(), Img.Size);
if (OS.has_error()) {
reportError(createFileError(FileName, OS.error()),
"Try a different Output File ('--stem' option)");
}
} // &Img: ImgInfo

// Fat binary is not expected to have more than one .tgtimg section.
break;
}

// Check if there are any sections with embedded binaries
if (OffloadSections.size() == 0) {
reportError(
createStringError(inconvertibleErrorCode(),
"Could not locate sections with offload binaries"),
"Fat Binary: '" + Input + "'\n");
}
// Check if we found the index section
if (OffloadIndex.size() == 0) {
reportError(
createStringError(inconvertibleErrorCode(),
"Could not locate index for embedded binaries"),
"Fat Binary: '" + Input + "'\n");
}
// Check if we have a valid index section
Expected<StringRef> DataOrErr = OffloadIndex[0].getContents();
if (auto E = DataOrErr.takeError()) {
reportError(std::move(E), "Input File: '" + Input + "'\n");
}

// The input binary contains embedded offload binaries
// The index section contains packed <address, size> pairs describing
// target images that are stored in the fat binary.
// Data type to store the index for an individual target image
struct ImgInfoType {
uintptr_t Addr;
uintptr_t Size;
};

// Store the metadata for all target images in an array of target
// image information descriptors
// This can be done by reinterpreting the content of the section
auto ImgInfo =
ArrayRef(reinterpret_cast<const ImgInfoType *>(DataOrErr->data()),
DataOrErr->size() / sizeof(ImgInfoType));

// Loop over the image information descriptors to extract each
// target image the object file data
unsigned FileNum = 0;
unsigned ImgCnt = 1;

for (const auto &Img : ImgInfo) {
// Ignore zero padding that can be inserted by the linker.
if (!Img.Addr)
continue;

// Find section which contains this image.
// /!\ There might be multiple images in a section
// std::find_if
// * searches for a true predicate in [first,last] =~ [first,end)
// * returns end if no predicate is true
// It is probably faster to track success through a bool (ImgFound)
bool ImgFound = false;
auto ImgSec = find_if(OffloadSections, [&Img, &ImgFound](auto Sec) {
bool pred = ( //
Sec.isData() //
&& (Img.Addr >= Sec.getAddress()) //
&& ((Img.Addr + Img.Size) <= (Sec.getAddress() + Sec.getSize())) //
);
ImgFound = ImgFound || pred;
return pred;
});
if (!ImgFound) {

reportError(
createStringError(inconvertibleErrorCode(),
"Target image (Address:0x%lx, Size:0x%lx>) is "
"not contained in any section",
Img.Addr, Img.Size),
"Fat Binary: '" + Input + "'\n");
}

Expected<StringRef> ImgSecNameOrErr = ImgSec->getName();
if (auto E = ImgSecNameOrErr.takeError()) {
reportError(std::move(E),
"Can not determine section name in Fat Binary '" + Input +
"'\n");
}
Expected<StringRef> SecDataOrErr = ImgSec->getContents();
if (auto E = SecDataOrErr.takeError()) {
reportError(std::move(E), "Can not extract contents for section '" +
*ImgSecNameOrErr + "' in Fat Binary '" +
Input + "'\n");
}

// Output file name is composed from the name prefix provided by the
// user and the image number which is appended to the prefix
std::string FileName = FileNameStem + "." + std::to_string(FileNum++);
std::string OffloadName = ImgSecNameOrErr->data();
std::string OffloadPrefix = isa<COFFObjectFile>(Binary)
? IMAGE_SECTION_NAME_PREFIX_COFF
: IMAGE_SECTION_NAME_PREFIX;
OffloadName.erase(0, OffloadPrefix.length());

// Tell user that we are saving an image.
outs() << "Section '" + OffloadName + "': Image " << ImgCnt++
<< "'-> File '" + FileName + "'\n";

// Write image data to the output
std::error_code EC;
raw_fd_ostream OS(FileName, EC);
if (EC) {
reportError(createFileError(FileName, EC),
"Specify a different Output File ('--stem' option)\n");
}

OS << SecDataOrErr->substr(Img.Addr - ImgSec->getAddress(), Img.Size);
if (OS.has_error()) {
reportError(createFileError(FileName, OS.error()),
"Try a different Output File ('--stem' option)");
}
} // &Img: ImgInfo

return 0;
}