Added shader reflection part to source-only PDB. (#3467)

include/dxc/DxilContainer/DxilContainerAssembler.h

@@ -15,11 +15,18 @@
 #include "dxc/DxilContainer/DxilContainer.h"
 #include "llvm/ADT/StringRef.h"
 
+struct IStream;
+
+namespace llvm {
+class Module;
+}
+
 namespace hlsl {
 
 class AbstractMemoryStream;
 class DxilModule;
 class RootSignatureHandle;
+class ShaderModel;
 namespace DXIL {
 enum class SignatureKind;
 }
@@ -46,6 +53,16 @@ DxilPartWriter *NewRDATWriter(const DxilModule &M);
 
 DxilContainerWriter *NewDxilContainerWriter();
 
+// Set validator version to 0,0 (not validated) then re-emit as much reflection metadata as possible.
+void ReEmitLatestReflectionData(llvm::Module *pReflectionM);
+
+// Strip functions and serialize module.
+void StripAndCreateReflectionStream(llvm::Module *pReflectionM, uint32_t *pReflectionPartSizeInBytes, AbstractMemoryStream **ppReflectionStreamOut);
+
+void WriteProgramPart(const hlsl::ShaderModel *pModel,
+                      AbstractMemoryStream *pModuleBitcode,
+                      IStream *pStream);
+
 void SerializeDxilContainerForModule(hlsl::DxilModule *pModule,
                                      AbstractMemoryStream *pModuleBitcode,
                                      AbstractMemoryStream *pStream,

include/dxc/Support/FileIOHelper.h

@@ -236,7 +236,7 @@ HRESULT CreateReadOnlyBlobStream(_In_ IDxcBlob *pSource, _COM_Outptr_ IStream**
 HRESULT CreateFixedSizeMemoryStream(_In_ LPBYTE pBuffer, size_t size, _COM_Outptr_ AbstractMemoryStream** ppResult) throw();
 
 template <typename T>
-HRESULT WriteStreamValue(AbstractMemoryStream *pStream, const T& value) {
+HRESULT WriteStreamValue(IStream *pStream, const T& value) {
   ULONG cb;
   return pStream->Write(&value, sizeof(value), &cb);
 }

lib/DxilContainer/DxilContainerAssembler.cpp

@@ -1536,9 +1536,9 @@ static void GetPaddedProgramPartSize(AbstractMemoryStream *pStream,
   bitcodeInUInt32 = (bitcodeInUInt32 / 4) + (bitcodePaddingBytes ? 1 : 0);
 }
 
-static void WriteProgramPart(const ShaderModel *pModel,
+void hlsl::WriteProgramPart(const ShaderModel *pModel,
                              AbstractMemoryStream *pModuleBitcode,
-                             AbstractMemoryStream *pStream) {
+                             IStream *pStream) {
   DXASSERT(pModel != nullptr, "else generation should have failed");
   DxilProgramHeader programHeader;
   uint32_t shaderVersion =
@@ -1579,6 +1579,60 @@ public:
 
 } // namespace
 
+
+void hlsl::ReEmitLatestReflectionData(llvm::Module *pM) {
+  // Retain usage information in metadata for reflection by:
+  // Upgrade validator version, re-emit metadata
+  // 0,0 = Not meant to be validated, support latest
+
+  DxilModule &DM = pM->GetOrCreateDxilModule();
+
+  DM.SetValidatorVersion(0, 0);
+  DM.ReEmitDxilResources();
+  DM.EmitDxilCounters();
+}
+
+static std::unique_ptr<Module> CloneModuleForReflection(Module *pM) {
+  DxilModule &DM = pM->GetOrCreateDxilModule();
+
+  unsigned ValMajor = 0, ValMinor = 0;
+  DM.GetValidatorVersion(ValMajor, ValMinor);
+
+  // Emit the latest reflection metadata
+  hlsl::ReEmitLatestReflectionData(pM);
+
+  // Clone module
+  std::unique_ptr<Module> reflectionModule( llvm::CloneModule(pM) );
+
+  // Now restore validator version on main module and re-emit metadata.
+  DM.SetValidatorVersion(ValMajor, ValMinor);
+  DM.ReEmitDxilResources();
+
+  return reflectionModule;
+}
+
+void hlsl::StripAndCreateReflectionStream(Module *pReflectionM, uint32_t *pReflectionPartSizeInBytes, AbstractMemoryStream **ppReflectionStreamOut) {
+  for (Function &F : pReflectionM->functions()) {
+    if (!F.isDeclaration()) {
+      F.deleteBody();
+    }
+  }
+
+  uint32_t reflectPartSizeInBytes = 0;
+  CComPtr<AbstractMemoryStream> pReflectionBitcodeStream;
+
+  IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pReflectionBitcodeStream));
+  raw_stream_ostream outStream(pReflectionBitcodeStream.p);
+  WriteBitcodeToFile(pReflectionM, outStream, false);
+  outStream.flush();
+  uint32_t reflectInUInt32 = 0, reflectPaddingBytes = 0;
+  GetPaddedProgramPartSize(pReflectionBitcodeStream, reflectInUInt32, reflectPaddingBytes);
+  reflectPartSizeInBytes = reflectInUInt32 * sizeof(uint32_t) + sizeof(DxilProgramHeader);
+
+  *pReflectionPartSizeInBytes = reflectPartSizeInBytes;
+  *ppReflectionStreamOut = pReflectionBitcodeStream.Detach();
+}
+
 void hlsl::SerializeDxilContainerForModule(DxilModule *pModule,
                                            AbstractMemoryStream *pModuleBitcode,
                                            AbstractMemoryStream *pFinalStream,
@@ -1714,7 +1768,7 @@ void hlsl::SerializeDxilContainerForModule(DxilModule *pModule,
     GetPaddedProgramPartSize(pInputProgramStream, debugInUInt32, debugPaddingBytes);
     if (Flags & SerializeDxilFlags::IncludeDebugInfoPart) {
       writer.AddPart(DFCC_ShaderDebugInfoDXIL, debugInUInt32 * sizeof(uint32_t) + sizeof(DxilProgramHeader), [&](AbstractMemoryStream *pStream) {
-        WriteProgramPart(pModule->GetShaderModel(), pInputProgramStream, pStream);
+        hlsl::WriteProgramPart(pModule->GetShaderModel(), pInputProgramStream, pStream);
       });
     }
 
@@ -1727,43 +1781,13 @@ void hlsl::SerializeDxilContainerForModule(DxilModule *pModule,
     Flags &= ~SerializeDxilFlags::DebugNameDependOnSource;
   }
 
-  // Clone module for reflection, strip function defs
-  std::unique_ptr<Module> reflectionModule;
-  if (bEmitReflection) {
-    // Retain usage information in metadata for reflection by:
-    // Upgrade validator version, re-emit metadata, then clone module for reflection.
-    // 0,0 = Not meant to be validated, support latest
-    pModule->SetValidatorVersion(0, 0);
-    pModule->ReEmitDxilResources();
-    pModule->EmitDxilCounters();
-
-    reflectionModule.reset(llvm::CloneModule(pModule->GetModule()));
-
-    // Now restore validator version on main module and re-emit metadata.
-    pModule->SetValidatorVersion(ValMajor, ValMinor);
-    pModule->ReEmitDxilResources();
-
-    for (Function &F : reflectionModule->functions()) {
-      if (!F.isDeclaration()) {
-        F.deleteBody();
-      }
-    }
-    // Just make sure this doesn't crash/assert on debug build:
-    DXASSERT_NOMSG(&reflectionModule->GetOrCreateDxilModule());
-  }
-
+  uint32_t reflectPartSizeInBytes = 0;
   CComPtr<AbstractMemoryStream> pReflectionBitcodeStream;
 
-  uint32_t reflectPartSizeInBytes = 0;
-  if (bEmitReflection)
-  {
-    IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pReflectionBitcodeStream));
-    raw_stream_ostream outStream(pReflectionBitcodeStream.p);
-    WriteBitcodeToFile(reflectionModule.get(), outStream, false);
-    outStream.flush();
-    uint32_t reflectInUInt32 = 0, reflectPaddingBytes = 0;
-    GetPaddedProgramPartSize(pReflectionBitcodeStream, reflectInUInt32, reflectPaddingBytes);
-    reflectPartSizeInBytes = reflectInUInt32 * sizeof(uint32_t) + sizeof(DxilProgramHeader);
+  if (bEmitReflection) {
+    // Clone module for reflection
+    std::unique_ptr<Module> reflectionModule = CloneModuleForReflection(pModule->GetModule());
+    hlsl::StripAndCreateReflectionStream(reflectionModule.get(), &reflectPartSizeInBytes, &pReflectionBitcodeStream);
   }
 
   if (pReflectionStreamOut) {

tools/clang/tools/dxcompiler/dxclibrary.cpp

@@ -351,11 +351,23 @@ public:
 
     DxcThreadMalloc TM(m_pMalloc);
     try {
+      CComPtr<IDxcBlob> pPdbContainerBlob;
       const DxilPartHeader *pModulePart = nullptr;
       const DxilPartHeader *pRDATPart = nullptr;
 
+      const DxilContainerHeader *pHeader = IsDxilContainerLike(pData->Ptr, pData->Size);
+      if (!pHeader) {
+        CComPtr<IDxcBlobEncoding> pBlob;
+        IFR(hlsl::DxcCreateBlobWithEncodingFromPinned(pData->Ptr, pData->Size, pData->Size, &pBlob));
+        CComPtr<IStream> pStream;
+        IFR(hlsl::CreateReadOnlyBlobStream(pBlob, &pStream));
+        if (SUCCEEDED(hlsl::pdb::LoadDataFromStream(m_pMalloc, pStream, &pPdbContainerBlob))) {
+          pHeader = IsDxilContainerLike(pPdbContainerBlob->GetBufferPointer(), pPdbContainerBlob->GetBufferSize());
+        }
+      }
+
       // Is this a valid DxilContainer?
-      if (const DxilContainerHeader *pHeader = IsDxilContainerLike(pData->Ptr, pData->Size)) {
+      if (pHeader) {
         if (!IsValidDxilContainer(pHeader, pData->Size))
           return E_INVALIDARG;
 
@@ -389,7 +401,6 @@ public:
         pModulePart = pStatsPart ? pStatsPart : pDebugDXILPart ? pDebugDXILPart : pDXILPart;
         if (nullptr == pModulePart)
           return DXC_E_MISSING_PART;
-
       } else {
         // Not a container, try a statistics part that holds a valid program part.
         // In the future, this will just be the RDAT part.

tools/clang/tools/dxcompiler/dxcompilerobj.cpp

@@ -160,7 +160,14 @@ struct CompilerVersionPartWriter {
   }
 };
 
-static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer, IDxcBlob *pDebugBlob, IDxcVersionInfo *pVersionInfo, const hlsl::DxilSourceInfo *pSourceInfo, IDxcBlob **ppNewContaner) {
+static HRESULT CreateContainerForPDB(IMalloc *pMalloc,
+  llvm::Module *pModule,
+  IDxcBlob *pOldContainer,
+  IDxcBlob *pDebugBlob, IDxcVersionInfo *pVersionInfo,
+  const hlsl::DxilSourceInfo *pSourceInfo,
+  AbstractMemoryStream *pReflectionStream, const uint32_t reflectionSizeInBytes,
+  IDxcBlob **ppNewContaner)
+{
   // If the pContainer is not a valid container, give up.
   if (!hlsl::IsValidDxilContainer((hlsl::DxilContainerHeader *)pOldContainer->GetBufferPointer(), pOldContainer->GetBufferSize()))
     return E_FAIL;
@@ -185,12 +192,16 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
   SmallVector<UINT32, 4> OffsetTable;
   SmallVector<Part, 4> PartWriters;
   UINT32 uTotalPartsSize = 0;
+
+  auto AddPart = [&PartWriters, &OffsetTable, &uTotalPartsSize](Part NewPart, UINT32 uSize) {
+    OffsetTable.push_back(uTotalPartsSize);
+    uTotalPartsSize += uSize + sizeof(hlsl::DxilPartHeader);
+    PartWriters.push_back(NewPart);
+  };
+
   for (unsigned i = 0; i < DxilHeader->PartCount; i++) {
     hlsl::DxilPartHeader *PartHeader = GetDxilContainerPart(DxilHeader, i);
     if (ShouldBeCopiedIntoPDB(PartHeader->PartFourCC)) {
-      OffsetTable.push_back(uTotalPartsSize);
-      uTotalPartsSize += PartHeader->PartSize + sizeof(*PartHeader);
-
       UINT32 uSize = PartHeader->PartSize;
       const void *pPartData = PartHeader+1;
       Part NewPart(
@@ -202,7 +213,7 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
           return S_OK;
         }
       );
-      PartWriters.push_back(NewPart);
+      AddPart(NewPart, uSize);
     }
 
     // Could use any of these. We're mostly after the header version and all that.
@@ -219,9 +230,6 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
   if (pSourceInfo) {
     const UINT32 uPartSize = pSourceInfo->AlignedSizeInBytes;
 
-    OffsetTable.push_back(uTotalPartsSize);
-    uTotalPartsSize += uPartSize + sizeof(hlsl::DxilPartHeader);
-
     Part NewPart(
       hlsl::DFCC_ShaderSourceInfo,
       uPartSize,
@@ -231,16 +239,26 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
         return S_OK;
       }
     );
-    PartWriters.push_back(NewPart);
+
+    AddPart(NewPart, uPartSize);
+  }
+
+  if (pReflectionStream) {
+    Part NewPart(
+      hlsl::DFCC_ShaderStatistics,
+      reflectionSizeInBytes,
+      [&pReflectionStream, pModule](IStream *pStream) {
+        hlsl::WriteProgramPart(pModule->GetOrCreateDxilModule().GetShaderModel(), pReflectionStream, pStream);
+        return S_OK;
+      }
+    );
+    AddPart(NewPart, reflectionSizeInBytes);
   }
 
   CompilerVersionPartWriter versionWriter;
   if (pVersionInfo) {
     versionWriter.Init(pVersionInfo);
 
-    OffsetTable.push_back(uTotalPartsSize);
-    uTotalPartsSize += versionWriter.GetSize() + sizeof(hlsl::DxilPartHeader);
-
     Part NewPart(
       hlsl::DFCC_CompilerVersion,
       versionWriter.GetSize(),
@@ -249,7 +267,7 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
         return S_OK;
       }
     );
-    PartWriters.push_back(NewPart);
+    AddPart(NewPart, versionWriter.GetSize());
   }
 
   if (pDebugBlob) {
@@ -262,9 +280,6 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
 
     UINT32 uPaddingSize = 0;
     UINT32 uPartSize = AlignByDword(sizeof(hlsl::DxilProgramHeader) + pDebugBlob->GetBufferSize(), &uPaddingSize);
-
-    OffsetTable.push_back(uTotalPartsSize);
-    uTotalPartsSize += uPartSize + sizeof(hlsl::DxilPartHeader);
     
     Part NewPart(
       hlsl::DFCC_ShaderDebugInfoDXIL,
@@ -286,7 +301,7 @@ static HRESULT CreateContainerForPDB(IMalloc *pMalloc, IDxcBlob *pOldContainer,
         return S_OK;
       }
     );
-    PartWriters.push_back(NewPart);
+    AddPart(NewPart, uPartSize);
   }
 
   // Offset the offset table by the offset table itself
@@ -1036,10 +1051,15 @@ public:
             pSourceInfo = debugSourceInfoWriter.GetPart();
           }
 
+          CComPtr<AbstractMemoryStream> pReflectionStream;
+          uint32_t reflectionSizeInBytes = 0;
+
           CComPtr<IDxcBlob> pDebugProgramBlob;
           // Don't include the debug part if using source only PDB
           if (opts.SourceOnlyDebug) {
             assert(pSourceInfo);
+            hlsl::ReEmitLatestReflectionData(compiledModule.get());
+            hlsl::StripAndCreateReflectionStream(compiledModule.get(), &reflectionSizeInBytes, &pReflectionStream);
           }
           else {
             if (!opts.SourceInDebugModule) {
@@ -1056,8 +1076,11 @@ public:
           }
 
           IFT(CreateContainerForPDB(
-            m_pMalloc, pOutputBlob, pDebugProgramBlob,
+            m_pMalloc,
+            compiledModule.get(),
+            pOutputBlob, pDebugProgramBlob,
             static_cast<IDxcVersionInfo *>(this), pSourceInfo,
+            pReflectionStream, reflectionSizeInBytes,
             &pStrippedContainer));
         }
 

tools/clang/unittests/HLSL/CompilerTest.cpp

@@ -28,6 +28,7 @@
 #include "dxc/dxcpix.h"
 #ifdef _WIN32
 #include <atlfile.h>
+#include <d3dcompiler.h>
 #include "dia2.h"
 #endif
 
@@ -988,6 +989,7 @@ static void VerifyPdbUtil(dxc::DxcDllSupport &dllSupport,
     bool HasVersion,
     bool IsFullPDB,
     bool HasHashAndPdbName,
+    bool TestReflection,
     const std::string &MainSource,
     const std::string &IncludedFile)
 {
@@ -1213,6 +1215,40 @@ static void VerifyPdbUtil(dxc::DxcDllSupport &dllSupport,
     TestArgumentPair(Args, ExpectedArgs);
   }
 
+  // Shader reflection
+  if (TestReflection) {
+    CComPtr<IDxcUtils> pUtils;
+    VERIFY_SUCCEEDED(dllSupport.CreateInstance(CLSID_DxcUtils, &pUtils));
+
+    DxcBuffer buf = {};
+    buf.Ptr = pBlob->GetBufferPointer();
+    buf.Size = pBlob->GetBufferSize();
+    buf.Encoding = CP_ACP;
+
+    CComPtr<ID3D12ShaderReflection> pRefl;
+    VERIFY_SUCCEEDED(pUtils->CreateReflection(&buf, IID_PPV_ARGS(&pRefl)));
+
+    D3D12_SHADER_DESC desc = {};
+    VERIFY_SUCCEEDED(pRefl->GetDesc(&desc));
+
+    VERIFY_ARE_EQUAL(desc.ConstantBuffers, 1);
+    ID3D12ShaderReflectionConstantBuffer *pCB = pRefl->GetConstantBufferByIndex(0);
+
+    D3D12_SHADER_BUFFER_DESC cbDesc = {};
+    VERIFY_SUCCEEDED(pCB->GetDesc(&cbDesc));
+
+    VERIFY_IS_TRUE(0 == strcmp(cbDesc.Name, "MyCbuffer"));
+    VERIFY_ARE_EQUAL(cbDesc.Variables, 1);
+
+    ID3D12ShaderReflectionVariable *pVar = pCB->GetVariableByIndex(0);
+    D3D12_SHADER_VARIABLE_DESC varDesc = {};
+    VERIFY_SUCCEEDED(pVar->GetDesc(&varDesc));
+
+    VERIFY_ARE_EQUAL(varDesc.uFlags, D3D_SVF_USED);
+    VERIFY_IS_TRUE(0 == strcmp(varDesc.Name, "my_cbuf_foo"));
+    VERIFY_ARE_EQUAL(varDesc.Size, sizeof(float) * 4);
+  }
+
   // Make the pix debug info
   if (IsFullPDB) {
     VERIFY_IS_TRUE(pPdbUtils->IsFullPDB());
@@ -1250,6 +1286,7 @@ static void VerifyPdbUtil(dxc::DxcDllSupport &dllSupport,
       pMainFileName,
       NewExpectedArgs, NewExpectedFlags, ExpectedDefines,
       pCompiler, HasVersion, /*IsFullPDB*/true,
+      /*TestReflection*/true,
       HasHashAndPdbName, MainSource, IncludedFile);
   }
 
@@ -1362,8 +1399,17 @@ void CompilerTest::TestPdbUtils(bool bSlim, bool bSourceInDebugModule, bool bStr
   CComPtr<IDxcBlobEncoding> pSource;
   CComPtr<IDxcOperationResult> pOperationResult;
 
-  std::string main_source = "#include \"helper.h\"\r\n"
-    "float4 PSMain() : SV_Target { return ZERO; }";
+  std::string main_source = R"x(
+      #include "helper.h"
+      cbuffer MyCbuffer : register(b1) {
+        float4 my_cbuf_foo;
+      }
+
+      [RootSignature("CBV(b1)")]
+      float4 PSMain() : SV_Target {
+        return ZERO + my_cbuf_foo;
+      }
+  )x";
   std::string included_File = "#define ZERO 0";
 
   VERIFY_SUCCEEDED(CreateCompiler(&pCompiler));
@@ -1449,7 +1495,8 @@ void CompilerTest::TestPdbUtils(bool bSlim, bool bSourceInDebugModule, bool bStr
       pCompiler,
       /*HasVersion*/ false,
       /*IsFullPDB*/  true,
-      /*hasHasAndPDBName*/false,
+      /*HasHashAndPdbName*/false,
+      /*TestReflection*/false, // Reflection creation interface doesn't support just the DxilProgramHeader.
       main_source, included_File);
   }
 
@@ -1460,7 +1507,8 @@ void CompilerTest::TestPdbUtils(bool bSlim, bool bSourceInDebugModule, bool bStr
     pCompiler,
     /*HasVersion*/ true,
     /*IsFullPDB*/ !bSlim,
-    /*hasHasAndPDBName*/true,
+    /*HasHashAndPdbName*/true,
+    /*TestReflection*/true,
     main_source, included_File);
 
   if (!bStrip) {
@@ -1471,7 +1519,8 @@ void CompilerTest::TestPdbUtils(bool bSlim, bool bSourceInDebugModule, bool bStr
       pCompiler,
       /*HasVersion*/ false,
       /*IsFullPDB*/ true,
-      /*hasHasAndPDBName*/true,
+      /*HasHashAndPdbName*/true,
+      /*TestReflection*/true,
       main_source, included_File);
   }
 }