refactoring root signature analysis to return a map instead

Author: joaosaffran

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -27,6 +27,7 @@
 #include "llvm/MC/DXContainerRootSignature.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/MD5.h"
+#include "llvm/TargetParser/Triple.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 #include <optional>
 
@@ -153,12 +154,23 @@ void DXContainerGlobals::addSignature(Module &M,
 void DXContainerGlobals::addRootSignature(Module &M,
                                           SmallVector<GlobalValue *> &Globals) {
 
+  dxil::ModuleMetadataInfo &MMI =
+      getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+
+  // Root Signature in Library shaders are different,
+  // since they don't use DXContainer to share it.
+  if (MMI.ShaderProfile == llvm::Triple::Library)
+    return;
+
+  assert(MMI.EntryPropertyVec.size() == 1);
+
   auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>();
+  const Function *&EntryFunction = MMI.EntryPropertyVec[0].Entry;
 
-  if (!RSA.getResult())
+  if (!RSA.hasForFunction(EntryFunction))
     return;
 
-  const ModuleRootSignature &MRS = RSA.getResult().value();
+  const ModuleRootSignature &MRS = RSA.getForFunction(EntryFunction);
   SmallString<256> Data;
   raw_svector_ostream OS(Data);
 

llvm/lib/Target/DirectX/DXILRootSignature.cpp

@@ -26,7 +26,9 @@
 #include "llvm/Pass.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <cstdint>
 #include <optional>
+#include <utility>
 
 using namespace llvm;
 using namespace llvm::dxil;
@@ -37,20 +39,20 @@ static bool reportError(LLVMContext *Ctx, Twine Message,
   return true;
 }
 
-static bool parseRootFlags(LLVMContext *Ctx, ModuleRootSignature *MRS,
+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();
+  MRS.Flags = Flag->getZExtValue();
 
   return false;
 }
 
 static bool parseRootSignatureElement(LLVMContext *Ctx,
-                                      ModuleRootSignature *MRS,
+                                      ModuleRootSignature &MRS,
                                       MDNode *Element) {
   MDString *ElementText = cast<MDString>(Element->getOperand(0));
   if (ElementText == nullptr)
@@ -73,8 +75,7 @@ static bool parseRootSignatureElement(LLVMContext *Ctx,
   llvm_unreachable("Root signature element kind not expected.");
 }
 
-static bool parse(LLVMContext *Ctx, ModuleRootSignature *MRS, NamedMDNode *Root,
-                  const Function *EntryFunction) {
+static bool parse(LLVMContext *Ctx, ModuleRootSignature &MRS, MDNode *Node) {
   bool HasError = false;
 
   /** Root Signature are specified as following in the metadata:
@@ -89,15 +90,46 @@ static bool parse(LLVMContext *Ctx, ModuleRootSignature *MRS, NamedMDNode *Root,
       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.");
+  // Get the Root Signature Description from the function signature pair.
+
+  // Loop through the Root Elements of the root signature.
+  for (const auto &Operand : Node->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, const ModuleRootSignature &MRS) {
+  if (!dxbc::RootSignatureValidations::isValidRootFlag(MRS.Flags)) {
+    return reportError(Ctx, "Invalid Root Signature flag value");
+  }
+  return false;
+}
+
+static SmallDenseMap<const Function *, ModuleRootSignature>
+analyzeModule(Module &M) {
+
+  LLVMContext *Ctx = &M.getContext();
+
+  SmallDenseMap<const Function *, ModuleRootSignature> MRSMap;
+
+  NamedMDNode *RootSignatureNode = M.getNamedMetadata("dx.rootsignatures");
+  if (RootSignatureNode == nullptr)
+    return MRSMap;
+
+  for (const auto &RSDefNode : RootSignatureNode->operands()) {
+    if (RSDefNode->getNumOperands() != 2) {
+      reportError(Ctx, "Invalid format for Root Signature Definition. Pairs "
+                       "of function, root signature expected.");
       continue;
     }
 
-    const MDOperand &FunctionPointerMdNode = Node->getOperand(0);
+    const MDOperand &FunctionPointerMdNode = RSDefNode->getOperand(0);
     if (FunctionPointerMdNode == nullptr) {
       // Function was pruned during compilation.
       continue;
@@ -106,97 +138,76 @@ static bool parse(LLVMContext *Ctx, ModuleRootSignature *MRS, NamedMDNode *Root,
     ValueAsMetadata *VAM =
         llvm::dyn_cast<ValueAsMetadata>(FunctionPointerMdNode.get());
     if (VAM == nullptr) {
-      HasError =
-          reportError(Ctx, "First element of root signature is not a value");
+      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");
+      reportError(Ctx, "First element of root signature is not a function");
       continue;
     }
 
-    if (F != EntryFunction)
-      continue;
+    MDNode *RootElementListNode =
+        dyn_cast<MDNode>(RSDefNode->getOperand(1).get());
 
-    // Get the Root Signature Description from the function signature pair.
-    MDNode *RS = dyn_cast<MDNode>(Node->getOperand(1).get());
-
-    if (RS == nullptr) {
+    if (RootElementListNode == 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.");
+    ModuleRootSignature MRS;
 
-      HasError = HasError || parseRootSignatureElement(Ctx, MRS, Element);
+    if (parse(Ctx, MRS, RootElementListNode) || validate(Ctx, MRS)) {
+      return MRSMap;
     }
-  }
-  return HasError;
-}
 
-static bool validate(LLVMContext *Ctx, ModuleRootSignature *MRS) {
-  if (!dxbc::RootSignatureValidations::isValidRootFlag(MRS->Flags)) {
-    return reportError(Ctx, "Invalid Root Signature flag value");
+    MRSMap.insert(std::make_pair(F, MRS));
   }
-  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;
+  return MRSMap;
 }
 
-std::optional<ModuleRootSignature>
-ModuleRootSignature::analyzeModule(Module &M, const Function *F) {
-
-  LLVMContext *Ctx = &M.getContext();
+AnalysisKey RootSignatureAnalysis::Key;
 
-  ModuleRootSignature MRS;
+SmallDenseMap<const Function *, ModuleRootSignature>
+RootSignatureAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
+  return analyzeModule(M);
+}
 
-  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;
+static void printSpaces(raw_ostream &Stream, unsigned int Count) {
+  for (unsigned int I = 0; I < Count; ++I) {
+    Stream << ' ';
   }
-
-  return MRS;
 }
 
-AnalysisKey RootSignatureAnalysis::Key;
+PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
+                                                    ModuleAnalysisManager &AM) {
+
+  SmallDenseMap<const Function *, ModuleRootSignature> &MRSMap =
+      AM.getResult<RootSignatureAnalysis>(M);
+  OS << "Root Signature Definitions"
+     << "\n";
+  uint8_t Space = 0;
+  for (const auto &P : MRSMap) {
+    const auto &[Function, MRS] = P;
+    OS << "Definition for '" << Function->getName() << "':\n";
+
+    // start root signature header
+    Space++;
+    printSpaces(OS, Space);
+    OS << "Flags: " << format_hex(MRS.Flags, 8) << ":\n";
+    Space--;
+    // end root signature header
+  }
 
-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));
+  return PreservedAnalyses::all();
 }
 
 //===----------------------------------------------------------------------===//
 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));
+  MRS = analyzeModule(M);
   return false;
 }
 
@@ -208,8 +219,8 @@ void RootSignatureAnalysisWrapper::getAnalysisUsage(AnalysisUsage &AU) const {
 char RootSignatureAnalysisWrapper::ID = 0;
 
 INITIALIZE_PASS_BEGIN(RootSignatureAnalysisWrapper,
-                      "dx-root-signature-analysis",
+                      "dxil-root-signature-analysis",
                       "DXIL Root Signature Analysis", true, true)
-INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
-INITIALIZE_PASS_END(RootSignatureAnalysisWrapper, "dx-root-signature-analysis",
+INITIALIZE_PASS_END(RootSignatureAnalysisWrapper,
+                    "dxil-root-signature-analysis",
                     "DXIL Root Signature Analysis", true, true)

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -11,6 +11,7 @@
 ///
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Metadata.h"
@@ -25,9 +26,8 @@ namespace dxil {
 enum class RootSignatureElementKind { None = 0, RootFlags = 1 };
 
 struct ModuleRootSignature {
+  ModuleRootSignature() = default;
   uint32_t Flags = 0;
-  static std::optional<ModuleRootSignature> analyzeModule(Module &M,
-                                                          const Function *F);
 };
 
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
@@ -37,9 +37,10 @@ class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
 public:
   RootSignatureAnalysis() = default;
 
-  using Result = std::optional<ModuleRootSignature>;
+  using Result = SmallDenseMap<const Function *, ModuleRootSignature>;
 
-  std::optional<ModuleRootSignature> run(Module &M, ModuleAnalysisManager &AM);
+  SmallDenseMap<const Function *, ModuleRootSignature>
+  run(Module &M, ModuleAnalysisManager &AM);
 };
 
 /// Wrapper pass for the legacy pass manager.
@@ -48,19 +49,34 @@ class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
 /// passes which run through the legacy pass manager.
 class RootSignatureAnalysisWrapper : public ModulePass {
 private:
-  std::optional<ModuleRootSignature> MRS;
+  SmallDenseMap<const Function *, ModuleRootSignature> MRS;
 
 public:
   static char ID;
 
   RootSignatureAnalysisWrapper() : ModulePass(ID) {}
 
-  const std::optional<ModuleRootSignature> &getResult() const { return MRS; }
+  bool hasForFunction(const Function *F) { return MRS.find(F) != MRS.end(); }
+
+  ModuleRootSignature getForFunction(const Function *F) {
+    assert(hasForFunction(F));
+    return MRS[F];
+  }
 
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;
 };
 
+/// Printer pass for RootSignatureAnalysis results.
+class RootSignatureAnalysisPrinter
+    : public PassInfoMixin<RootSignatureAnalysisPrinter> {
+  raw_ostream &OS;
+
+public:
+  explicit RootSignatureAnalysisPrinter(raw_ostream &OS) : OS(OS) {}
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
 } // namespace dxil
 } // namespace llvm

llvm/lib/Target/DirectX/DirectXPassRegistry.def

@@ -18,6 +18,7 @@
 #endif
 MODULE_ANALYSIS("dx-shader-flags", dxil::ShaderFlagsAnalysis())
 MODULE_ANALYSIS("dxil-resource-md", DXILResourceMDAnalysis())
+MODULE_ANALYSIS("dxil-root-signature-analysis", dxil::RootSignatureAnalysis())
 #undef MODULE_ANALYSIS
 
 #ifndef MODULE_PASS
@@ -31,6 +32,7 @@ MODULE_PASS("dxil-pretty-printer", DXILPrettyPrinterPass(dbgs()))
 MODULE_PASS("dxil-translate-metadata", DXILTranslateMetadata())
 // TODO: rename to print<foo> after NPM switch
 MODULE_PASS("print-dx-shader-flags", dxil::ShaderFlagsAnalysisPrinter(dbgs()))
+MODULE_PASS("print<dxil-root-signature>", dxil::RootSignatureAnalysisPrinter(dbgs()))
 #undef MODULE_PASS
 
 #ifndef FUNCTION_PASS

llvm/lib/Target/DirectX/DirectXTargetMachine.cpp

@@ -19,6 +19,7 @@
 #include "DXILPrettyPrinter.h"
 #include "DXILResourceAccess.h"
 #include "DXILResourceAnalysis.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DXILTranslateMetadata.h"
 #include "DXILWriter/DXILWriterPass.h"

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Error.ll

@@ -1,8 +1,9 @@
-; RUN: not llc %s --filetype=obj -o - 2>&1 | FileCheck %s
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
 
 target triple = "dxil-unknown-shadermodel6.0-compute"
 
 ; CHECK: error: Invalid format for Root Signature Definition. Pairs of function, root signature expected.
+; CHECK-NO: Root Signature Definitions
 
 
 define void @main() #0 {

llvm/test/CodeGen/DirectX/ContainerData/RootSignature-Flags-Error.ll

@@ -1,8 +1,9 @@
-; RUN: not llc %s --filetype=obj -o - 2>&1 | FileCheck %s
+; RUN: not opt -passes='print<dxil-root-signature>' %s -S -o - 2>&1 | FileCheck %s
 
 target triple = "dxil-unknown-shadermodel6.0-compute"
 
 ; CHECK: error: Invalid Root Element: NOTRootFlags
+; CHECK-NO: Root Signature Definitions
 
 
 define void @main() #0 {