[DirectX backend] support empty program signature parts for compute shader

[python3kgae]

Summary

New parts ISG1 and OSG1 are added but only support 0 parameter case.
2 uint32_t will be emited for empty program signature. 0 for parameter count and 8 for parameter offset.

This is enough for computer shader.

[beanz]

All the other changes going into the DXContainer object and yaml support are aiming to be more full featured, and the object and yaml tooling changes are separate from the implementation change.

Please follow that structure.

llvm/include/llvm/BinaryFormat/DXContainer.h

133

nit: please add whitespace separating the members from the method.

[python3kgae]

Add test for new globals generated in DXContatinerGlobal pass.

llvm/include/llvm/BinaryFormat/DXContainer.h

void swapBytes() {
     sys::swapByteOrder(ShaderKind);
     sys::swapByteOrder(Size);
     Bitcode.swapBytes();
   }
 };
 
 static_assert(sizeof(ProgramHeader) == 24, "ProgramHeader Size incorrect!");
 
+struct ProgramSignature {
+  uint32_t ParamCount;
+  uint32_t ParamOffset;
+
+  void swapBytes() {
+    sys::swapByteOrder(ParamCount);
+    sys::swapByteOrder(ParamOffset);
+  }
+};
+
+static_assert(sizeof(ProgramSignature) == 8,
+              "ProgramSignature Size incorrect!");
+
 #define CONTAINER_PART(Part) Part,
 enum class PartType {
   Unknown = 0,
 #include "DXContainerConstants.def"
 };
 
 #define SHADER_FLAG(Num, Val, Str) Val = 1ull << Num,
 enum class FeatureFlags : uint64_t {

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

 
 #ifdef CONTAINER_PART
 CONTAINER_PART(DXIL)
 CONTAINER_PART(SFI0)
 CONTAINER_PART(HASH)
+CONTAINER_PART(ISG1)
+CONTAINER_PART(OSG1)
 
 #undef CONTAINER_PART
 #endif 
 
 #ifdef SHADER_FLAG
 
 SHADER_FLAG(0, Doubles, "Double-precision floating point")
 SHADER_FLAG(1, ComputeShadersPlusRawAndStructuredBuffers, "Raw and Structured buffers")

llvm/include/llvm/Object/DXContainer.h

private:
   DXContainer(MemoryBufferRef O);
 
   MemoryBufferRef Data;
   dxbc::Header Header;
   SmallVector<uint32_t, 4> PartOffsets;
   std::optional<DXILData> DXIL;
   std::optional<uint64_t> ShaderFlags;
   std::optional<dxbc::ShaderHash> Hash;
+  std::optional<dxbc::ProgramSignature> InputSig;
+  std::optional<dxbc::ProgramSignature> OutputSig;
 
   Error parseHeader();
   Error parsePartOffsets();
   Error parseDXILHeader(StringRef Part);
   Error parseShaderFlags(StringRef Part);
   Error parseHash(StringRef Part);
+  Error parseProgramSignature(StringRef Part, dxbc::ProgramSignature &Sig);
   friend class PartIterator;
 
 public:
   // The PartIterator is a wrapper around the iterator for the PartOffsets
   // member of the DXContainer. It contains a refernce to the container, and the
   // current iterator value, as well as storage for a parsed part header.
   class PartIterator {
     const DXContainer &Container;
public:
 
   const dxbc::Header &getHeader() const { return Header; }
 
   std::optional<DXILData> getDXIL() const { return DXIL; }
 
   std::optional<uint64_t> getShaderFlags() const { return ShaderFlags; }
 
   std::optional<dxbc::ShaderHash> getShaderHash() const { return Hash; }
+
+  std::optional<dxbc::ProgramSignature> getInputSignature() const {
+    return InputSig;
+  }
+  std::optional<dxbc::ProgramSignature> getOutputSignature() const {
+    return OutputSig;
+  }
 };
 
 } // namespace object
 } // namespace llvm
 
 #endif // LLVM_OBJECT_DXCONTAINER_H

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

 struct ShaderHash {
   ShaderHash() = default;
   ShaderHash(const dxbc::ShaderHash &Data);
 
   bool IncludesSource;
   std::vector<llvm::yaml::Hex8> Digest;
 };
 
