Add some higher level SPIR-V editing helpers

* These both automate some common (and verbose) tasks and also add extra helpers
  for declaring a new output buffer that might come via BDA or via a binding.

renderdoc/driver/shaders/spirv/spirv_common.h

@@ -455,3 +455,32 @@ enum class ShaderBuiltin : uint32_t;
 
 ShaderStage MakeShaderStage(rdcspv::ExecutionModel model);
 ShaderBuiltin MakeShaderBuiltin(ShaderStage stage, const rdcspv::BuiltIn el);
+
+enum class BufferStorageMode : uint32_t
+{
+  Unknown,
+  // patched in plain descriptor
+  Descriptor,
+  // legacy path - EXT_buffer_device_address only
+  EXT_bda,
+  // KHR_buffer_device_address, but with uint2 addresses to not require int64 types/maths
+  KHR_bda32,
+  // like above, but requiring int64
+  KHR_bda64,
+};
+
+constexpr bool IsBDA(BufferStorageMode mode)
+{
+  return mode == BufferStorageMode::EXT_bda || mode == BufferStorageMode::KHR_bda32 ||
+         mode == BufferStorageMode::KHR_bda64;
+}
+
+constexpr bool IsKHRBDA(BufferStorageMode mode)
+{
+  return mode == BufferStorageMode::KHR_bda32 || mode == BufferStorageMode::KHR_bda64;
+}
+
+constexpr bool IsBinding(BufferStorageMode mode)
+{
+  return mode == BufferStorageMode::Descriptor;
+}

renderdoc/driver/shaders/spirv/spirv_editor.cpp

@@ -80,6 +80,8 @@ void Editor::Prepare()
   // instead of Uniform + BufferBlock
   if(m_MajorVersion > 1 || m_MinorVersion >= 3)
     m_StorageBufferClass = rdcspv::StorageClass::StorageBuffer;
+  else
+    m_StorageBufferClass = rdcspv::StorageClass::Uniform;
 
   // find any empty sections and insert a nop into the stream there. We need to fixup later section
   // offsets by hand as addWords doesn't handle empty sections properly (it thinks we're inserting
@@ -192,11 +194,16 @@ Id Editor::MakeId()
 
 void Editor::DecorateStorageBufferStruct(Id id)
 {
-  // set bufferblock if needed
   if(m_StorageBufferClass == rdcspv::StorageClass::Uniform)
-    AddDecoration(rdcspv::OpDecorate(id, rdcspv::Decoration::BufferBlock));
+  {
+    if(!m_BufferBlockTypes.contains(id))
+      AddDecoration(rdcspv::OpDecorate(id, rdcspv::Decoration::BufferBlock));
+  }
   else
-    AddDecoration(rdcspv::OpDecorate(id, rdcspv::Decoration::Block));
+  {
+    if(!m_BlockTypes.contains(id))
+      AddDecoration(rdcspv::OpDecorate(id, rdcspv::Decoration::Block));
+  }
 }
 
 void Editor::SetName(Id id, const rdcstr &name)
@@ -403,6 +410,38 @@ Iter Editor::GetEntry(Id id)
   return Iter();
 }
 
+Id Editor::FindEntryID(ShaderEntryPoint entry)
+{
+  rdcspv::Id entryID;
+  for(rdcspv::Iter it = Begin(rdcspv::Section::EntryPoints), end = End(rdcspv::Section::EntryPoints);
+      it < end; ++it)
+  {
+    rdcspv::OpEntryPoint e(it);
+    if(e.name == entry.name && MakeShaderStage(e.executionModel) == entry.stage)
+      return e.entryPoint;
+  }
+  return rdcspv::Id();
+}
+
+void Editor::AddEntryGlobals(Id entry, const rdcarray<Id> &newGlobals)
+{
+  if(!newGlobals.empty())
+  {
+    rdcspv::Iter it = GetEntry(entry);
+
+    // this copies into the helper struct
+    rdcspv::OpEntryPoint e(it);
+
+    // add our IDs
+    e.iface.append(newGlobals);
+
+    // erase the old one
+    Remove(it);
+
+    AddOperation(it, e);
+  }
+}
+
 rdcpair<Id, Id> Editor::AddBuiltinInputLoad(OperationList &ops, ShaderStage stage, BuiltIn builtin,
                                             Id type)
 {
@@ -458,6 +497,33 @@ Id Editor::DeclareStructType(const rdcarray<Id> &members)
   return typeId;
 }
 
