PIX: Modify root sigs in place (plus fix root sig memory leak) (microsoft#4876)

PIX is unique in that it needs to deserialize, modify and then reserialize root sigs. The focus of this checkin is adding such modifications to PixPassHelpers.cpp. (This closes a gap in PIX support: PIX can now support shaders that use dxil-defined and attribute-style root signatures.)

But this required some work outside of the purely PIX-focused areas. Deserialized root sigs are described by a C-like structure with embedded arrays that are new/delete-ed by routines in DxilRootSignature.cpp. Since modifying these structs requires more new/delete calls, I chose to add a new entry point in DxilRootSignature.cpp to do the bare minimum that PIX needs: extend root params by one descriptor. This approach keeps all those raw new/deletes in one file.

I found a leak in DxilRootSignatureSerialzier.cpp, which I fixed. There appear to be no unit tests that exercise this path. I welcome feedback on adding one.

There were other leaks in class CVersionedRootSignatureDeserializer, but this class is unused so I deleted it.

Oh, and there are bazillions of commits because I was cherry-picking from a recent change (microsoft#4845) as it eveolved, since I needed that change and this to test PIX.

Author: jeffnn

include/dxc/DxilRootSignature/DxilRootSignature.h

@@ -385,6 +385,43 @@ bool VerifyRootSignature(_In_ const DxilVersionedRootSignatureDesc *pDesc,
                          _In_ llvm::raw_ostream &DiagStream,
                          _In_ bool bAllowReservedRegisterSpace);
 
+class DxilVersionedRootSignature {
+  DxilVersionedRootSignatureDesc *m_pRootSignature;
+
+public:
+    // Non-copyable:
+  DxilVersionedRootSignature(DxilVersionedRootSignature const &) = delete;
+  DxilVersionedRootSignature const &
+  operator=(DxilVersionedRootSignature const &) = delete;
+
+  // but movable:
+  DxilVersionedRootSignature(DxilVersionedRootSignature &&) = default;
+  DxilVersionedRootSignature &
+  operator=(DxilVersionedRootSignature &&) = default;
+
+  DxilVersionedRootSignature() : m_pRootSignature(nullptr) {}
+  explicit DxilVersionedRootSignature(
+      const DxilVersionedRootSignatureDesc *pRootSignature)
+      : m_pRootSignature(
+            const_cast<DxilVersionedRootSignatureDesc *> (pRootSignature)) {}
+  ~DxilVersionedRootSignature() { 
+      DeleteRootSignature(m_pRootSignature);
+  }
+  const DxilVersionedRootSignatureDesc* operator -> () const {
+    return m_pRootSignature;
+  }
+  const DxilVersionedRootSignatureDesc ** get_address_of() {
+    if (m_pRootSignature != nullptr)
+      return nullptr; // You're probably about to leak...
+    return const_cast<const DxilVersionedRootSignatureDesc **> (&m_pRootSignature);
+  }
+  const DxilVersionedRootSignatureDesc* get() const { 
+      return m_pRootSignature;
+  }
+  DxilVersionedRootSignatureDesc* get_mutable() const { 
+      return m_pRootSignature;
+  }
+};
 } // namespace hlsl
 
 #endif // __DXC_ROOTSIGNATURE__

lib/DxilPIXPasses/PixPassHelpers.cpp

@@ -9,10 +9,12 @@
 
 #include "dxc/DXIL/DxilOperations.h"
 #include "dxc/DXIL/DxilInstructions.h"
+#include "dxc/DXIL/DxilFunctionProps.h"
 #include "dxc/DXIL/DxilModule.h"
 #include "dxc/DXIL/DxilResourceBinding.h"
 #include "dxc/DXIL/DxilResourceProperties.h"
 #include "dxc/HLSL/DxilSpanAllocator.h"
+#include "dxc/DxilRootSignature/DxilRootSignature.h"
 
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
@@ -21,6 +23,10 @@
 
 #include "PixPassHelpers.h"
 
+#include "dxc/Support/Global.h"
+#include "dxc/Support/WinIncludes.h"
+#include "dxc/dxcapi.h"
+
 #ifdef PIX_DEBUG_DUMP_HELPER
 #include <iostream>
 #include "llvm/IR/DebugInfo.h"
@@ -159,7 +165,93 @@ llvm::CallInst *CreateHandleForResource(hlsl::DxilModule &DM,
   }
 }
 