+struct ProgramSignature {
+  ProgramSignature() = default;
+  ProgramSignature(const dxbc::ProgramSignature &Data);
+  // FIXME: support case when ParamCount is not 0.
+  //  see issue https://github.com/llvm/llvm-project/issues/57743 and
+  //  https://github.com/llvm/llvm-project/issues/57744.
+  uint32_t ParamCount;
+  uint32_t ParamOffset;
+};
+
 struct Part {
   Part() = default;
   Part(std::string N, uint32_t S) : Name(N), Size(S) {}
   std::string Name;
   uint32_t Size;
   std::optional<DXILProgram> Program;
   std::optional<ShaderFlags> Flags;
   std::optional<ShaderHash> Hash;
+  std::optional<ProgramSignature> InputSig;
+  std::optional<ProgramSignature> OutputSig;
 };
 
 struct Object {
   FileHeader Header;
   std::vector<Part> Parts;
 };
 
 } // namespace DXContainerYAML

 template <> struct MappingTraits<DXContainerYAML::ShaderFlags> {
   static void mapping(IO &IO, DXContainerYAML::ShaderFlags &Flags);
 };
 
 template <> struct MappingTraits<DXContainerYAML::ShaderHash> {
   static void mapping(IO &IO, DXContainerYAML::ShaderHash &Hash);
 };
 
+template <> struct MappingTraits<DXContainerYAML::ProgramSignature> {
+  static void mapping(IO &IO, DXContainerYAML::ProgramSignature &Sig);
+};
+
 template <> struct MappingTraits<DXContainerYAML::Part> {
   static void mapping(IO &IO, DXContainerYAML::Part &Version);
 };
 
 template <> struct MappingTraits<DXContainerYAML::Object> {
   static void mapping(IO &IO, DXContainerYAML::Object &Obj);
 };
 
 } // namespace yaml
 
 } // namespace llvm
 
 #endif // LLVM_OBJECTYAML_DXCONTAINERYAML_H

llvm/lib/Object/DXContainer.cpp

if (Hash)
     return parseFailed("More than one HASH part is present in the file");
   dxbc::ShaderHash ReadHash;
   if (Error Err = readStruct(Part, Part.begin(), ReadHash))
     return Err;
   Hash = ReadHash;
   return Error::success();
 }
 
