[DirectX] generate resource table for PSV part (#106607)

python3kgae · web-flow · commit 981bb9dcc926 · 2024-09-12T20:41:20.000-04:00
Use DXILResourceWrapperPass to build the resource table. Since DXILResourceWrapperPass operates on LLVM intrinsics rather than DXIL operations, add addPreserved for DXILResourceWrapperPass in the passes before DXContainerGlobals Fixes #103275
diff --git a/llvm/lib/Target/DirectX/DXContainerGlobals.cpp b/llvm/lib/Target/DirectX/DXContainerGlobals.cpp
@@ -16,6 +16,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
+#include "llvm/Analysis/DXILResource.h"
 #include "llvm/BinaryFormat/DXContainer.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Constants.h"
@@ -40,6 +41,7 @@ class DXContainerGlobals : public llvm::ModulePass {
   GlobalVariable *buildSignature(Module &M, Signature &Sig, StringRef Name,
                                  StringRef SectionName);
   void addSignature(Module &M, SmallVector<GlobalValue *> &Globals);
+  void addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV);
   void addPipelineStateValidationInfo(Module &M,
                                       SmallVector<GlobalValue *> &Globals);
 
@@ -59,6 +61,7 @@ class DXContainerGlobals : public llvm::ModulePass {
     AU.setPreservesAll();
     AU.addRequired<ShaderFlagsAnalysisWrapper>();
     AU.addRequired<DXILMetadataAnalysisWrapperPass>();
+    AU.addRequired<DXILResourceWrapperPass>();
   }
 };
 
@@ -140,6 +143,56 @@ void DXContainerGlobals::addSignature(Module &M,
   Globals.emplace_back(buildSignature(M, OutputSig, "dx.osg1", "OSG1"));
 }
 