+rdcspv::Id Editor::DeclareStructType(const rdcstr &name, const rdcarray<StructMember> &members)
+{
+  Id typeId = MakeId();
+
+  rdcarray<Id> memberTypes;
+
+  for(uint32_t i = 0; i < members.size(); i++)
+    memberTypes.push_back(members[i].type);
+
+  AddType(OpTypeStruct(typeId, memberTypes));
+
+  for(uint32_t i = 0; i < members.size(); i++)
+  {
+    if(!members[i].name.empty())
+      SetMemberName(typeId, i, members[i].name);
+
+    if(members[i].offset != ~0U)
+      AddDecoration(rdcspv::OpMemberDecorate(
+          typeId, i, rdcspv::DecorationParam<rdcspv::Decoration::Offset>(members[i].offset)));
+  }
+
+  if(!name.empty())
+    SetName(typeId, name);
+
+  return typeId;
+}
+
 Id Editor::AddOperation(Iter iter, const Operation &op)
 {
   if(!iter)
@@ -542,6 +608,10 @@ void Editor::RegisterOp(Iter it)
       bindings[decorate.target].set = decorate.decoration.descriptorSet;
     if(decorate.decoration == Decoration::Binding)
       bindings[decorate.target].binding = decorate.decoration.binding;
+    if(decorate.decoration == Decoration::Block)
+      m_BlockTypes.push_back(decorate.target);
+    if(decorate.decoration == Decoration::BufferBlock)
+      m_BufferBlockTypes.push_back(decorate.target);
   }
 }
 
@@ -600,6 +670,10 @@ void Editor::UnregisterOp(Iter it)
       bindings[decorate.target].set = Binding().set;
     if(decorate.decoration == Decoration::Binding)
       bindings[decorate.target].binding = Binding().binding;
+    if(decorate.decoration == Decoration::Block)
+      m_BlockTypes.removeOne(decorate.target);
+    if(decorate.decoration == Decoration::BufferBlock)
+      m_BufferBlockTypes.removeOne(decorate.target);
   }
 }
 
@@ -730,6 +804,187 @@ Operation Editor::MakeDeclaration(const FunctionType &f)
   return OpTypeFunction(Id(), f.returnId, f.argumentIds);
 }
 