+Error DXContainer::parseProgramSignature(StringRef Part,
+                                         dxbc::ProgramSignature &Sig) {
+  if (Error Err = readStruct(Part, Part.begin(), Sig))
+    return Err;
+
+  return Error::success();
+}
+
 Error DXContainer::parsePartOffsets() {
   uint32_t LastOffset =
       sizeof(dxbc::Header) + (Header.PartCount * sizeof(uint32_t));
   const char *Current = Data.getBuffer().data() + sizeof(dxbc::Header);
   for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     uint32_t PartOffset;
     if (Error Err = readInteger(Data.getBuffer(), Current, PartOffset))
       return Err;
for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     case dxbc::PartType::SFI0:
       if (Error Err = parseShaderFlags(PartData))
         return Err;
       break;
     case dxbc::PartType::HASH:
       if (Error Err = parseHash(PartData))
         return Err;
       break;
+    case dxbc::PartType::ISG1: {
+      dxbc::ProgramSignature Sig;
+      if (Error Err = parseProgramSignature(PartData, Sig))
+        return Err;
+      InputSig = Sig;
+    } break;
+    case dxbc::PartType::OSG1: {
+      dxbc::ProgramSignature Sig;
+      if (Error Err = parseProgramSignature(PartData, Sig))
+        return Err;
+      OutputSig = Sig;
+    } break;
     case dxbc::PartType::Unknown:
       break;
     }
   }
   return Error::success();
 }
 
 Expected<DXContainer> DXContainer::create(MemoryBufferRef Object) {

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

case dxbc::PartType::HASH: {
       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::ISG1:
+    case dxbc::PartType::OSG1: {
+      bool IsInput = PT == dxbc::PartType::ISG1;
+      dxbc::ProgramSignature Sig = {
+          IsInput ? P.InputSig->ParamCount : P.OutputSig->ParamCount,
+          IsInput ? P.InputSig->ParamOffset : P.OutputSig->ParamOffset};
+      if (sys::IsBigEndianHost)
+        Sig.swapBytes();
+
+      OS.write(reinterpret_cast<char *>(&Sig), sizeof(dxbc::ProgramSignature));
+    } 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;

llvm/lib/ObjectYAML/DXContainerYAML.cpp

 
 DXContainerYAML::ShaderHash::ShaderHash(const dxbc::ShaderHash &Data)
     : IncludesSource((Data.Flags & static_cast<uint32_t>(
                                        dxbc::HashFlags::IncludesSource)) != 0),
       Digest(16, 0) {
   memcpy(Digest.data(), &Data.Digest[0], 16);
 }
 
+DXContainerYAML::ProgramSignature::ProgramSignature(
+    const dxbc::ProgramSignature &Data)
+    : ParamCount(Data.ParamCount), ParamOffset(Data.ParamOffset) {}
+
 namespace yaml {
 
 void MappingTraits<DXContainerYAML::VersionTuple>::mapping(
     IO &IO, DXContainerYAML::VersionTuple &Version) {
   IO.mapRequired("Major", Version.Major);
   IO.mapRequired("Minor", Version.Minor);
 }
 

 }
 
 void MappingTraits<DXContainerYAML::ShaderFlags>::mapping(
     IO &IO, DXContainerYAML::ShaderFlags &Flags) {
 #define SHADER_FLAG(Num, Val, Str) IO.mapRequired(#Val, Flags.Val);
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
+void MappingTraits<DXContainerYAML::ProgramSignature>::mapping(
+    IO &IO, DXContainerYAML::ProgramSignature &Sig) {
+  IO.mapRequired("ParamCount", Sig.ParamCount);
+  IO.mapRequired("ParamOffset", Sig.ParamOffset);
+}
+
 void MappingTraits<DXContainerYAML::ShaderHash>::mapping(
     IO &IO, DXContainerYAML::ShaderHash &Hash) {
   IO.mapRequired("IncludesSource", Hash.IncludesSource);
   IO.mapRequired("Digest", Hash.Digest);
 }
 
 void MappingTraits<DXContainerYAML::Part>::mapping(IO &IO,
                                                    DXContainerYAML::Part &P) {
   IO.mapRequired("Name", P.Name);
   IO.mapRequired("Size", P.Size);
   IO.mapOptional("Program", P.Program);
   IO.mapOptional("Flags", P.Flags);
   IO.mapOptional("Hash", P.Hash);
+  IO.mapOptional("InputSig", P.InputSig);
+  IO.mapOptional("OutputSig", P.OutputSig);
 }
 
 void MappingTraits<DXContainerYAML::Object>::mapping(
     IO &IO, DXContainerYAML::Object &Obj) {
   IO.mapTag("!dxcontainer", true);
   IO.mapRequired("Header", Obj.Header);
   IO.mapRequired("Parts", Obj.Parts);
 }
 
 } // namespace yaml
 } // namespace llvm

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MD5.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
+#define DEBUG_TYPE "dxil-globals"
+
 using namespace llvm;
 using namespace llvm::dxil;
 
 namespace {
 class DXContainerGlobals : public llvm::ModulePass {
+  GlobalVariable *buildContainerGlobal(Module &M, Constant *Content,
+                                       StringRef Name, StringRef SectionName);
 
   GlobalVariable *getShaderFlags(Module &M);
   GlobalVariable *computeShaderHash(Module &M);
+  GlobalVariable *buildInputSingature(Module &M, Constant *Content);
+  GlobalVariable *buildOutputSingature(Module &M, Constant *Content);
+  void addSingature(Module &M, SmallVector<GlobalValue *> &Globals);
 
 public:
   static char ID; // Pass identification, replacement for typeid
   DXContainerGlobals() : ModulePass(ID) {
     initializeDXContainerGlobalsPass(*PassRegistry::getPassRegistry());
   }
 
   StringRef getPassName() const override {

 };
 
 } // namespace
 
 bool DXContainerGlobals::runOnModule(Module &M) {
   llvm::SmallVector<GlobalValue *> Globals;
   Globals.push_back(getShaderFlags(M));
   Globals.push_back(computeShaderHash(M));
-
+  addSingature(M, Globals);
   appendToCompilerUsed(M, Globals);
   return true;
 }
 
 GlobalVariable *DXContainerGlobals::getShaderFlags(Module &M) {
   const uint64_t Flags =
       (uint64_t)(getAnalysis<ShaderFlagsAnalysisWrapper>().getShaderFlags());
 
GlobalVariable *DXContainerGlobals::computeShaderHash(Module &M) {
   auto *GV = new llvm::GlobalVariable(M, ModuleConstant->getType(), true,
                                       GlobalValue::PrivateLinkage,
                                       ModuleConstant, "dx.hash");
   GV->setSection("HASH");
   GV->setAlignment(Align(4));
   return GV;
 }
 
+GlobalVariable *DXContainerGlobals::buildContainerGlobal(
+    Module &M, Constant *Content, StringRef Name, StringRef SectionName) {
+  auto *GV = new llvm::GlobalVariable(
+      M, Content->getType(), true, GlobalValue::PrivateLinkage, Content, Name);
+  GV->setSection(SectionName);
+  GV->setAlignment(Align(4));
+  return GV;
+}
+
+GlobalVariable *DXContainerGlobals::buildInputSingature(Module &M,
+                                                        Constant *Content) {
+  return buildContainerGlobal(M, Content, "dx.isg1", "ISG1");
+}
+GlobalVariable *DXContainerGlobals::buildOutputSingature(Module &M,
+                                                         Constant *Content) {
+  return buildContainerGlobal(M, Content, "dx.osg1", "OSG1");
+}
+void DXContainerGlobals::addSingature(Module &M,
+                                      SmallVector<GlobalValue *> &Globals) {
+  dxbc::ProgramSignature Sig;
+  Sig.ParamCount = 0;
+  Sig.ParamOffset = sizeof(dxbc::ProgramSignature);
+  // FIXME: support case when ParamCount is not 0.
+  //  see issue https://github.com/llvm/llvm-project/issues/57743 and
+  //  https://github.com/llvm/llvm-project/issues/57744.
+  Type *Int32Ty = Type::getInt32Ty(M.getContext());
+  Constant *EmptySig = ConstantStruct::get(
+      StructType::get(Int32Ty, Int32Ty),
+      {ConstantInt::get(M.getContext(), APInt(32, Sig.ParamCount)),
+       ConstantInt::get(M.getContext(), APInt(32, Sig.ParamOffset))});
+  Globals.emplace_back(buildInputSingature(M, EmptySig));
+  Globals.emplace_back(buildOutputSingature(M, EmptySig));
+}
+
 char DXContainerGlobals::ID = 0;
-INITIALIZE_PASS_BEGIN(DXContainerGlobals, "dxil-globals",
+INITIALIZE_PASS_BEGIN(DXContainerGlobals, DEBUG_TYPE,
                       "DXContainer Global Emitter", false, true)
 INITIALIZE_PASS_DEPENDENCY(ShaderFlagsAnalysisWrapper)
-INITIALIZE_PASS_END(DXContainerGlobals, "dxil-globals",
+INITIALIZE_PASS_END(DXContainerGlobals, DEBUG_TYPE,
                     "DXContainer Global Emitter", false, true)
 
 ModulePass *llvm::createDXContainerGlobalsPass() {
   return new DXContainerGlobals();
 }

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeDirectXTarget() {
   auto *PR = PassRegistry::getPassRegistry();
   initializeDXILPrepareModulePass(*PR);
   initializeEmbedDXILPassPass(*PR);
   initializeWriteDXILPassPass(*PR);
   initializeDXILOpLoweringLegacyPass(*PR);
   initializeDXILTranslateMetadataPass(*PR);
   initializeDXILResourceWrapperPass(*PR);
   initializeShaderFlagsAnalysisWrapperPass(*PR);
+  initializeDXContainerGlobalsPass(*PR);
 }
 
 class DXILTargetObjectFile : public TargetLoweringObjectFile {
 public:
   DXILTargetObjectFile() = default;
 
   MCSection *getExplicitSectionGlobal(const GlobalObject *GO, SectionKind Kind,
                                       const TargetMachine &TM) const override {

llvm/test/CodeGen/DirectX/ContainerData/EmptyProgramSignature.ll

+; RUN:opt -S -dxil-globals 2>&1 %s | FileCheck %s
+
+; Make sure empty program signature generated for compute shader.
+; CHECK:@dx.isg1 = private constant { i32, i32 } { i32 0, i32 8 }, section "ISG1", align 4
+; CHECK:@dx.osg1 = private constant { i32, i32 } { i32 0, i32 8 }, section "OSG1", align 4
+
+target datalayout = "e-m:e-p:32:32-i1:32-i8:8-i16:16-i32:32-i64:64-f16:16-f32:32-f64:64-n8:16:32:64"
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
[e-mail] = private constant [1 x i8] c"0", section "DXIL", align 4
[e-mail] = appending global [1 x ptr] [ptr @dx.dxil], section "llvm.metadata"
+
+; Function Attrs: nounwind memory(readwrite, inaccessiblemem: read)
+define void @CSMain() local_unnamed_addr #0 {
+  ret void
+}
+
+
+attributes #0 = { nounwind memory(readwrite, inaccessiblemem: read) }
+
+!dx.valver = !{!0}
+!llvm.ident = !{!1}
+!dx.shaderModel = !{!2}
+!dx.version = !{!3}
+
+!dx.entryPoints = !{!4}
+!llvm.module.flags = !{!7}
+
+!0 = !{i32 1, i32 7}
+!1 = !{!"clang version 16.0.0 (https://github.com/llvm/llvm-project.git 9641897a6b7b727fd5139d7bf92cb6b027037193)"}
+!2 = !{!"cs", i32 6, i32 0}
+!3 = !{i32 1, i32 0}
+
+!4 = !{ptr @CSMain, !"CSMain", null, null, !5}
+!5 = !{i32 4, !6}
+!6 = !{i32 1024, i32 1, i32 1}
+!7 = !{i32 1, !"wchar_size", i32 4}

llvm/test/tools/obj2yaml/DXContainer/EmptyProgramSignature.yaml

+# RUN: yaml2obj %s | obj2yaml | FileCheck %s
+
+# This test verifies that empty program signature works.
+--- !dxcontainer
+Header:
+  Hash:            [ 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
+                     0x0, 0x0, 0x0, 0x0, 0x0, 0x0 ]
+  Version:
+    Major:           1
+    Minor:           0
+  PartCount:       2
+Parts:
+  - Name:            ISG1
+    Size:            8
+    InputSig:
+      ParamCount:    0
+      ParamOffset:   8
+  - Name:            OSG1
+    Size:            8
+    OutputSig:
+      ParamCount:    0
+      ParamOffset:   8
+...
+
+# CHECK: Parts:
+# CHECK-NEXT:   - Name:            ISG1
+# CHECK-NEXT:     Size:            8
+# CHECK-NEXT:     InputSig:
+# CHECK-NEXT:       ParamCount:    0
+# CHECK-NEXT:       ParamOffset:   8
+# CHECK-NEXT:   - Name:            OSG1
+# CHECK-NEXT:     Size:            8
+# CHECK-NEXT:     OutputSig:
+# CHECK-NEXT:       ParamCount:    0
+# CHECK-NEXT:       ParamOffset:   8

llvm/tools/obj2yaml/dxcontainer2yaml.cpp

case dxbc::PartType::SFI0: {
       break;
     }
     case dxbc::PartType::HASH: {
       std::optional<dxbc::ShaderHash> Hash = Container.getShaderHash();
       if (Hash && Hash->isPopulated())
         NewPart.Hash = DXContainerYAML::ShaderHash(*Hash);
       break;
     }
+    case dxbc::PartType::ISG1: {
+      std::optional<dxbc::ProgramSignature> Sig = Container.getInputSignature();
+      if (Sig)
+        NewPart.InputSig = DXContainerYAML::ProgramSignature(*Sig);
+    } break;
+    case dxbc::PartType::OSG1: {
+      std::optional<dxbc::ProgramSignature> Sig =
+          Container.getOutputSignature();
+      if (Sig)
+        NewPart.OutputSig = DXContainerYAML::ProgramSignature(*Sig);
+    } break;
     case dxbc::PartType::Unknown:
       break;
     }
   }
 
   return Obj.release();
 }