[DXIL] Add support for root signature flag element in DXContainer (#123147)

Adding support for Root Signature Flags Element extraction and writing
to DXContainer.
- Adding an analysis to deal with RootSignature metadata definition
- Adding validation for Flag
- writing RootSignature blob into DXIL

Closes: [126632](#126632)

---------

Co-authored-by: joaosaffran <[email protected]>

Author: joaosaffran

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -14,8 +14,6 @@
 #define LLVM_BINARYFORMAT_DXCONTAINER_H
 
 #include "llvm/ADT/StringRef.h"
-#include "llvm/Support/BinaryStreamError.h"
-#include "llvm/Support/Error.h"
 #include "llvm/Support/SwapByteOrder.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -550,18 +548,10 @@ static_assert(sizeof(ProgramSignatureElement) == 32,
 
 struct RootSignatureValidations {
 
-  static Expected<uint32_t> validateRootFlag(uint32_t Flags) {
-    if ((Flags & ~0x80000fff) != 0)
-      return llvm::make_error<BinaryStreamError>("Invalid Root Signature flag");
-    return Flags;
-  }
-
-  static Expected<uint32_t> validateVersion(uint32_t Version) {
-    if (Version == 1 || Version == 2)
-      return Version;
+  static bool isValidRootFlag(uint32_t Flags) { return (Flags & ~0xfff) == 0; }
 
-    return llvm::make_error<BinaryStreamError>(
-        "Invalid Root Signature Version");
+  static bool isValidVersion(uint32_t Version) {
+    return (Version == 1 || Version == 2);
   }
 };
 

llvm/include/llvm/MC/DXContainerRootSignature.h

@@ -14,15 +14,15 @@ namespace llvm {
 class raw_ostream;
 
 namespace mcdxbc {
-struct RootSignatureHeader {
+struct RootSignatureDesc {
   uint32_t Version = 2;
   uint32_t NumParameters = 0;
   uint32_t RootParametersOffset = 0;
   uint32_t NumStaticSamplers = 0;
   uint32_t StaticSamplersOffset = 0;
   uint32_t Flags = 0;
 
-  void write(raw_ostream &OS);
+  void write(raw_ostream &OS) const;
 };
 } // namespace mcdxbc
 } // namespace llvm

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -74,9 +74,9 @@ struct ShaderHash {
 };
 
 #define ROOT_ELEMENT_FLAG(Num, Val) bool Val = false;
-struct RootSignatureDesc {
-  RootSignatureDesc() = default;
-  RootSignatureDesc(const object::DirectX::RootSignature &Data);
+struct RootSignatureYamlDesc {
+  RootSignatureYamlDesc() = default;
+  RootSignatureYamlDesc(const object::DirectX::RootSignature &Data);
 
   uint32_t Version;
   uint32_t NumParameters;
@@ -176,7 +176,7 @@ struct Part {
   std::optional<ShaderHash> Hash;
   std::optional<PSVInfo> Info;
   std::optional<DXContainerYAML::Signature> Signature;
-  std::optional<DXContainerYAML::RootSignatureDesc> RootSignature;
+  std::optional<DXContainerYAML::RootSignatureYamlDesc> RootSignature;
 };
 
 struct Object {
@@ -259,9 +259,9 @@ template <> struct MappingTraits<DXContainerYAML::Signature> {
   static void mapping(IO &IO, llvm::DXContainerYAML::Signature &El);
 };
 
-template <> struct MappingTraits<DXContainerYAML::RootSignatureDesc> {
+template <> struct MappingTraits<DXContainerYAML::RootSignatureYamlDesc> {
   static void mapping(IO &IO,
-                      DXContainerYAML::RootSignatureDesc &RootSignature);
+                      DXContainerYAML::RootSignatureYamlDesc &RootSignature);
 };
 
 } // namespace yaml

llvm/lib/MC/DXContainerRootSignature.cpp

@@ -12,7 +12,7 @@
 using namespace llvm;
 using namespace llvm::mcdxbc;
 
-void RootSignatureHeader::write(raw_ostream &OS) {
+void RootSignatureDesc::write(raw_ostream &OS) const {
 
   support::endian::write(OS, Version, llvm::endianness::little);
   support::endian::write(OS, NumParameters, llvm::endianness::little);

llvm/lib/Object/DXContainer.cpp

@@ -20,6 +20,10 @@ static Error parseFailed(const Twine &Msg) {
   return make_error<GenericBinaryError>(Msg.str(), object_error::parse_failed);
 }
 
+static Error validationFailed(const Twine &Msg) {
+  return make_error<StringError>(Msg.str(), inconvertibleErrorCode());
+}
+
 template <typename T>
 static Error readStruct(StringRef Buffer, const char *Src, T &Struct) {
   // Don't read before the beginning or past the end of the file
@@ -254,11 +258,10 @@ Error DirectX::RootSignature::parse(StringRef Data) {
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
-  Expected<uint32_t> MaybeVersion =
-      dxbc::RootSignatureValidations::validateVersion(VValue);
-  if (Error E = MaybeVersion.takeError())
-    return E;
-  Version = MaybeVersion.get();
+  if (!dxbc::RootSignatureValidations::isValidVersion(VValue))
+    return validationFailed("unsupported root signature version read: " +
+                            llvm::Twine(VValue));
+  Version = VValue;
 
   NumParameters =
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
@@ -280,11 +283,10 @@ Error DirectX::RootSignature::parse(StringRef Data) {
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
   Current += sizeof(uint32_t);
 
-  Expected<uint32_t> MaybeFlag =
-      dxbc::RootSignatureValidations::validateRootFlag(FValue);
-  if (Error E = MaybeFlag.takeError())
-    return E;
-  Flags = MaybeFlag.get();
+  if (!dxbc::RootSignatureValidations::isValidRootFlag(FValue))
+    return validationFailed("unsupported root signature flag value read: " +
+                            llvm::Twine(FValue));
+  Flags = FValue;
 
   return Error::success();
 }

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -266,15 +266,15 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       if (!P.RootSignature.has_value())
         continue;
 
-      mcdxbc::RootSignatureHeader Header;
-      Header.Flags = P.RootSignature->getEncodedFlags();
-      Header.Version = P.RootSignature->Version;
-      Header.NumParameters = P.RootSignature->NumParameters;
-      Header.RootParametersOffset = P.RootSignature->RootParametersOffset;
-      Header.NumStaticSamplers = P.RootSignature->NumStaticSamplers;
-      Header.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
-
-      Header.write(OS);
+      mcdxbc::RootSignatureDesc RS;
+      RS.Flags = P.RootSignature->getEncodedFlags();
+      RS.Version = P.RootSignature->Version;
+      RS.NumParameters = P.RootSignature->NumParameters;
+      RS.RootParametersOffset = P.RootSignature->RootParametersOffset;
+      RS.NumStaticSamplers = P.RootSignature->NumStaticSamplers;
+      RS.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
+
+      RS.write(OS);
       break;
     }
     uint64_t BytesWritten = OS.tell() - DataStart;

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -29,7 +29,7 @@ DXContainerYAML::ShaderFeatureFlags::ShaderFeatureFlags(uint64_t FlagData) {
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
-DXContainerYAML::RootSignatureDesc::RootSignatureDesc(
+DXContainerYAML::RootSignatureYamlDesc::RootSignatureYamlDesc(
     const object::DirectX::RootSignature &Data)
     : Version(Data.getVersion()), NumParameters(Data.getNumParameters()),
       RootParametersOffset(Data.getRootParametersOffset()),
@@ -41,7 +41,7 @@ DXContainerYAML::RootSignatureDesc::RootSignatureDesc(
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 }
 
-uint32_t DXContainerYAML::RootSignatureDesc::getEncodedFlags() {
+uint32_t DXContainerYAML::RootSignatureYamlDesc::getEncodedFlags() {
   uint64_t Flag = 0;
 #define ROOT_ELEMENT_FLAG(Num, Val)                                            \
   if (Val)                                                                     \
@@ -209,8 +209,8 @@ void MappingTraits<DXContainerYAML::Signature>::mapping(
   IO.mapRequired("Parameters", S.Parameters);
 }
 
-void MappingTraits<DXContainerYAML::RootSignatureDesc>::mapping(
-    IO &IO, DXContainerYAML::RootSignatureDesc &S) {
+void MappingTraits<DXContainerYAML::RootSignatureYamlDesc>::mapping(
+    IO &IO, DXContainerYAML::RootSignatureYamlDesc &S) {
   IO.mapRequired("Version", S.Version);
   IO.mapRequired("NumParameters", S.NumParameters);
   IO.mapRequired("RootParametersOffset", S.RootParametersOffset);

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -33,7 +33,8 @@ add_llvm_target(DirectXCodeGen
   DXILResourceAccess.cpp
   DXILShaderFlags.cpp
   DXILTranslateMetadata.cpp
-
+  DXILRootSignature.cpp
+  
   LINK_COMPONENTS
   Analysis
   AsmPrinter

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
 #include "llvm/ADT/SmallVector.h"
@@ -25,7 +26,9 @@
 #include "llvm/MC/DXContainerPSVInfo.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MD5.h"
+#include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <optional>
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -41,6 +44,7 @@ class DXContainerGlobals : public llvm::ModulePass {
   GlobalVariable *buildSignature(Module &M, Signature &Sig, StringRef Name,
                                  StringRef SectionName);
   void addSignature(Module &M, SmallVector<GlobalValue *> &Globals);
+  void addRootSignature(Module &M, SmallVector<GlobalValue *> &Globals);
   void addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV);
   void addPipelineStateValidationInfo(Module &M,
                                       SmallVector<GlobalValue *> &Globals);
@@ -60,6 +64,7 @@ class DXContainerGlobals : public llvm::ModulePass {
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     AU.addRequired<ShaderFlagsAnalysisWrapper>();
+    AU.addRequired<RootSignatureAnalysisWrapper>();
     AU.addRequired<DXILMetadataAnalysisWrapperPass>();
     AU.addRequired<DXILResourceTypeWrapperPass>();
     AU.addRequired<DXILResourceBindingWrapperPass>();
@@ -73,6 +78,7 @@ bool DXContainerGlobals::runOnModule(Module &M) {
   Globals.push_back(getFeatureFlags(M));
   Globals.push_back(computeShaderHash(M));
   addSignature(M, Globals);
+  addRootSignature(M, Globals);
   addPipelineStateValidationInfo(M, Globals);
   appendToCompilerUsed(M, Globals);
   return true;
@@ -144,6 +150,36 @@ void DXContainerGlobals::addSignature(Module &M,
   Globals.emplace_back(buildSignature(M, OutputSig, "dx.osg1", "OSG1"));
 }
 
+void DXContainerGlobals::addRootSignature(Module &M,
+                                          SmallVector<GlobalValue *> &Globals) {
+
+  dxil::ModuleMetadataInfo &MMI =
+      getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+
+  // Root Signature in Library don't compile to DXContainer.
+  if (MMI.ShaderProfile == llvm::Triple::Library)
+    return;
+
+  assert(MMI.EntryPropertyVec.size() == 1);
+
+  auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>();
+  const Function *EntryFunction = MMI.EntryPropertyVec[0].Entry;
+  const auto &FuncRs = RSA.find(EntryFunction);
+
+  if (FuncRs == RSA.end())
+    return;
+
+  const RootSignatureDesc &RS = FuncRs->second;
+  SmallString<256> Data;
+  raw_svector_ostream OS(Data);
+
+  RS.write(OS);
+
+  Constant *Constant =
+      ConstantDataArray::getString(M.getContext(), Data, /*AddNull*/ false);
+  Globals.emplace_back(buildContainerGlobal(M, Constant, "dx.rts0", "RTS0"));
+}
+
 void DXContainerGlobals::addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV) {
   const DXILBindingMap &DBM =
       getAnalysis<DXILResourceBindingWrapperPass>().getBindingMap();