init refactoring

Author: joaosaffran

clang/test/SemaHLSL/RootSignature-Validation.hlsl

@@ -0,0 +1,42 @@
+// RUN: %clang_dxc -triple dxil-pc-shadermodel6.3-library -x hlsl -o - %s -verify
+
+#define ROOT_SIGNATURE \
+    "RootFlags(ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT), " \
+    "CBV(b0, visibility=SHADER_VISIBILITY_ALL), " \
+    "DescriptorTable(SRV(t0, numDescriptors=3), visibility=SHADER_VISIBILITY_PIXEL), " \
+    "DescriptorTable(Sampler(s0, numDescriptors=2), visibility=SHADER_VISIBILITY_PIXEL), " \
+    "DescriptorTable(UAV(u0, numDescriptors=1), visibility=SHADER_VISIBILITY_ALL)"
+
+cbuffer CB : register(b3, space2) {
+  float a;
+}
+
+StructuredBuffer<int> In : register(t0);
+RWStructuredBuffer<int> Out : register(u0);
+
+RWBuffer<float> UAV : register(u3);
+
+RWBuffer<float> UAV1 : register(u2), UAV2 : register(u4);
+
+RWBuffer<float> UAV3 : register(space5);
+
+float f : register(c5);
+
+int4 intv : register(c2);
+
+double dar[5] :  register(c3);
+
+struct S {
+  int a;
+};
+
+S s : register(c10);
+
+// Compute Shader for UAV testing
+[numthreads(8, 8, 1)]
+[RootSignature(ROOT_SIGNATURE)]
+void CSMain(uint3 id : SV_DispatchThreadID)
+{
+    In[0] = id;
+    Out[0] = In[0];
+}

llvm/lib/Target/DirectX/DXContainerGlobals.cpp