-// Set up a UAV with structure of a single int
+static std::vector<uint8_t> SerializeRootSignatureToVector(DxilVersionedRootSignatureDesc const *rootSignature) {
+  CComPtr<IDxcBlob> serializedRootSignature;
+  CComPtr<IDxcBlobEncoding> errorBlob;
+  constexpr bool allowReservedRegisterSpace = true;
+  SerializeRootSignature(rootSignature, &serializedRootSignature, &errorBlob,
+                         allowReservedRegisterSpace);
+  std::vector<uint8_t> ret;
+  auto const *serializedData = reinterpret_cast<const uint8_t *>(
+      serializedRootSignature->GetBufferPointer());
+  ret.assign(serializedData,
+             serializedData + serializedRootSignature->GetBufferSize());
+
+  return ret;
+}
+
+constexpr uint32_t toolsRegisterSpace = static_cast<uint32_t>(-2);
+constexpr uint32_t toolsUAVRegister = 0;
+
+template<typename RootSigDesc, typename RootParameterDesc>
+void ExtendRootSig(RootSigDesc &rootSigDesc) {
+  auto *existingParams = rootSigDesc.pParameters;
+  auto *newParams = new RootParameterDesc[rootSigDesc.NumParameters + 1];
+  if (existingParams != nullptr) {
+    memcpy(newParams, existingParams,
+           rootSigDesc.NumParameters * sizeof(RootParameterDesc));
+    delete[] existingParams;
+  }
+  rootSigDesc.pParameters = newParams;
+  rootSigDesc.pParameters[rootSigDesc.NumParameters].ParameterType = DxilRootParameterType::UAV;
+  rootSigDesc.pParameters[rootSigDesc.NumParameters].Descriptor.RegisterSpace = toolsRegisterSpace;
+  rootSigDesc.pParameters[rootSigDesc.NumParameters].Descriptor.ShaderRegister = toolsUAVRegister;
+  rootSigDesc.pParameters[rootSigDesc.NumParameters].ShaderVisibility = DxilShaderVisibility::All;
+  rootSigDesc.NumParameters++;
+} 
+
+static std::vector<uint8_t> AddUAVParamterToRootSignature(const void *Data,
+                                                       uint32_t Size) {
+  DxilVersionedRootSignature rootSignature;
+  DeserializeRootSignature(Data, Size, rootSignature.get_address_of());
+  auto *rs = rootSignature.get_mutable();
+  switch (rootSignature->Version) {
+  case DxilRootSignatureVersion::Version_1_0:
+    ExtendRootSig<DxilRootSignatureDesc, DxilRootParameter>(rs->Desc_1_0);
+    break;
+  case DxilRootSignatureVersion::Version_1_1:
+    ExtendRootSig<DxilRootSignatureDesc1, DxilRootParameter1>(rs->Desc_1_1);
+    rs->Desc_1_1.pParameters[rs->Desc_1_1.NumParameters - 1].Descriptor.Flags =
+        hlsl::DxilRootDescriptorFlags::None;
+    break;
+  }
+  return SerializeRootSignatureToVector(rs);
+}
+
+static void AddUAVToShaderAttributeRootSignature(DxilModule &DM) {
+  auto rs = DM.GetSerializedRootSignature();
+  if(!rs.empty()) {
+    std::vector<uint8_t> asVector = AddUAVParamterToRootSignature(rs.data(), static_cast<uint32_t>(rs.size()));
+    DM.ResetSerializedRootSignature(asVector);
+  }
+}
+
+static void AddUAVToDxilDefinedGlobalRootSignatures(DxilModule& DM) {
+  auto *subObjects = DM.GetSubobjects();
+  if (subObjects != nullptr) {
+    for (auto const &subObject : subObjects->GetSubobjects()) {
+      if (subObject.second->GetKind() ==
+          DXIL::SubobjectKind::GlobalRootSignature) {
+        const void *Data = nullptr;
+        uint32_t Size = 0;
+        constexpr bool notALocalRS = false;
+        if (subObject.second->GetRootSignature(notALocalRS, Data, Size,
+                                               nullptr)) {
+          auto extendedRootSig = AddUAVParamterToRootSignature(Data, Size);
+          auto rootSignatureSubObjectName = subObject.first;
+          subObjects->RemoveSubobject(rootSignatureSubObjectName);
+          subObjects->CreateRootSignature(rootSignatureSubObjectName,
+                                          notALocalRS,
+                                          extendedRootSig.data(),
+                                          static_cast<uint32_t>(extendedRootSig.size()));
+          break;
+        }
+      }
+    }
+  }
+}
+
+    // Set up a UAV with structure of a single int
 llvm::CallInst *CreateUAV(DxilModule &DM, IRBuilder<> &Builder,
                           unsigned int registerId, const char *name) {
   LLVMContext &Ctx = DM.GetModule()->getContext();
@@ -170,6 +262,11 @@ llvm::CallInst *CreateUAV(DxilModule &DM, IRBuilder<> &Builder,
   if (UAVStructTy == nullptr) {
     SmallVector<llvm::Type *, 1> Elements{Type::getInt32Ty(Ctx)};
     UAVStructTy = llvm::StructType::create(Elements, PIXStructTypeName);
+
+    // Since we only have to do this once per module, we can do it now when
+    // we're adding the singular UAV structure type to the module:
+    AddUAVToDxilDefinedGlobalRootSignatures(DM);
+    AddUAVToShaderAttributeRootSignature(DM);
   }
 
   std::unique_ptr<DxilResource> pUAV = llvm::make_unique<DxilResource>();

lib/DxilRootSignature/DxilRootSignatureSerializer.cpp

@@ -332,93 +332,6 @@ void SerializeRootSignature(const DxilVersionedRootSignatureDesc *pRootSignature
   }
 }
 
-//=============================================================================
-//
-// CVersionedRootSignatureDeserializer.
-//
-//=============================================================================
-class CVersionedRootSignatureDeserializer {
-protected:
-  const DxilVersionedRootSignatureDesc *m_pRootSignature;
-  const DxilVersionedRootSignatureDesc *m_pRootSignature10;
-  const DxilVersionedRootSignatureDesc *m_pRootSignature11;
-
-public:
-  CVersionedRootSignatureDeserializer();
-  ~CVersionedRootSignatureDeserializer();
-
-  void Initialize(_In_reads_bytes_(SrcDataSizeInBytes) const void *pSrcData,
-                  _In_ uint32_t SrcDataSizeInBytes);
-
-  const DxilVersionedRootSignatureDesc *GetRootSignatureDescAtVersion(DxilRootSignatureVersion convertToVersion);
-
-  const DxilVersionedRootSignatureDesc *GetUnconvertedRootSignatureDesc();
-};
-
-CVersionedRootSignatureDeserializer::CVersionedRootSignatureDeserializer()
-  : m_pRootSignature(nullptr)
-  , m_pRootSignature10(nullptr)
-  , m_pRootSignature11(nullptr) {
-}
-
-CVersionedRootSignatureDeserializer::~CVersionedRootSignatureDeserializer() {
-  DeleteRootSignature(m_pRootSignature10);
-  DeleteRootSignature(m_pRootSignature11);
-}
-
-void CVersionedRootSignatureDeserializer::Initialize(_In_reads_bytes_(SrcDataSizeInBytes) const void *pSrcData,
-                                                     _In_ uint32_t SrcDataSizeInBytes) {
-  const DxilVersionedRootSignatureDesc *pRootSignature = nullptr;
-  DeserializeRootSignature(pSrcData, SrcDataSizeInBytes, &pRootSignature);
-
-  switch (pRootSignature->Version) {
-  case DxilRootSignatureVersion::Version_1_0:
-    m_pRootSignature10 = pRootSignature;
-    break;
-
-  case DxilRootSignatureVersion::Version_1_1:
-    m_pRootSignature11 = pRootSignature;
-    break;
-
-  default:
-    DeleteRootSignature(pRootSignature);
-    return;
-  }
-
-  m_pRootSignature = pRootSignature;
-}
-
-const DxilVersionedRootSignatureDesc *
-CVersionedRootSignatureDeserializer::GetUnconvertedRootSignatureDesc() {
-  return m_pRootSignature;
-}
-
-const DxilVersionedRootSignatureDesc *
-CVersionedRootSignatureDeserializer::GetRootSignatureDescAtVersion(DxilRootSignatureVersion ConvertToVersion) {
-  switch (ConvertToVersion) {
-  case DxilRootSignatureVersion::Version_1_0:
-    if (m_pRootSignature10 == nullptr) {
-      ConvertRootSignature(m_pRootSignature,
-                           ConvertToVersion,
-                           (const DxilVersionedRootSignatureDesc **)&m_pRootSignature10);
-    }
-    return m_pRootSignature10;
-
-  case DxilRootSignatureVersion::Version_1_1:
-    if (m_pRootSignature11 == nullptr) {
-      ConvertRootSignature(m_pRootSignature,
-                           ConvertToVersion,
-                           (const DxilVersionedRootSignatureDesc **)&m_pRootSignature11);
-    }
-    return m_pRootSignature11;
-
-  default:
-    IFTBOOL(false, E_FAIL);
-  }
-
-  return nullptr;
-}
-
 template<typename T_ROOT_SIGNATURE_DESC,
          typename T_ROOT_PARAMETER,
          typename T_ROOT_DESCRIPTOR,

lib/HLSL/DxilValidation.cpp

@@ -5850,13 +5850,14 @@ HRESULT ValidateDxilBitcode(
     IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pOutputStream));
     pOutputStream->Reserve(pWriter->size());
     pWriter->write(pOutputStream);
+    DxilVersionedRootSignature desc;
     try {
-      const DxilVersionedRootSignatureDesc* pDesc = nullptr;
-      DeserializeRootSignature(SerializedRootSig.data(), SerializedRootSig.size(), &pDesc);
-      if (!pDesc) {
+      DeserializeRootSignature(SerializedRootSig.data(),
+                               SerializedRootSig.size(), desc.get_address_of());
+      if (!desc.get()) {
         return DXC_E_INCORRECT_ROOT_SIGNATURE;
       }
-      IFTBOOL(VerifyRootSignatureWithShaderPSV(pDesc,
+      IFTBOOL(VerifyRootSignatureWithShaderPSV(desc.get(),
                                                dxilModule.GetShaderModel()->GetKind(),
                                                pOutputStream->GetPtr(), pWriter->size(),
                                                DiagStream), DXC_E_INCORRECT_ROOT_SIGNATURE);