llvm/lib/MC/MCDXContainerWriter.cpp

//===- llvm/MC/MCDXContainerWriter.cpp - DXContainer Writer -----*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "llvm/MC/MCDXContainerWriter.h"
#include "llvm/BinaryFormat/DXContainer.h"
#include "llvm/MC/MCAssembler.h"
#include "llvm/MC/MCContext.h"
#include "llvm/MC/MCSection.h"
#include "llvm/MC/MCValue.h"
#include "llvm/Support/Alignment.h"
#include "llvm/Support/CommandLine.h"

using namespace llvm;

cl::opt<bool> EmbedDebug("dx-embed-debug",
                         cl::desc("Embed PDB in shader container"));
cl::opt<bool>
    StripDebug("dx-strip-debug",
               cl::desc("Strip debug information from shader bytecode"));
cl::opt<bool> SlimDebug("dx-slim-debug",
                        cl::desc("Generate slim PDB without ILDB part"));

MCDXContainerTargetWriter::~MCDXContainerTargetWriter() = default;

MCDXContainerBaseWriter::~MCDXContainerBaseWriter() = default;

bool MCDXContainerBaseWriter::shouldSkipSection(StringRef SectionName,
                                                size_t SectionSize) {
  // Skip empty and auxiliary sections.
  if (SectionSize == 0 || SectionName == PdbFileNameSectionName ||
      SectionName == ModuleHashSectionName)
    return true;
  // Slim debug omits ILDB from all DXContainer outputs.
  return SlimDebug && SectionName == "ILDB";
}

void MCDXContainerBaseWriter::write(raw_ostream &OS, const Triple &TT) {
  ArrayRef<MCDXContainerPart> Parts = collectParts();

  support::endian::Writer W(OS, llvm::endianness::little);

  // Start the file size as the header plus the size of the part offsets.
  // Presently DXContainer files usually contain 7-10 parts. Reserving space for
  // 16 part offsets gives us a little room for growth.
  llvm::SmallVector<uint64_t, 16> PartOffsets;
  uint64_t PartOffset = 0;
  bool HasPrivate = false;
  for (const MCDXContainerPart &Part : Parts) {
    if (HasPrivate)
      reportFatalInternalError(
          "PRIV must be the last section in a DXContainer");
    if (Part.Name == "PRIV")
      HasPrivate = true;

    uint64_t SectionSize = Part.Data.size();
    assert(SectionSize < std::numeric_limits<uint32_t>::max() &&
           "Section size too large for DXContainer");

    PartOffsets.push_back(PartOffset);
    PartOffset += sizeof(dxbc::PartHeader) + SectionSize;
    if (!HasPrivate)
      PartOffset = alignTo(PartOffset, Align(4ul));
    // The DXIL part also writes a program header, so we need to include its
    // size when computing the offset for a part after the DXIL part.
    if (dxbc::isProgramPart(Part.Name))
      PartOffset += sizeof(dxbc::ProgramHeader);
  }
  assert(PartOffset < std::numeric_limits<uint32_t>::max() &&
         "Part data too large for DXContainer");

  uint64_t PartStart =
      sizeof(dxbc::Header) + (PartOffsets.size() * sizeof(uint32_t));
  uint64_t FileSize = PartStart + PartOffset;
  assert(FileSize < std::numeric_limits<uint32_t>::max() &&
         "File size too large for DXContainer");

  // Write the header.
  W.write<char>({'D', 'X', 'B', 'C'});
  // Write 16-bytes of 0's for the hash.
  W.OS.write_zeros(16);
  // Write 1.0 for file format version.
  W.write<uint16_t>(1u);
  W.write<uint16_t>(0u);
  // Write the file size.
  W.write<uint32_t>(static_cast<uint32_t>(FileSize));
  // Write the number of parts.
  W.write<uint32_t>(static_cast<uint32_t>(PartOffsets.size()));
  // Write the offsets for the part headers for each part.
  for (uint64_t Offset : PartOffsets)
    W.write<uint32_t>(static_cast<uint32_t>(PartStart + Offset));

  for (const MCDXContainerPart &Part : Parts) {
    uint64_t SectionSize = Part.Data.size();
    unsigned Start = W.OS.tell();
    // Write section header.
    W.write<char>(ArrayRef<char>(Part.Name.data(), 4));

    uint64_t PartSize = SectionSize;

    if (dxbc::isProgramPart(Part.Name))
      PartSize += sizeof(dxbc::ProgramHeader);
    // DXContainer part should be 4-byte aligned, unless it is PRIV part.
    bool IsPrivate = Part.Name == "PRIV";
    if (!IsPrivate)
      PartSize = alignTo(PartSize, Align(4));
    W.write<uint32_t>(static_cast<uint32_t>(PartSize));
    if (dxbc::isProgramPart(Part.Name)) {
      dxbc::ProgramHeader Header;
      memset(reinterpret_cast<void *>(&Header), 0, sizeof(dxbc::ProgramHeader));

      VersionTuple Version = TT.getOSVersion();
      uint8_t MajorVersion = static_cast<uint8_t>(Version.getMajor());
      uint8_t MinorVersion =
          static_cast<uint8_t>(Version.getMinor().value_or(0));
      Header.Version =
          dxbc::ProgramHeader::getVersion(MajorVersion, MinorVersion);
      if (TT.hasEnvironment())
        Header.ShaderKind =
            static_cast<uint16_t>(TT.getEnvironment() - Triple::Pixel);

      // The program header's size field is in 32-bit words.
      Header.Size = (SectionSize + sizeof(dxbc::ProgramHeader) + 3) / 4;
      memcpy(Header.Bitcode.Magic, "DXIL", 4);
      VersionTuple DXILVersion = TT.getDXILVersion();
      Header.Bitcode.MajorVersion = DXILVersion.getMajor();
      Header.Bitcode.MinorVersion = DXILVersion.getMinor().value_or(0);
      Header.Bitcode.Offset = sizeof(dxbc::BitcodeHeader);
      Header.Bitcode.Size = SectionSize;
      if (sys::IsBigEndianHost)
        Header.swapBytes();
      W.write<char>(ArrayRef<char>(reinterpret_cast<char *>(&Header),
                                   sizeof(dxbc::ProgramHeader)));
    }
    W.write<char>(Part.Data);
    if (!IsPrivate) {
      unsigned Size = W.OS.tell() - Start;
      W.OS.write_zeros(offsetToAlignment(Size, Align(4)));
    }
  }
}

void DXContainerObjectWriter::clearParts() {
  Parts.clear();
  SectionBuffers.clear();
}

ArrayRef<MCDXContainerPart> DXContainerObjectWriter::collectParts() {
  clearParts();
  for (const MCSection &Sec : *Asm) {
    if (shouldSkipSection(Sec.getName(), Asm->getSectionAddressSize(Sec)))
      continue;

    SectionBuffers.emplace_back();
    raw_svector_ostream OS(SectionBuffers.back());
    Asm->writeSectionData(OS, &Sec);
    Parts.push_back({Sec.getName(), StringRef(SectionBuffers.back())});
  }
  return Parts;
}

bool DXContainerObjectWriter::shouldSkipSection(StringRef SectionName,
                                                size_t SectionSize) {
  // Do not write ILDB part if we're not embedding it.
  if (SectionName == "ILDB" && (!EmbedDebug || StripDebug))
    return true;
  if (SectionName == "SRCI")
    return true;
  return MCDXContainerBaseWriter::shouldSkipSection(SectionName, SectionSize);
}

uint64_t DXContainerObjectWriter::writeObject() {
  write(W.OS, getContext().getTargetTriple());
  clearParts();

  return 0;
}