[DirectX][NFC] Refactor DXILRootSignature to follow the same pattern as other analysis

[joaosaffran]

When implementing #146785, notice DXILRootSignature had some design changes that made it harder to integrate with other analysis. This change refactors DXILRootSignature to solve this issue.

[llvmbot]

@llvm/pr-subscribers-backend-directx

Author: None (joaosaffran)

Changes

As part of #126645, DXILRootSignature analysis needs to be referenced in multiple passes, which can lead to duplication of data. To fix that, this patch makes sure the analysis only parses the data once.

---

Full diff: https://github.com/llvm/llvm-project/pull/146783.diff

3 Files Affected:

• (modified) llvm/lib/Target/DirectX/DXContainerGlobals.cpp (+4-5)
• (modified) llvm/lib/Target/DirectX/DXILRootSignature.cpp (+9-5)
• (modified) llvm/lib/Target/DirectX/DXILRootSignature.h (+37-12)

diff --git a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
index 9c38901f6821f..6c8ae8eaaea77 100644
--- a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
+++ b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
@@ -160,18 +160,17 @@ void DXContainerGlobals::addRootSignature(Module &M,
 
   assert(MMI.EntryPropertyVec.size() == 1);
 
-  auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>();
+  auto &RSA = getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
   const Function *EntryFunction = MMI.EntryPropertyVec[0].Entry;
-  const auto &FuncRs = RSA.find(EntryFunction);
+  const auto &RS = RSA.getDescForFunction(EntryFunction);
 
-  if (FuncRs == RSA.end())
+  if (!RS)
     return;
 
-  const RootSignatureDesc &RS = FuncRs->second;
   SmallString<256> Data;
   raw_svector_ostream OS(Data);
 
-  RS.write(OS);
+  RS->write(OS);
 
   Constant *Constant =
       ConstantDataArray::getString(M.getContext(), Data, /*AddNull*/ false);
diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.cpp b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
index 29e78fcce5262..5a53ea8a3631b 100644
--- a/llvm/lib/Target/DirectX/DXILRootSignature.cpp
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.cpp
@@ -554,9 +554,12 @@ analyzeModule(Module &M) {
 
 AnalysisKey RootSignatureAnalysis::Key;
 
-SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>
+RootSignatureAnalysis::Result
 RootSignatureAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
-  return analyzeModule(M);
+  if (!AnalysisResult)
+    AnalysisResult = std::make_unique<RootSignatureBindingInfo>(
+        RootSignatureBindingInfo(analyzeModule(M)));
+  return *AnalysisResult;
 }
 
 //===----------------------------------------------------------------------===//
@@ -564,8 +567,7 @@ RootSignatureAnalysis::run(Module &M, ModuleAnalysisManager &AM) {
 PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
                                                     ModuleAnalysisManager &AM) {
 
-  SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> &RSDMap =
-      AM.getResult<RootSignatureAnalysis>(M);
+  RootSignatureBindingInfo &RSDMap = AM.getResult<RootSignatureAnalysis>(M);
 
   OS << "Root Signature Definitions"
      << "\n";
@@ -636,7 +638,9 @@ PreservedAnalyses RootSignatureAnalysisPrinter::run(Module &M,
 
 //===----------------------------------------------------------------------===//
 bool RootSignatureAnalysisWrapper::runOnModule(Module &M) {
-  FuncToRsMap = analyzeModule(M);
+  if (!FuncToRsMap)
+    FuncToRsMap = std::make_unique<RootSignatureBindingInfo>(
+        RootSignatureBindingInfo(analyzeModule(M)));
   return false;
 }
 
diff --git a/llvm/lib/Target/DirectX/DXILRootSignature.h b/llvm/lib/Target/DirectX/DXILRootSignature.h
index b45cebc15fd39..8a057404e01c1 100644
--- a/llvm/lib/Target/DirectX/DXILRootSignature.h
+++ b/llvm/lib/Target/DirectX/DXILRootSignature.h
@@ -33,17 +33,46 @@ enum class RootSignatureElementKind {
   CBV = 5,
   DescriptorTable = 6,
 };
+
+class RootSignatureBindingInfo {
+  private:
+    SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> FuncToRsMap;
+
+  public:
+  using iterator =
+        SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>::iterator;
+
+  RootSignatureBindingInfo () = default;
+  RootSignatureBindingInfo(SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> Map) : FuncToRsMap(Map) {};
+
+  iterator find(const Function *F) { return FuncToRsMap.find(F); }
+
+  iterator end() { return FuncToRsMap.end(); }
+
+  std::optional<mcdxbc::RootSignatureDesc> getDescForFunction(const Function* F) {
+    const auto FuncRs = find(F);
+    if (FuncRs == end())
+      return std::nullopt;
+
+    return FuncRs->second;
+  }
+  
+};
+
 class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
   friend AnalysisInfoMixin<RootSignatureAnalysis>;
   static AnalysisKey Key;
 
 public:
-  RootSignatureAnalysis() = default;
 
-  using Result = SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>;
+RootSignatureAnalysis() = default;
+
+  using Result = RootSignatureBindingInfo;
+
+  Result run(Module &M, ModuleAnalysisManager &AM);
 
-  SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>
-  run(Module &M, ModuleAnalysisManager &AM);
+private:
+  std::unique_ptr<RootSignatureBindingInfo> AnalysisResult;
 };
 
 /// Wrapper pass for the legacy pass manager.
@@ -52,20 +81,16 @@ class RootSignatureAnalysis : public AnalysisInfoMixin<RootSignatureAnalysis> {
 /// passes which run through the legacy pass manager.
 class RootSignatureAnalysisWrapper : public ModulePass {
 private:
-  SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> FuncToRsMap;
+  std::unique_ptr<RootSignatureBindingInfo> FuncToRsMap;
 
 public:
   static char ID;
+  using Result = RootSignatureBindingInfo;
 
   RootSignatureAnalysisWrapper() : ModulePass(ID) {}
 
-  using iterator =
-      SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>::iterator;
-
-  iterator find(const Function *F) { return FuncToRsMap.find(F); }
-
-  iterator end() { return FuncToRsMap.end(); }
-
+  RootSignatureBindingInfo& getRSInfo() {return *FuncToRsMap;}
+  
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[spall]

isn't RS a reference, why are you using the arrow?

[joaosaffran]

Updated the type to make it clearer, but the type is const std::optional<mcdxbc::RootSignatureDesc> &RS

[spall]

Why create a type alias here? Why not just put 'RootSignatureBindingInfo' inline? It looks like its only used one place.

[joaosaffran]

Result type alias is required in analysis passes or passes that can return something, like in this case, it can contain the Root Signature data representation.

[spall]

This type alias doesn't even appear to be used?

[bogner]

This change isn't sound. If we're running the analysis pass again, that means that the pass manager has said it's been invalidated, and if it's been invalidated then it isn't safe to return the previous result.

What we actually need to do here is either (1) ensure that passes that don't modify the metadata don't invalidate this pass, or (2) make this an immutable pass. I suspect (1) is more correct - if passes were to modify the root signature metadata they would indeed invalidate this pass.