llvm-project/llvm/lib/ObjectYAML/DXContainerEmitter.cpp
Chris Bieneman 2556ba4a52 [ObjectYAML] Add support for DXContainer HASH
DXContainer files contain a part that has an MD5 of the generated
shader. This adds support to the ObjectYAML tooling to expand the hash
part data and hash iteself in preparation for adding hashing support to
DirectX code generation.

Reviewed By: python3kgae

Differential Revision: https://reviews.llvm.org/D136632
2022-10-27 12:28:45 -05:00

231 lines
7.6 KiB
C++

//===- DXContainerEmitter.cpp - Convert YAML to a DXContainer -------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Binary emitter for yaml to DXContainer binary
///
//===----------------------------------------------------------------------===//
#include "llvm/BinaryFormat/DXContainer.h"
#include "llvm/ObjectYAML/ObjectYAML.h"
#include "llvm/ObjectYAML/yaml2obj.h"
#include "llvm/Support/Errc.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;
namespace {
class DXContainerWriter {
public:
DXContainerWriter(DXContainerYAML::Object &ObjectFile)
: ObjectFile(ObjectFile) {}
Error write(raw_ostream &OS);
private:
DXContainerYAML::Object &ObjectFile;
Error computePartOffsets();
Error validatePartOffsets();
Error validateSize(uint32_t Computed);
void writeHeader(raw_ostream &OS);
void writeParts(raw_ostream &OS);
};
} // namespace
Error DXContainerWriter::validateSize(uint32_t Computed) {
if (!ObjectFile.Header.FileSize)
ObjectFile.Header.FileSize = Computed;
else if (*ObjectFile.Header.FileSize < Computed)
return createStringError(errc::result_out_of_range,
"File size specified is too small.");
return Error::success();
}
Error DXContainerWriter::validatePartOffsets() {
if (ObjectFile.Parts.size() != ObjectFile.Header.PartOffsets->size())
return createStringError(
errc::invalid_argument,
"Mismatch between number of parts and part offsets.");
uint32_t RollingOffset =
sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
if (RollingOffset > std::get<1>(I))
return createStringError(errc::invalid_argument,
"Offset mismatch, not enough space for data.");
RollingOffset =
std::get<1>(I) + sizeof(dxbc::PartHeader) + std::get<0>(I).Size;
}
if (Error Err = validateSize(RollingOffset))
return Err;
return Error::success();
}
Error DXContainerWriter::computePartOffsets() {
if (ObjectFile.Header.PartOffsets)
return validatePartOffsets();
uint32_t RollingOffset =
sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
ObjectFile.Header.PartOffsets = std::vector<uint32_t>();
for (const auto &Part : ObjectFile.Parts) {
ObjectFile.Header.PartOffsets->push_back(RollingOffset);
RollingOffset += sizeof(dxbc::PartHeader) + Part.Size;
}
if (Error Err = validateSize(RollingOffset))
return Err;
return Error::success();
}
void DXContainerWriter::writeHeader(raw_ostream &OS) {
dxbc::Header Header;
memcpy(Header.Magic, "DXBC", 4);
memcpy(Header.FileHash.Digest, ObjectFile.Header.Hash.data(), 16);
Header.Version.Major = ObjectFile.Header.Version.Major;
Header.Version.Minor = ObjectFile.Header.Version.Minor;
Header.FileSize = *ObjectFile.Header.FileSize;
Header.PartCount = ObjectFile.Parts.size();
if (sys::IsBigEndianHost)
Header.swapBytes();
OS.write(reinterpret_cast<char *>(&Header), sizeof(Header));
SmallVector<uint32_t> Offsets(ObjectFile.Header.PartOffsets->begin(),
ObjectFile.Header.PartOffsets->end());
if (sys::IsBigEndianHost)
for (auto &O : Offsets)
sys::swapByteOrder(O);
OS.write(reinterpret_cast<char *>(Offsets.data()),
Offsets.size() * sizeof(uint32_t));
}
void DXContainerWriter::writeParts(raw_ostream &OS) {
uint32_t RollingOffset =
sizeof(dxbc::Header) + (ObjectFile.Header.PartCount * sizeof(uint32_t));
for (auto I : llvm::zip(ObjectFile.Parts, *ObjectFile.Header.PartOffsets)) {
if (RollingOffset < std::get<1>(I)) {
uint32_t PadBytes = std::get<1>(I) - RollingOffset;
OS.write_zeros(PadBytes);
}
DXContainerYAML::Part P = std::get<0>(I);
RollingOffset = std::get<1>(I) + sizeof(dxbc::PartHeader);
uint32_t PartSize = P.Size;
OS.write(P.Name.c_str(), 4);
if (sys::IsBigEndianHost)
sys::swapByteOrder(P.Size);
OS.write(reinterpret_cast<const char *>(&P.Size), sizeof(uint32_t));
dxbc::PartType PT = dxbc::parsePartType(P.Name);
uint64_t DataStart = OS.tell();
switch (PT) {
case dxbc::PartType::DXIL: {
if (!P.Program)
continue;
dxbc::ProgramHeader Header;
Header.MajorVersion = P.Program->MajorVersion;
Header.MinorVersion = P.Program->MinorVersion;
Header.Unused = 0;
Header.ShaderKind = P.Program->ShaderKind;
memcpy(Header.Bitcode.Magic, "DXIL", 4);
Header.Bitcode.MajorVersion = P.Program->DXILMajorVersion;
Header.Bitcode.MinorVersion = P.Program->DXILMinorVersion;
Header.Bitcode.Unused = 0;
// Compute the optional fields if needed...
if (P.Program->DXILOffset)
Header.Bitcode.Offset = P.Program->DXILOffset.value();
else
Header.Bitcode.Offset = sizeof(dxbc::BitcodeHeader);
if (P.Program->DXILSize)
Header.Bitcode.Size = P.Program->DXILSize.value();
else
Header.Bitcode.Size = P.Program->DXIL ? P.Program->DXIL->size() : 0;
if (P.Program->Size)
Header.Size = P.Program->Size.value();
else
Header.Size = sizeof(dxbc::ProgramHeader) + Header.Bitcode.Size;
uint32_t BitcodeOffset = Header.Bitcode.Offset;
if (sys::IsBigEndianHost)
Header.swapBytes();
OS.write(reinterpret_cast<const char *>(&Header),
sizeof(dxbc::ProgramHeader));
if (P.Program->DXIL) {
if (BitcodeOffset > sizeof(dxbc::BitcodeHeader)) {
uint32_t PadBytes = BitcodeOffset - sizeof(dxbc::BitcodeHeader);
OS.write_zeros(PadBytes);
}
OS.write(reinterpret_cast<char *>(P.Program->DXIL->data()),
P.Program->DXIL->size());
}
break;
}
case dxbc::PartType::SFI0: {
// If we don't have any flags we can continue here and the data will be
// zeroed out.
if (!P.Flags.has_value())
continue;
uint64_t Flags = P.Flags->getEncodedFlags();
if (sys::IsBigEndianHost)
sys::swapByteOrder(Flags);
OS.write(reinterpret_cast<char *>(&Flags), sizeof(uint64_t));
break;
}
case dxbc::PartType::HASH: {
if (!P.Hash.has_value())
continue;
dxbc::ShaderHash Hash = {0, {0}};
if (P.Hash->IncludesSource)
Hash.Flags |= static_cast<uint32_t>(dxbc::HashFlags::IncludesSource);
memcpy(&Hash.Digest[0], &P.Hash->Digest[0], 16);
if (sys::IsBigEndianHost)
Hash.swapBytes();
OS.write(reinterpret_cast<char *>(&Hash), sizeof(dxbc::ShaderHash));
break;
}
case dxbc::PartType::Unknown:
break; // Skip any handling for unrecognized parts.
}
uint64_t BytesWritten = OS.tell() - DataStart;
RollingOffset += BytesWritten;
if (BytesWritten < PartSize)
OS.write_zeros(PartSize - BytesWritten);
RollingOffset += PartSize;
}
}
Error DXContainerWriter::write(raw_ostream &OS) {
if (Error Err = computePartOffsets())
return Err;
writeHeader(OS);
writeParts(OS);
return Error::success();
}
namespace llvm {
namespace yaml {
bool yaml2dxcontainer(DXContainerYAML::Object &Doc, raw_ostream &Out,
ErrorHandler EH) {
DXContainerWriter Writer(Doc);
if (Error Err = Writer.write(Out)) {
handleAllErrors(std::move(Err),
[&](const ErrorInfoBase &Err) { EH(Err.message()); });
return false;
}
return true;
}
} // namespace yaml
} // namespace llvm