+void DXContainerGlobals::addResourcesForPSV(Module &M, PSVRuntimeInfo &PSV) {
+  const DXILResourceMap &ResMap =
+      getAnalysis<DXILResourceWrapperPass>().getResourceMap();
+
+  for (const dxil::ResourceInfo &ResInfo : ResMap) {
+    const dxil::ResourceInfo::ResourceBinding &Binding = ResInfo.getBinding();
+    dxbc::PSV::v2::ResourceBindInfo BindInfo;
+    BindInfo.LowerBound = Binding.LowerBound;
+    BindInfo.UpperBound = Binding.LowerBound + Binding.Size - 1;
+    BindInfo.Space = Binding.Space;
+
+    dxbc::PSV::ResourceType ResType = dxbc::PSV::ResourceType::Invalid;
+    bool IsUAV = ResInfo.getResourceClass() == dxil::ResourceClass::UAV;
+    switch (ResInfo.getResourceKind()) {
+    case dxil::ResourceKind::Sampler:
+      ResType = dxbc::PSV::ResourceType::Sampler;
+      break;
+    case dxil::ResourceKind::CBuffer:
+      ResType = dxbc::PSV::ResourceType::CBV;
+      break;
+    case dxil::ResourceKind::StructuredBuffer:
+      ResType = IsUAV ? dxbc::PSV::ResourceType::UAVStructured
+                      : dxbc::PSV::ResourceType::SRVStructured;
+      if (IsUAV && ResInfo.getUAV().HasCounter)
+        ResType = dxbc::PSV::ResourceType::UAVStructuredWithCounter;
+      break;
+    case dxil::ResourceKind::RTAccelerationStructure:
+      ResType = dxbc::PSV::ResourceType::SRVRaw;
+      break;
+    case dxil::ResourceKind::RawBuffer:
+      ResType = IsUAV ? dxbc::PSV::ResourceType::UAVRaw
+                      : dxbc::PSV::ResourceType::SRVRaw;
+      break;
+    default:
+      ResType = IsUAV ? dxbc::PSV::ResourceType::UAVTyped
+                      : dxbc::PSV::ResourceType::SRVTyped;
+      break;
+    }
+    BindInfo.Type = ResType;
+
+    BindInfo.Kind =
+        static_cast<dxbc::PSV::ResourceKind>(ResInfo.getResourceKind());
+    // TODO: Add support for dxbc::PSV::ResourceFlag::UsedByAtomic64, tracking
+    // with https://github.com/llvm/llvm-project/issues/104392
+    BindInfo.Flags.Flags = 0u;
+
+    PSV.Resources.emplace_back(BindInfo);
+  }
+}
+
 void DXContainerGlobals::addPipelineStateValidationInfo(
     Module &M, SmallVector<GlobalValue *> &Globals) {
   SmallString<256> Data;
@@ -155,6 +208,8 @@ void DXContainerGlobals::addPipelineStateValidationInfo(
   PSV.BaseData.ShaderStage =
       static_cast<uint8_t>(MMI.ShaderStage - Triple::Pixel);
 
+  addResourcesForPSV(M, PSV);
+
   // Hardcoded values here to unblock loading the shader into D3D.
   //
   // TODO: Lots more stuff to do here!
@@ -185,6 +240,7 @@ INITIALIZE_PASS_BEGIN(DXContainerGlobals, "dxil-globals",
                       "DXContainer Global Emitter", false, true)
 INITIALIZE_PASS_DEPENDENCY(ShaderFlagsAnalysisWrapper)
 INITIALIZE_PASS_DEPENDENCY(DXILMetadataAnalysisWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DXILResourceWrapperPass)
 INITIALIZE_PASS_END(DXContainerGlobals, "dxil-globals",
                     "DXContainer Global Emitter", false, true)
 
diff --git a/llvm/lib/Target/DirectX/DXILFinalizeLinkage.cpp b/llvm/lib/Target/DirectX/DXILFinalizeLinkage.cpp
@@ -8,6 +8,7 @@
 
 #include "DXILFinalizeLinkage.h"
 #include "DirectX.h"
+#include "llvm/Analysis/DXILResource.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/Metadata.h"
@@ -48,6 +49,10 @@ bool DXILFinalizeLinkageLegacy::runOnModule(Module &M) {
   return finalizeLinkage(M);
 }
 
+void DXILFinalizeLinkageLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addPreserved<DXILResourceWrapperPass>();
+}
+
 char DXILFinalizeLinkageLegacy::ID = 0;
 
 INITIALIZE_PASS_BEGIN(DXILFinalizeLinkageLegacy, DEBUG_TYPE,
diff --git a/llvm/lib/Target/DirectX/DXILFinalizeLinkage.h b/llvm/lib/Target/DirectX/DXILFinalizeLinkage.h
@@ -32,6 +32,7 @@ class DXILFinalizeLinkageLegacy : public ModulePass {
   DXILFinalizeLinkageLegacy() : ModulePass(ID) {}
   bool runOnModule(Module &M) override;
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
   static char ID; // Pass identification.
 };
 } // namespace llvm
diff --git a/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp b/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.cpp
@@ -14,6 +14,7 @@
 #include "DirectX.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/DXILResource.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
@@ -495,6 +496,10 @@ bool DXILIntrinsicExpansionLegacy::runOnModule(Module &M) {
   return expansionIntrinsics(M);
 }
 
+void DXILIntrinsicExpansionLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
+  AU.addPreserved<DXILResourceWrapperPass>();
+}
+
 char DXILIntrinsicExpansionLegacy::ID = 0;
 
 INITIALIZE_PASS_BEGIN(DXILIntrinsicExpansionLegacy, DEBUG_TYPE,
diff --git a/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.h b/llvm/lib/Target/DirectX/DXILIntrinsicExpansion.h
@@ -26,6 +26,7 @@ class DXILIntrinsicExpansionLegacy : public ModulePass {
   bool runOnModule(Module &M) override;
   DXILIntrinsicExpansionLegacy() : ModulePass(ID) {}
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
   static char ID; // Pass identification.
 };
 } // namespace llvm
diff --git a/llvm/lib/Target/DirectX/DXILPrepare.cpp b/llvm/lib/Target/DirectX/DXILPrepare.cpp
@@ -20,6 +20,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Analysis/DXILMetadataAnalysis.h"
+#include "llvm/Analysis/DXILResource.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/AttributeMask.h"
 #include "llvm/IR/IRBuilder.h"
@@ -249,6 +250,7 @@ class DXILPrepareModule : public ModulePass {
     AU.addPreserved<ShaderFlagsAnalysisWrapper>();
     AU.addPreserved<DXILResourceMDWrapper>();
     AU.addPreserved<DXILMetadataAnalysisWrapperPass>();
+    AU.addPreserved<DXILResourceWrapperPass>();
   }
   static char ID; // Pass identification.
 };