+void Editor::OffsetBindingsToMatchReservation(size_t numReservedBindings)
+{
+  if(!IsBinding(m_StorageMode))
+    return;
+
+  // patch all bindings up trivially to account for the extra reservation
+  for(Iter it = Begin(rdcspv::Section::Annotations), end = End(rdcspv::Section::Annotations);
+      it < end; ++it)
+  {
+    // we will use descriptor set 0 for our own purposes if we don't have a buffer address.
+    //
+    // Since bindings are arbitrary, we just increase all user bindings to make room, and we'll
+    // redeclare the descriptor set layouts and pipeline layout. This is inevitable in the case
+    // where all descriptor sets are already used. In theory we only have to do this with set 0,
+    // but that requires knowing which variables are in set 0 and it's simpler to increase all
+    // bindings.
+    if(it.opcode() == Op::Decorate)
+    {
+      OpDecorate dec(it);
+      if(dec.decoration == Decoration::Binding)
+      {
+        RDCASSERT(dec.decoration.binding < (0xffffffff - numReservedBindings));
+        dec.decoration.binding += (uint32_t)numReservedBindings;
+        it = dec;
+      }
+    }
+  }
+}
+
+StorageClass Editor::PrepareAddedBufferAccess()
+{
+  if(IsBinding(m_StorageMode))
+  {
+    return StorageBufferClass();
+  }
+  else if(IsBDA(m_StorageMode))
+  {
+    // add the extension
+    AddExtension(m_StorageMode == BufferStorageMode::EXT_bda ? "SPV_EXT_physical_storage_buffer"
+                                                             : "SPV_KHR_physical_storage_buffer");
+
+    // change the memory model to physical storage buffer 64
+    Iter it = Begin(Section::MemoryModel);
+    OpMemoryModel model(it);
+    model.addressingModel = AddressingModel::PhysicalStorageBuffer64;
+    it = model;
+
+    // add capabilities
+    AddCapability(Capability::PhysicalStorageBufferAddresses);
+
+    // for simplicity on KHR we always load from uint2 so we're compatible with the case where int64
+    // isn't supported
+    if(m_StorageMode == BufferStorageMode::EXT_bda || m_StorageMode == BufferStorageMode::KHR_bda64)
+    {
+      AddCapability(Capability::Int64);
+    }
+
+    return StorageClass::PhysicalStorageBuffer;
+  }
+  else
+  {
+    RDCERR("Added buffer access can't be used until storage mode is set");
+    return StorageClass::Invalid;
+  }
+}
+
+Id Editor::LoadBufferVariable(OperationList &ops, rdcpair<Id, Id> var)
+{
+  if(IsBinding(m_StorageMode))
+  {
+    return var.second;
+  }
+  else if(IsBDA(m_StorageMode))
+  {
+    Id ret;
+    // if we don't have the struct as a bind, we need to cast it from the pointer. In
+    // KHR_buffer_device_address we bitcast since we store it as a uint2
+    if(m_StorageMode == BufferStorageMode::KHR_bda32)
+      ret = ops.add(OpBitcast(var.first, MakeId(), var.second));
+    else
+      ret = ops.add(OpConvertUToPtr(var.first, MakeId(), var.second));
+
+    SetName(ret, "loaded_buf");
+
+    return ret;
+  }
+  else
+  {
+    RDCERR("Added buffer access can't be used until storage mode is set");
+    return Id();
+  }
+}
+
+rdcpair<Id, Id> Editor::AddBufferVariable(rdcarray<Id> &addedGlobals, Id varType, const rdcstr &name,
+                                          uint32_t binding, uint32_t specID, uint64_t fixedAddr)
+{
+  rdcpair<Id, Id> ret;
+
+  if(IsBinding(m_StorageMode))
+  {
+    StorageClass bufferClass = StorageBufferClass();
+
+    // the pointers are SSBO pointers
+    ret.first = DeclareType(Pointer(varType, bufferClass));
+
+    // add our SSBO variable, at set 0 binding 0
+    ret.second = MakeId();
+    AddVariable(OpVariable(ret.first, ret.second, bufferClass));
+    AddDecoration(OpDecorate(ret.second, DecorationParam<Decoration::DescriptorSet>(0)));
+    AddDecoration(OpDecorate(ret.second, DecorationParam<Decoration::Binding>(binding)));
+
+    if(EntryPointAllGlobals())
+      addedGlobals.push_back(ret.second);
+
+    SetName(ret.second, name);
+
+    if(GetDataType(varType).type == DataType::StructType)
+      DecorateStorageBufferStruct(varType);
+  }
+  else if(IsBDA(m_StorageMode))
+  {
+    StorageClass bufferClass = StorageClass::PhysicalStorageBuffer;
+
+    ret.first = DeclareType(Pointer(varType, bufferClass));
+
+    if(fixedAddr != 0)
+    {
+      // for simplicity on KHR we always load from uint2 so we're compatible with the case where
+      // int64 isn't supported
+      if(m_StorageMode == BufferStorageMode::KHR_bda32)
+      {
+        Id addressConstantLSB = AddConstantImmediate<uint32_t>(fixedAddr & 0xffffffffu);
+        Id addressConstantMSB = AddConstantImmediate<uint32_t>((fixedAddr >> 32) & 0xffffffffu);
+        SetName(addressConstantLSB, name + "_addressLSB");
+        SetName(addressConstantMSB, name + "_addressMSB");
+
+        Id uint2 = DeclareType(Vector(scalar<uint32_t>(), 2));
+
+        ret.second = AddConstant(
+            OpConstantComposite(uint2, MakeId(), {addressConstantLSB, addressConstantMSB}));
+      }
+      else
+      {
+        // declare the address constants and make our pointers physical storage buffer pointers
+        ret.second = AddConstantImmediate<uint64_t>(fixedAddr);
+      }
+    }
+    else
+    {
+      if(m_StorageMode == BufferStorageMode::KHR_bda32)
+      {
+        Id addressConstantLSB = AddSpecConstantImmediate<uint32_t>(0U, specID);
+        Id addressConstantMSB = AddSpecConstantImmediate<uint32_t>(0U, specID + 1);
+        SetName(addressConstantLSB, name + "_addressLSB");
+        SetName(addressConstantMSB, name + "_addressMSB");
+
+        Id uint2 = DeclareType(Vector(scalar<uint32_t>(), 2));
+
+        ret.second = AddConstant(
+            OpSpecConstantComposite(uint2, MakeId(), {addressConstantLSB, addressConstantMSB}));
+      }
+      else
+      {
+        ret.second = AddSpecConstantImmediate<uint64_t>(0ULL, specID);
+      }
+    }
+
+    SetName(ret.second, name + "_address");
+
+    // structs are block decorated
+    if(GetDataType(varType).type == DataType::StructType && !m_BlockTypes.contains(varType))
+      AddDecoration(OpDecorate(varType, Decoration::Block));
+  }
+  else
+  {
+    RDCERR("Added buffer access can't be used until storage mode is set");
+  }
+
+  return ret;
+}
+
 #define TYPETABLE(StructType, variable)                                \
   template <>                                                          \
   std::map<StructType, Id> &Editor::GetTable<StructType>()             \