@@ -164,7 +164,7 @@ void DXContainerGlobals::addRootSignature(Module &M,
   const Function *EntryFunction = MMI.EntryPropertyVec[0].Entry;
   const auto &RS = RSA.getDescForFunction(EntryFunction);
 
-  if (!RS )
+  if (!RS)
     return;
 
   SmallString<256> Data;

llvm/lib/Target/DirectX/DXILPostOptimizationValidation.cpp

@@ -7,11 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "DXILPostOptimizationValidation.h"
+#include "DXILRootSignature.h"
 #include "DXILShaderFlags.h"
 #include "DirectX.h"
+#include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
 #include "llvm/Analysis/DXILResource.h"
+#include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicsDirectX.h"
@@ -85,7 +88,9 @@ static void reportOverlappingBinding(Module &M, DXILResourceMap &DRM) {
 }
 
 static void reportErrors(Module &M, DXILResourceMap &DRM,
-                         DXILResourceBindingInfo &DRBI) {
+                         DXILResourceBindingInfo &DRBI,
+                         RootSignatureBindingInfo &RSBI,
+                         dxil::ModuleMetadataInfo &MMI) {
   if (DRM.hasInvalidCounterDirection())
     reportInvalidDirection(M, DRM);
 
@@ -94,14 +99,41 @@ static void reportErrors(Module &M, DXILResourceMap &DRM,
 
   assert(!DRBI.hasImplicitBinding() && "implicit bindings should be handled in "
                                        "DXILResourceImplicitBinding pass");
+  // Assuming this is used to validate only the root signature assigned to the
+  // entry function.
+  std::optional<mcdxbc::RootSignatureDesc> RootSigDesc =
+      RSBI.getDescForFunction(MMI.EntryPropertyVec[0].Entry);
+  if (!RootSigDesc)
+    return;
+
+  for (const mcdxbc::RootParameterInfo &Info :
+       RootSigDesc->ParametersContainer) {
+    const auto &[Type, Loc] =
+        RootSigDesc->ParametersContainer.getTypeAndLocForParameter(
+            Info.Location);
+    switch (Type) {
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+      dxbc::RTS0::v2::RootDescriptor Desc =
+          RootSigDesc->ParametersContainer.getRootDescriptor(Loc);
+
+      llvm::dxil::ResourceInfo::ResourceBinding Binding;
+      Binding.LowerBound = Desc.ShaderRegister;
+      Binding.Space = Desc.RegisterSpace;
+      Binding.Size = 1;
+      break;
+    }
+  }
 }
 } // namespace
 
 PreservedAnalyses
 DXILPostOptimizationValidation::run(Module &M, ModuleAnalysisManager &MAM) {
   DXILResourceMap &DRM = MAM.getResult<DXILResourceAnalysis>(M);
   DXILResourceBindingInfo &DRBI = MAM.getResult<DXILResourceBindingAnalysis>(M);
-  reportErrors(M, DRM, DRBI);
+  RootSignatureBindingInfo &RSBI = MAM.getResult<RootSignatureAnalysis>(M);
+  ModuleMetadataInfo &MMI = MAM.getResult<DXILMetadataAnalysis>(M);
+
+  reportErrors(M, DRM, DRBI, RSBI, MMI);
   return PreservedAnalyses::all();
 }
 
@@ -113,7 +145,13 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
         getAnalysis<DXILResourceWrapperPass>().getResourceMap();
     DXILResourceBindingInfo &DRBI =
         getAnalysis<DXILResourceBindingWrapperPass>().getBindingInfo();
-    reportErrors(M, DRM, DRBI);
+
+    RootSignatureBindingInfo &RSBI =
+        getAnalysis<RootSignatureAnalysisWrapper>().getRSInfo();
+    dxil::ModuleMetadataInfo &MMI =
+        getAnalysis<DXILMetadataAnalysisWrapperPass>().getModuleMetadata();
+
+    reportErrors(M, DRM, DRBI, RSBI, MMI);
     return false;
   }
   StringRef getPassName() const override {
@@ -125,10 +163,13 @@ class DXILPostOptimizationValidationLegacy : public ModulePass {
   void getAnalysisUsage(llvm::AnalysisUsage &AU) const override {
     AU.addRequired<DXILResourceWrapperPass>();
     AU.addRequired<DXILResourceBindingWrapperPass>();
+    AU.addRequired<RootSignatureAnalysisWrapper>();
+    AU.addRequired<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<DXILResourceWrapperPass>();
     AU.addPreserved<DXILResourceBindingWrapperPass>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
+    AU.addPreserved<RootSignatureAnalysisWrapper>();
   }
 };
 char DXILPostOptimizationValidationLegacy::ID = 0;

llvm/lib/Target/DirectX/DXILRootSignature.h

@@ -35,42 +35,42 @@ enum class RootSignatureElementKind {
 };
 
 class RootSignatureBindingInfo {
-  private:
-    SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> FuncToRsMap;
+private:
+  SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> FuncToRsMap;
 
-  public:
+public:
   using iterator =
-        SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>::iterator;
+      SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc>::iterator;
 
-  RootSignatureBindingInfo () = default;
-  RootSignatureBindingInfo(SmallDenseMap<const Function *, mcdxbc::RootSignatureDesc> Map) : FuncToRsMap(Map) {};
+  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) {
+  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;
+  RootSignatureAnalysis() = default;
 
   using Result = RootSignatureBindingInfo;
-  
-  RootSignatureBindingInfo
-  run(Module &M, ModuleAnalysisManager &AM);
+
+  RootSignatureBindingInfo run(Module &M, ModuleAnalysisManager &AM);
 };
 
 /// Wrapper pass for the legacy pass manager.
@@ -87,8 +87,8 @@ class RootSignatureAnalysisWrapper : public ModulePass {
 
   RootSignatureAnalysisWrapper() : ModulePass(ID) {}
 
-  RootSignatureBindingInfo& getRSInfo() {return *FuncToRsMap;}
-  
+  RootSignatureBindingInfo &getRSInfo() { return *FuncToRsMap; }
+
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;