[DXIL] Add support for root signature flag element in DXContainer

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/lib/Object/DXContainer.cpp

@@ -254,11 +254,9 @@ 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 make_error<GenericBinaryError>("Invalid Root Signature Version");
+  Version = VValue;
 
   NumParameters =
       support::endian::read<uint32_t, llvm::endianness::little>(Current);
@@ -280,11 +278,9 @@ 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 make_error<GenericBinaryError>("Invalid Root Signature flag");
+  Flags = FValue;
 
   return Error::success();
 }

llvm/lib/Target/DirectX/CMakeLists.txt

@@ -33,7 +33,7 @@ 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"
@@ -23,9 +24,11 @@
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/MC/DXContainerPSVInfo.h"
+#include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MD5.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,28 @@ void DXContainerGlobals::addSignature(Module &M,
   Globals.emplace_back(buildSignature(M, OutputSig, "dx.osg1", "OSG1"));
 }
 
+void DXContainerGlobals::addRootSignature(Module &M,
+                                          SmallVector<GlobalValue *> &Globals) {
+
+  auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>();
+
+  if (!RSA.getResult())
+    return;
+
+  const ModuleRootSignature &MRS = RSA.getResult().value();
+  SmallString<256> Data;
+  raw_svector_ostream OS(Data);
+
+  RootSignatureHeader RSH;
+  RSH.Flags = MRS.Flags;
+
+  RSH.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();

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -0,0 +1,215 @@
+//===- DXILRootSignature.cpp - DXIL Root Signature helper objects ----===//
+//
+// 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 This file contains helper objects and APIs for working with DXIL
+///       Root Signatures.
+///
+//===----------------------------------------------------------------------===//
+#include "DXILRootSignature.h"
+#include "DirectX.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/Twine.h"
+#include "llvm/Analysis/DXILMetadataAnalysis.h"
+#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DiagnosticInfo.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include <optional>
+
+using namespace llvm;
+using namespace llvm::dxil;
+
+static bool reportError(LLVMContext *Ctx, Twine Message,
+                        DiagnosticSeverity Severity = DS_Error) {
+  Ctx->diagnose(DiagnosticInfoGeneric(Message, Severity));
+  return true;
+}
+
+static bool parseRootFlags(LLVMContext *Ctx, ModuleRootSignature *MRS,
+                           MDNode *RootFlagNode) {
+
+  if (RootFlagNode->getNumOperands() != 2)
+    return reportError(Ctx, "Invalid format for RootFlag Element");
+
+  auto *Flag = mdconst::extract<ConstantInt>(RootFlagNode->getOperand(1));
+  MRS->Flags = Flag->getZExtValue();
+
+  return false;
+}
+
+static bool parseRootSignatureElement(LLVMContext *Ctx,
+                                      ModuleRootSignature *MRS,
+                                      MDNode *Element) {
+  MDString *ElementText = cast<MDString>(Element->getOperand(0));
+  if (ElementText == nullptr)
+    return reportError(Ctx, "Invalid format for Root Element");
+
+  RootSignatureElementKind ElementKind =
+      StringSwitch<RootSignatureElementKind>(ElementText->getString())
+          .Case("RootFlags", RootSignatureElementKind::RootFlags)
+          .Default(RootSignatureElementKind::None);
+
+  switch (ElementKind) {
+
+  case RootSignatureElementKind::RootFlags:
+    return parseRootFlags(Ctx, MRS, Element);
+  case RootSignatureElementKind::None:
+    return reportError(Ctx,
+                       "Invalid Root Element: " + ElementText->getString());
+  }
+
+  llvm_unreachable("Root signature element kind not expected.");
+}
+
+static bool parse(LLVMContext *Ctx, ModuleRootSignature *MRS, NamedMDNode *Root,
+                  const Function *EntryFunction) {
+  bool HasError = false;
+
+  /** Root Signature are specified as following in the metadata:
+
+      !dx.rootsignatures = !{!2} ; list of function/root signature pairs
+      !2 = !{ ptr @main, !3 } ; function, root signature
+      !3 = !{ !4, !5, !6, !7 } ; list of root signature elements
+
+      So for each MDNode inside dx.rootsignatures NamedMDNode
+      (the Root parameter of this function), the parsing process needs
+      to loop through each of its operands and process the function,
+      signature pair.
+   */
+
+  for (const MDNode *Node : Root->operands()) {
+    if (Node->getNumOperands() != 2) {
+      HasError = reportError(
+          Ctx, "Invalid format for Root Signature Definition. Pairs "
+               "of function, root signature expected.");
+      continue;
+    }
+
+    const MDOperand &FunctionPointerMdNode = Node->getOperand(0);
+    if (FunctionPointerMdNode == nullptr) {
+      // Function was pruned during compilation.
+      continue;
+    }
+
+    ValueAsMetadata *VAM =
+        llvm::dyn_cast<ValueAsMetadata>(FunctionPointerMdNode.get());
+    if (VAM == nullptr) {
+      HasError =
+          reportError(Ctx, "First element of root signature is not a value");
+      continue;
+    }
+
+    Function *F = dyn_cast<Function>(VAM->getValue());
+    if (F == nullptr) {
+      HasError =
+          reportError(Ctx, "First element of root signature is not a function");
+      continue;
+    }
+
+    if (F != EntryFunction)
+      continue;
+
+    // Get the Root Signature Description from the function signature pair.
+    MDNode *RS = dyn_cast<MDNode>(Node->getOperand(1).get());
+
+    if (RS == nullptr) {
+      reportError(Ctx, "Missing Root Element List Metadata node.");
+      continue;
+    }
+
+    // Loop through the Root Elements of the root signature.
+    for (const auto &Operand : RS->operands()) {
+      MDNode *Element = dyn_cast<MDNode>(Operand);
+      if (Element == nullptr)
+        return reportError(Ctx, "Missing Root Element Metadata Node.");
+
+      HasError = HasError || parseRootSignatureElement(Ctx, MRS, Element);
+    }
+  }
+  return HasError;
+}
+
+static bool validate(LLVMContext *Ctx, ModuleRootSignature *MRS) {
+  if (!dxbc::RootSignatureValidations::isValidRootFlag(MRS->Flags)) {
+    return reportError(Ctx, "Invalid Root Signature flag value");
+  }
+  return false;
+}
+
+static const Function *getEntryFunction(Module &M, ModuleMetadataInfo MMI) {
+
+  LLVMContext *Ctx = &M.getContext();
+  if (MMI.EntryPropertyVec.size() != 1) {
+    reportError(Ctx, "More than one entry function defined.");
+    // needed to stop compilation
+    report_fatal_error("Invalid Root Signature Definition", false);
+    return nullptr;
+  }
+  return MMI.EntryPropertyVec[0].Entry;
+}
+
+std::optional<ModuleRootSignature>
+ModuleRootSignature::analyzeModule(Module &M, const Function *F) {
+
+  LLVMContext *Ctx = &M.getContext();
+
+  ModuleRootSignature MRS;
+
+  NamedMDNode *RootSignatureNode = M.getNamedMetadata("dx.rootsignatures");
+  if (RootSignatureNode == nullptr)
+    return std::nullopt;
+
+  if (parse(Ctx, &MRS, RootSignatureNode, F) || validate(Ctx, &MRS)) {
+    // needed to stop compilation
+    report_fatal_error("Invalid Root Signature Definition", false);
+    return std::nullopt;
+  }
+
+  return MRS;
+}
+
+AnalysisKey RootSignatureAnalysis::Key;
+
+std::optional<ModuleRootSignature>
+RootSignatureAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
+  ModuleMetadataInfo MMI = AM.getResult<DXILMetadataAnalysis>(M);
+  if (MMI.ShaderProfile == Triple::Library)
+    return std::nullopt;
+  return ModuleRootSignature::analyzeModule(M, getEntryFunction(M, MMI));
+}
+
+//===----------------------------------------------------------------------===//
+bool RootSignatureAnalysisWrapper::runOnModule(Module &M) {
+  dxil::ModuleMetadataInfo &MMI =
+      getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+  if (MMI.ShaderProfile == Triple::Library)
+    return false;
+  MRS = ModuleRootSignature::analyzeModule(M, getEntryFunction(M, MMI));
+  return false;
+}
+
+void RootSignatureAnalysisWrapper::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.setPreservesAll();
+  AU.addRequired<DXILMetadataAnalysisWrapperPass>();
+}
+
+char RootSignatureAnalysisWrapper::ID = 0;
+
+INITIALIZE_PASS_BEGIN(RootSignatureAnalysisWrapper,
+                      "dx-root-signature-analysis",
+                      "DXIL Root Signature Analysis", true, true)
+INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
+INITIALIZE_PASS_END(RootSignatureAnalysisWrapper, "dx-root-signature-analysis",
+                    "DXIL Root Signature Analysis", true, true)