renderdoc/driver/shaders/spirv/spirv_editor.h

@@ -58,6 +58,13 @@ struct Binding
   bool operator==(const Binding &o) const { return set == o.set && binding == o.binding; }
 };
 
+struct StructMember
+{
+  Id type;
+  rdcstr name;
+  uint32_t offset;
+};
+
 template <typename SPIRVType>
 using TypeToId = std::pair<SPIRVType, Id>;
 
@@ -71,6 +78,7 @@ public:
   ~Editor();
 
   void Prepare();
+  void SetBufferStorageMode(BufferStorageMode mode) { m_StorageMode = mode; }
   void CreateEmpty(uint32_t major, uint32_t minor);
 
   Id MakeId();
@@ -92,6 +100,26 @@ public:
     iter.nopRemove();
   }
 
+  void OffsetBindingsToMatchReservation(size_t numReservedBindings);
+  StorageClass PrepareAddedBufferAccess();
+  rdcpair<Id, Id> AddBufferVariable(rdcarray<Id> &addedGlobals, Id varType, const rdcstr &name,
+                                    uint32_t binding, uint32_t specID, uint64_t fixedAddr);
+  Id LoadBufferVariable(OperationList &ops, rdcpair<Id, Id> var);
+
+  Id FindEntryID(ShaderEntryPoint entry);
+  void AddEntryGlobals(Id entry, const rdcarray<Id> &newGlobals);
+
+  rdcpair<Id, Id> AddBuiltinInputLoad(OperationList &ops, ShaderStage stage, BuiltIn builtin,
+                                      Id type);
+  Id AddBuiltinInputLoad(OperationList &ops, rdcarray<Id> &addedGlobals, ShaderStage stage,
+                         BuiltIn builtin, Id type)
+  {
+    rdcpair<Id, Id> ret = AddBuiltinInputLoad(ops, stage, builtin, type);
+    if(ret.second != rdcspv::Id())
+      addedGlobals.push_back(ret.second);
+    return ret.first;
+  }
+
   StorageClass StorageBufferClass() { return m_StorageBufferClass; }
   bool EntryPointAllGlobals() { return m_MajorVersion > 1 || m_MinorVersion >= 4; }
   void DecorateStorageBufferStruct(Id id);
@@ -175,10 +203,8 @@ public:
     return it->second;
   }
 
-  rdcpair<Id, Id> AddBuiltinInputLoad(OperationList &ops, ShaderStage stage, BuiltIn builtin,
-                                      Id type);
-
   Id DeclareStructType(const rdcarray<Id> &members);
+  Id DeclareStructType(const rdcstr &name, const rdcarray<StructMember> &members);
 
   // helper for AddConstant
   template <typename T>
@@ -270,7 +296,11 @@ private:
 
   std::map<BuiltIn, BuiltinInputData> builtinInputs;
 
+  BufferStorageMode m_StorageMode = BufferStorageMode::Unknown;
+
   std::map<Id, Binding> bindings;
+  rdcarray<Id> m_BufferBlockTypes;
+  rdcarray<Id> m_BlockTypes;
 
   std::map<Scalar, Id> scalarTypeToId;
   std::map<Vector, Id> vectorTypeToId;