diff --git a/llvm/test/CodeGen/DirectX/ContainerData/PSVResources.ll b/llvm/test/CodeGen/DirectX/ContainerData/PSVResources.ll
@@ -0,0 +1,93 @@
+; RUN: llc %s --filetype=obj -o - | obj2yaml | FileCheck %s
+
+; Make sure resource table is created correctly.
+; CHECK: Resources:
+target triple = "dxil-unknown-shadermodel6.0-compute"
+
+define void @main() #0 {
+
+  ; ByteAddressBuffer Buf : register(t8, space1)
+; CHECK:        - Type:            SRVRaw
+; CHECK:          Space:           1
+; CHECK:          LowerBound:      8
+; CHECK:          UpperBound:      8
+; CHECK:          Kind:            RawBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  %srv0 = call target("dx.RawBuffer", i8, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.RawBuffer_i8_0_0t(
+          i32 1, i32 8, i32 1, i32 0, i1 false)
+
+  ; struct S { float4 a; uint4 b; };
+  ; StructuredBuffer<S> Buf : register(t2, space4)
+; CHECK:        - Type:            SRVStructured
+; CHECK:          Space:           4
+; CHECK:          LowerBound:      2
+; CHECK:          UpperBound:      2
+; CHECK:          Kind:            StructuredBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  %srv1 = call target("dx.RawBuffer", {<4 x float>, <4 x i32>}, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.RawBuffer_sl_v4f32v4i32s_0_0t(
+          i32 4, i32 2, i32 1, i32 0, i1 false)
+
+  ; Buffer<uint4> Buf[24] : register(t3, space5)
+; CHECK:        - Type:            SRVTyped
+; CHECK:          Space:           5
+; CHECK:          LowerBound:      3
+; CHECK:          UpperBound:      26
+; CHECK:          Kind:            TypedBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  %srv2 = call target("dx.TypedBuffer", <4 x i32>, 0, 0, 0)
+      @llvm.dx.handle.fromBinding.tdx.TypedBuffer_i32_0_0t(
+          i32 5, i32 3, i32 24, i32 0, i1 false)
+
+  ; RWBuffer<int> Buf : register(u7, space2)
+; CHECK:        - Type:            UAVTyped
+; CHECK:          Space:           2
+; CHECK:          LowerBound:      7
+; CHECK:          UpperBound:      7
+; CHECK:          Kind:            TypedBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  %uav0 = call target("dx.TypedBuffer", i32, 1, 0, 1)
+      @llvm.dx.handle.fromBinding.tdx.TypedBuffer_i32_1_0t(
+          i32 2, i32 7, i32 1, i32 0, i1 false)
+
+  ; RWBuffer<float4> Buf : register(u5, space3)
+; CHECK:        - Type:            UAVTyped
+; CHECK:          Space:           3
+; CHECK:          LowerBound:      5
+; CHECK:          UpperBound:      5
+; CHECK:          Kind:            TypedBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  %uav1 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
+              @llvm.dx.handle.fromBinding.tdx.TypedBuffer_f32_1_0(
+                  i32 3, i32 5, i32 1, i32 0, i1 false)
+
+  ; RWBuffer<float4> BufferArray[10] : register(u0, space4)
+; CHECK:        - Type:            UAVTyped
+; CHECK:          Space:           4
+; CHECK:          LowerBound:      0
+; CHECK:          UpperBound:      9
+; CHECK:          Kind:            TypedBuffer
+; CHECK:          Flags:
+; CHECK:            UsedByAtomic64:  false
+  ; RWBuffer<float4> Buf = BufferArray[0]
+  %uav2_1 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
+              @llvm.dx.handle.fromBinding.tdx.TypedBuffer_f32_1_0(
+                  i32 4, i32 0, i32 10, i32 0, i1 false)
+  ; RWBuffer<float4> Buf = BufferArray[5]
+  %uav2_2 = call target("dx.TypedBuffer", <4 x float>, 1, 0, 0)
+              @llvm.dx.handle.fromBinding.tdx.TypedBuffer_f32_1_0(
+                  i32 4, i32 0, i32 10, i32 5, i1 false)
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+
+!dx.valver = !{!0}
+
+!0 = !{i32 1, i32 7}
diff --git a/llvm/test/CodeGen/DirectX/llc-pipeline.ll b/llvm/test/CodeGen/DirectX/llc-pipeline.ll
@@ -14,12 +14,11 @@
 ; CHECK-NEXT:   DXIL Resource analysis
 ; CHECK-NEXT:   DXIL Op Lowering
 ; CHECK-NEXT:   DXIL Finalize Linkage
-; CHECK-NEXT:   DXIL Resource analysis
 ; CHECK-NEXT:   DXIL resource Information
 ; CHECK-NEXT:   DXIL Shader Flag Analysis
+; CHECK-NEXT:   DXIL Module Metadata analysis
 ; CHECK-NEXT:   DXIL Translate Metadata
 ; CHECK-NEXT:   DXIL Prepare Module
-; CHECK-NEXT:   DXIL Resource analysis
 ; CHECK-NEXT:   DXIL Metadata Pretty Printer
 ; CHECK-NEXT:   Print Module IR
  
