Patch PSV chunk to fix dxil validation warning

renderdoc/driver/d3d12/d3d12_shader_feedback.cpp

@@ -458,6 +458,10 @@ static bool AnnotateDXILShader(const DXBC::DXBCContainer *dxbc, uint32_t space,
     entryPoints->children[0] = entry;
   }
 
+  // get the editor to patch PSV0 with our extra UAV
+  editor.RegisterUAV(DXIL::DXILResourceType::ByteAddressUAV, space, 0, 0,
+                     DXIL::ResourceKind::RawBuffer);
+
   DXIL::Function *f = editor.GetFunctionByName(entryName);
 
   if(!f)

renderdoc/driver/shaders/dxbc/dxbc_bytecode_editor.cpp

@@ -55,7 +55,10 @@ ProgramEditor::ProgramEditor(const DXBC::DXBCContainer *container, bytebuf &outB
 
 ProgramEditor::~ProgramEditor()
 {
-  DXBC::DXBCContainer::ReplaceDXBCBytecode(m_OutBlob, EncodeProgram());
+  rdcarray<uint32_t> encoded = EncodeProgram();
+  // only one of these fourcc's will be present, so we just try to replace both
+  DXBC::DXBCContainer::ReplaceChunk(m_OutBlob, MAKE_FOURCC('S', 'H', 'E', 'X'), encoded);
+  DXBC::DXBCContainer::ReplaceChunk(m_OutBlob, MAKE_FOURCC('S', 'H', 'D', 'R'), encoded);
 }
 
 /*

renderdoc/driver/shaders/dxbc/dxbc_container.cpp

@@ -654,66 +654,6 @@ void DXBCContainer::FillStateInstructionInfo(ShaderDebugState &state) const
   }
 }
 
-void DXBCContainer::ReplaceDXBCBytecode(bytebuf &ByteCode, const rdcarray<uint32_t> &replacement)
-{
-  FileHeader *header = (FileHeader *)ByteCode.data();
-
-  if(header->fourcc != FOURCC_DXBC)
-    return;
-
-  if(header->fileLength != (uint32_t)ByteCode.size())
-    return;
-
-  uint32_t *chunkOffsets =
-      (uint32_t *)(ByteCode.data() + sizeof(FileHeader));    // right after the header
-
-  for(uint32_t chunkIdx = 0; chunkIdx < header->numChunks; chunkIdx++)
-  {
-    uint32_t offs = chunkOffsets[chunkIdx];
-
-    uint32_t *fourcc = (uint32_t *)(ByteCode.data() + offs);
-    uint32_t *chunkSize = (uint32_t *)(fourcc + 1);
-
-    if(*fourcc == FOURCC_SHEX || *fourcc == FOURCC_SHDR)
-    {
-      int32_t diff = int32_t(replacement.byteSize()) - int32_t(*chunkSize);
-
-      *chunkSize = (uint32_t)replacement.byteSize();
-
-      if(diff == 0)
-      {
-        memcpy(ByteCode.data() + offs + 8, replacement.data(), replacement.byteSize());
-      }
-      else if(diff > 0)
-      {
-        const byte *replaceBytes = (const byte *)replacement.data();
-        ByteCode.insert(offs + 8, replaceBytes, diff);
-        memcpy(ByteCode.data() + offs + 8 + diff, replaceBytes + diff, replacement.byteSize() - diff);
-      }
-      else if(diff < 0)
-      {
-        ByteCode.erase(offs + 8, -diff);
-        memcpy(ByteCode.data() + offs + 8, replacement.data(), replacement.byteSize());
-      }
-
-      // fixup offsets of chunks after this point
-      header = (FileHeader *)ByteCode.data();
-      chunkOffsets = (uint32_t *)(ByteCode.data() + sizeof(FileHeader));
-
-      header->fileLength += diff;
-
-      chunkIdx++;
-
-      for(; chunkIdx < header->numChunks; chunkIdx++)
-        chunkOffsets[chunkIdx] += diff;
-
-      break;
-    }
-  }
-
-  HashContainer(ByteCode.data(), ByteCode.size());
-}
-
 void DXBCContainer::StripDXILDebugInfo(bytebuf &ByteCode)
 {
   FileHeader *header = (FileHeader *)ByteCode.data();
@@ -753,7 +693,8 @@ void DXBCContainer::StripDXILDebugInfo(bytebuf &ByteCode)
   HashContainer(ByteCode.data(), ByteCode.size());
 }
 
-void DXBCContainer::ReplaceDXILBytecode(bytebuf &ByteCode, const bytebuf &replacement)
+void DXBCContainer::ReplaceChunk(bytebuf &ByteCode, uint32_t fourcc, const byte *replacement,
+                                 size_t size)
 {
   FileHeader *header = (FileHeader *)ByteCode.data();
 
@@ -770,29 +711,29 @@ void DXBCContainer::ReplaceDXILBytecode(bytebuf &ByteCode, const bytebuf &replac
   {
     uint32_t offs = chunkOffsets[chunkIdx];
 
-    uint32_t *fourcc = (uint32_t *)(ByteCode.data() + offs);
-    uint32_t *chunkSize = (uint32_t *)(fourcc + 1);
+    uint32_t *chunkFourcc = (uint32_t *)(ByteCode.data() + offs);
+    uint32_t *chunkSize = (uint32_t *)(chunkFourcc + 1);
 
-    if(*fourcc == FOURCC_DXIL)
+    if(*chunkFourcc == fourcc)
     {
-      int32_t diff = int32_t(replacement.size()) - int32_t(*chunkSize);
+      int32_t diff = int32_t(size) - int32_t(*chunkSize);
 
-      *chunkSize = (uint32_t)replacement.size();
+      *chunkSize = (uint32_t)size;
 
       if(diff == 0)
       {
-        memcpy(ByteCode.data() + offs + 8, replacement.data(), replacement.size());
+        memcpy(ByteCode.data() + offs + 8, replacement, size);
       }
       else if(diff > 0)
       {
-        const byte *replaceBytes = (const byte *)replacement.data();
+        const byte *replaceBytes = (const byte *)replacement;
         ByteCode.insert(offs + 8, replaceBytes, diff);
-        memcpy(ByteCode.data() + offs + 8 + diff, replaceBytes + diff, replacement.size() - diff);
+        memcpy(ByteCode.data() + offs + 8 + diff, replaceBytes + diff, size - diff);
       }
       else if(diff < 0)
       {
         ByteCode.erase(offs + 8, -diff);
-        memcpy(ByteCode.data() + offs + 8, replacement.data(), replacement.size());
+        memcpy(ByteCode.data() + offs + 8, replacement, size);
       }
 
       // fixup offsets of chunks after this point
@@ -806,11 +747,42 @@ void DXBCContainer::ReplaceDXILBytecode(bytebuf &ByteCode, const bytebuf &replac
       for(; chunkIdx < header->numChunks; chunkIdx++)
         chunkOffsets[chunkIdx] += diff;
 
-      break;
+      HashContainer(ByteCode.data(), ByteCode.size());
+      return;
+    }
+  }
+}
+
+const byte *DXBCContainer::FindChunk(const bytebuf &ByteCode, uint32_t fourcc, size_t &size)
+{
+  const FileHeader *header = (const FileHeader *)ByteCode.data();
+
+  size = 0;
+
+  if(header->fourcc != FOURCC_DXBC)
+    return NULL;
+
+  if(header->fileLength != (uint32_t)ByteCode.size())
+    return NULL;
+
+  const uint32_t *chunkOffsets =
+      (const uint32_t *)(ByteCode.data() + sizeof(FileHeader));    // right after the header
+
+  for(uint32_t chunkIdx = 0; chunkIdx < header->numChunks; chunkIdx++)
+  {
+    uint32_t offs = chunkOffsets[chunkIdx];
+
+    const uint32_t *chunkFourcc = (uint32_t *)(ByteCode.data() + offs);
+    const uint32_t *chunkSize = (uint32_t *)(chunkFourcc + 1);
+
+    if(*chunkFourcc == fourcc)
+    {
+      size = *chunkSize;
+      return ByteCode.data() + offs + 8;
     }
   }
 
-  HashContainer(ByteCode.data(), ByteCode.size());
+  return NULL;
 }
 
 void DXBCContainer::GetHash(uint32_t hash[4], const void *ByteCode, size_t BytecodeLength)

renderdoc/driver/shaders/dxbc/dxbc_container.h

@@ -163,9 +163,16 @@ public:
   void FillTraceLineInfo(ShaderDebugTrace &trace) const;
   void FillStateInstructionInfo(ShaderDebugState &state) const;
 
-  static void ReplaceDXBCBytecode(bytebuf &ByteCode, const rdcarray<uint32_t> &replacement);
   static void StripDXILDebugInfo(bytebuf &ByteCode);
-  static void ReplaceDXILBytecode(bytebuf &ByteCode, const bytebuf &replacement);
+  static void ReplaceChunk(bytebuf &ByteCode, uint32_t fourcc, const byte *replacement, size_t size);
+
+  template <typename T>
+  static void ReplaceChunk(bytebuf &ByteCode, uint32_t fourcc, const rdcarray<T> &replacement)
+  {
+    ReplaceChunk(ByteCode, fourcc, (byte *)replacement.data(), replacement.byteSize());
+  }
+
+  static const byte *FindChunk(const bytebuf &ByteCode, uint32_t fourcc, size_t &size);
 
   const DXBCBytecode::Program *GetDXBCByteCode() const { return m_DXBCByteCode; }
   DXBCBytecode::Program *GetDXBCByteCode() { return m_DXBCByteCode; }

renderdoc/driver/shaders/dxil/dxil_bytecode_editor.cpp

@@ -471,7 +471,7 @@ ProgramEditor::~ProgramEditor()
   GetInt32Type();
 
   // replace the DXIL bytecode in the container with
-  DXBC::DXBCContainer::ReplaceDXILBytecode(m_OutBlob, EncodeProgram());
+  DXBC::DXBCContainer::ReplaceChunk(m_OutBlob, MAKE_FOURCC('D', 'X', 'I', 'L'), EncodeProgram());
 
 #if ENABLED(RDOC_DEVEL)
   // on debug builds, run through dxil for "validation" if it's available.
@@ -1897,6 +1897,87 @@ bytebuf ProgramEditor::EncodeProgram() const
   return ret;
 }
 
+struct ResourceBind0
+{
+  DXILResourceType type;
+  uint32_t space;
+  uint32_t regBase;
+  uint32_t regEnd;
+};
+
+struct ResourceBind1 : ResourceBind0
+{
+  ResourceKind kind;
+  uint32_t flags;
+};
+
+// this function should be expanded in future, maybe to automatically re-write the PSV0 from the
+// DXIL data in ~ProgramEditor()
+void ProgramEditor::RegisterUAV(DXILResourceType type, uint32_t space, uint32_t regBase,
+                                uint32_t regEnd, ResourceKind kind)
+{
+  size_t sz = 0;
+  const byte *psv0 = DXBC::DXBCContainer::FindChunk(m_OutBlob, MAKE_FOURCC('P', 'S', 'V', '0'), sz);
+
+  ResourceBind1 bind = {};
+  bind.type = type;
+  bind.space = space;
+  bind.regBase = regBase;
+  bind.regEnd = regEnd;
+  bind.kind = kind;
+
+  if(psv0)
+  {
+    bytebuf psv0blob(psv0, sz);
+
+    byte *begin = psv0blob.data();
+    byte *end = begin + sz;
+
+    byte *cur = begin;
+
+    uint32_t *headerSize = (uint32_t *)cur;
+    cur += sizeof(uint32_t);
+    if(cur >= end)
+      return;
+
+    // don't need to patch the header
+    cur += *headerSize;
+    if(cur >= end)
+      return;
+
+    uint32_t *numResources = (uint32_t *)cur;
+    cur += sizeof(uint32_t);
+    if(cur >= end)
+      return;
+
+    uint32_t *resourceBindSize = (uint32_t *)cur;
+    cur += sizeof(uint32_t);
+    if(cur >= end)
+      return;
+
+    // fortunately UAVs are the last entry so we don't need to walk the list to insert in the right
+    // place, we can just add it at the end
+    cur += (*resourceBindSize) * (*numResources);
+    if(cur >= end)
+      return;
+
+    // add an extra resource
+    (*numResources)++;
+
+    if(*resourceBindSize == sizeof(ResourceBind1) || *resourceBindSize == sizeof(ResourceBind0))
+    {
+      psv0blob.insert(cur - begin, (byte *)&bind, *resourceBindSize);
+    }
+    else
+    {
+      RDCERR("Unexpected resource bind size %u", *resourceBindSize);
+      return;
+    }
+
+    DXBC::DXBCContainer::ReplaceChunk(m_OutBlob, MAKE_FOURCC('P', 'S', 'V', '0'), psv0blob);
+  }
+}
+
 void ProgramEditor::EncodeConstants(LLVMBC::BitcodeWriter &writer, const rdcarray<Value> &values,
                                     const rdcarray<Constant> &constants) const
 {

renderdoc/driver/shaders/dxil/dxil_bytecode_editor.h

@@ -25,6 +25,7 @@
 #pragma once
 
 #include "dxil_bytecode.h"
+#include "dxil_common.h"
 
 namespace DXBC
 {
@@ -38,6 +39,20 @@ class BitcodeWriter;
 
 namespace DXIL
 {
+enum class DXILResourceType
+{
+  Unknown,
+  Sampler,
+  CBuffer,
+  TypedSRV,
+  ByteAddressSRV,
+  StructuredSRV,
+  TypedUAV,
+  ByteAddressUAV,
+  StructuredUAV,
+  StructuredUAVWithCounter,
+};
+
 class ProgramEditor : public Program
 {
 public:
@@ -70,6 +85,9 @@ public:
 
   Instruction *AddInstruction(Function *f, size_t idx, const Instruction &inst);
 
+  void RegisterUAV(DXILResourceType type, uint32_t space, uint32_t regBase, uint32_t regEnd,
+                   ResourceKind kind);
+
 private:
   bytebuf &m_OutBlob;