Add ability to re-use or declare new types

renderdoc/driver/shaders/spirv/spirv_editor.cpp

@@ -23,44 +23,177 @@
  ******************************************************************************/
 
 #include "spirv_editor.h"
+#include <utility>
 #include "common/common.h"
 
-SPIRVEditor::SPIRVEditor(std::vector<uint32_t> &spirvWords) : spirv(spirvWords)
+SPIRVScalar::SPIRVScalar(SPIRVIterator it)
 {
-  if(spirvWords.size() < 5 || spirvWords[0] != spv::MagicNumber)
+  type = it.opcode();
+
+  if(type != spv::OpTypeBool)
+    width = it.word(2);
+  else
+    width = 0;
+
+  if(type == spv::OpTypeInt)
+    signedness = it.word(3) == 1;
+  else
+    signedness = false;
+}
+
+SPIRVEditor::SPIRVEditor(const std::vector<uint32_t> &spirvWords) : spirv(spirvWords)
+{
+  if(spirv.size() < 5 || spirv[0] != spv::MagicNumber)
   {
     RDCERR("Empty or invalid SPIR-V module");
     return;
   }
 
-  moduleVersion.major = uint8_t((spirvWords[1] & 0x00ff0000) >> 16);
-  moduleVersion.minor = uint8_t((spirvWords[1] & 0x0000ff00) >> 8);
-  generator = spirvWords[2];
-  idBound = SPIRVIterator(spirvWords, 3);
+  moduleVersion.major = uint8_t((spirv[1] & 0x00ff0000) >> 16);
+  moduleVersion.minor = uint8_t((spirv[1] & 0x0000ff00) >> 8);
+  generator = spirv[2];
+  idBound = SPIRVIterator(spirv, 3);
+  ids.resize(*idBound);
 
   // [4] is reserved
   RDCASSERT(spirv[4] == 0);
 
-  SPIRVIterator it(spirvWords, 5);
+  bool decorations = false;
 
-  while(it)
+  bool headerSections = false;
+
+  for(SPIRVIterator it(spirv, 5); it; it++)
   {
     spv::Op opcode = it.opcode();
 
+    // identify entry points
     if(opcode == spv::OpEntryPoint)
     {
       SPIRVEntry entry;
-      entry.iter = it;
+      entry.entryPoint = it;
       entry.id = it.word(2);
       entry.name = (const char *)&it.word(3);
 
       entries.push_back(entry);
+
+      headerSections = true;
     }
 
-    it++;
+    // identify functions
+    if(opcode == spv::OpFunction)
+    {
+      // if we don't have the types/variables iter yet, this is the point to set it
+      if(!typeVarSection.iter)
+        typeVarSection.iter = it;
+
+      uint32_t id = it.word(2);
+      ids[id] = it;
+
+      // if this is an entry point function, point the iter
+      for(SPIRVEntry &entry : entries)
+      {
+        if(entry.id == id)
+        {
+          entry.function = it;
+          break;
+        }
+      }
+
+      SPIRVFunction func;
+      func.id = id;
+      func.iter = it;
+
+      functions.push_back(func);
+    }
+
+    if(opcode == spv::OpDecorate || opcode == spv::OpMemberDecorate ||
+       opcode == spv::OpGroupDecorate || opcode == spv::OpGroupMemberDecorate ||
+       opcode == spv::OpDecorationGroup)
+    {
+      // when we hit the first decoration, this is the insert point for new debug instructions
+      if(!decorations)
+      {
+        debugSection.iter = it;
+      }
+
+      decorations = true;
+    }
+    else
+    {
+      // when we stop seeing decoration instructions, this is the insert point for them
+      if(decorations)
+      {
+        decorationSection.iter = it;
+        decorations = false;
+      }
+    }
+
+    // identify declared scalar/vector/matrix types
+    if(opcode == spv::OpTypeBool || opcode == spv::OpTypeInt || opcode == spv::OpTypeFloat)
+    {
+      uint32_t id = it.word(1);
+      ids[id] = it;
+
+      SPIRVScalar scalar(it);
+      scalarTypes[scalar] = id;
+    }
+
+    if(opcode == spv::OpTypeVector)
+    {
+      uint32_t id = it.word(1);
+      ids[id] = it;
+
+      SPIRVIterator scalarIt = ids[it.word(2)];
+
+      if(!scalarIt)
+      {
+        RDCERR("Vector type declared with unknown scalar component type %u", it.word(2));
+        continue;
+      }
+
+      vectorTypes[SPIRVVector(scalarIt, it.word(3))] = id;
+    }
+
+    if(opcode == spv::OpTypeMatrix)
+    {
+      uint32_t id = it.word(1);
+      ids[id] = it;
+
+      SPIRVIterator vectorIt = ids[it.word(2)];
+
+      if(!vectorIt)
+      {
+        RDCERR("Matrix type declared with unknown vector component type %u", it.word(2));
+        continue;
+      }
+
+      SPIRVIterator scalarIt = ids[vectorIt.word(2)];
+      uint32_t vectorDim = vectorIt.word(3);
+
+      matrixTypes[SPIRVMatrix(SPIRVVector(scalarIt, vectorDim), it.word(3))] = id;
+    }
+
+    if(opcode == spv::OpTypePointer)
+    {
+      uint32_t id = it.word(1);
+      ids[id] = it;
+
+      pointerTypes[SPIRVPointer(it.word(3), (spv::StorageClass)it.word(2))] = id;
+    }
   }
 }
 
+std::vector<uint32_t> SPIRVEditor::GetWords()
+{
+  std::vector<uint32_t> ret(spirv);
+
+  // add sections in reverse order so that each one doesn't perturb the offset for the next
+  for(const LogicalSection &section : {typeVarSection, decorationSection, debugSection})
+    ret.insert(ret.begin() + section.iter.offset, section.additions.begin(), section.additions.end());
+
+  return ret;
+}
+
 uint32_t SPIRVEditor::MakeId()
 {
   if(!idBound)
@@ -68,5 +201,109 @@ uint32_t SPIRVEditor::MakeId()
 
   uint32_t ret = *idBound;
   (*idBound)++;
+  ids.push_back(SPIRVIterator());
   return ret;
 }
+
+void SPIRVEditor::SetName(uint32_t id, const char *name)
+{
+  size_t sz = strlen(name);
+  std::vector<uint32_t> uintName((sz / 4) + 1);
+  memcpy(&uintName[0], name, sz);
+
+  uintName.insert(uintName.begin(), id);
+
+  SPIRVOperation op(spv::OpName, uintName);
+
+  debugSection.additions.insert(debugSection.additions.end(), op.begin(), op.end());
+}
+
+void SPIRVEditor::AddDecoration(const SPIRVOperation &op)
+{
+  decorationSection.additions.insert(decorationSection.additions.end(), op.begin(), op.end());
+}
+
+void SPIRVEditor::AddType(const SPIRVOperation &op)
+{
+  ids[op[1]] = SPIRVIterator(typeVarSection.additions, typeVarSection.additions.size());
+  typeVarSection.additions.insert(typeVarSection.additions.end(), op.begin(), op.end());
+}
+
+void SPIRVEditor::AddVariable(const SPIRVOperation &op)
+{
+  ids[op[2]] = SPIRVIterator(typeVarSection.additions, typeVarSection.additions.size());
+  typeVarSection.additions.insert(typeVarSection.additions.end(), op.begin(), op.end());
+}
+
+void SPIRVEditor::AddFunction(const SPIRVOperation *ops, size_t count)
+{
+  // because this is added to the end, we can add this immediately
+  ids[ops[0][2]] = SPIRVIterator(spirv, spirv.size());
+
+  auto insertIter = spirv.end();
+  for(size_t i = 0; i < count; i++)
+  {
+    spirv.insert(insertIter, ops[i].begin(), ops[i].end());
+    insertIter += ops[i].size();
+  }
+}
+
+uint32_t SPIRVEditor::DeclareType(const SPIRVScalar &scalar)
+{
+  auto it = scalarTypes.lower_bound(scalar);
+  if(it != scalarTypes.end() && it->first == scalar)
+    return it->second;
+
+  SPIRVOperation decl = scalar.decl();
+  uint32_t id = decl[1] = MakeId();
+  AddType(decl);
+
+  scalarTypes.insert(it, std::make_pair(scalar, id));
+
+  return id;
+}
+
+uint32_t SPIRVEditor::DeclareType(const SPIRVVector &vector)
+{
+  auto it = vectorTypes.lower_bound(vector);
+  if(it != vectorTypes.end() && it->first == vector)
+    return it->second;
+
+  uint32_t id = MakeId();
+  SPIRVOperation decl(spv::OpTypeVector, {id, DeclareType(vector.scalar), vector.count});
+  AddType(decl);
+
+  vectorTypes.insert(it, std::make_pair(vector, id));
+
+  return id;
+}
+
+uint32_t SPIRVEditor::DeclareType(const SPIRVMatrix &matrix)
+{
+  auto it = matrixTypes.lower_bound(matrix);
+  if(it != matrixTypes.end() && it->first == matrix)
+    return it->second;
+
+  uint32_t id = MakeId();
+  SPIRVOperation decl(spv::OpTypeVector, {id, DeclareType(matrix.vector), matrix.count});
+  AddType(decl);
+
+  matrixTypes.insert(it, std::make_pair(matrix, id));
+
+  return id;
+}
+
+uint32_t SPIRVEditor::DeclareType(const SPIRVPointer &pointer)
+{
+  auto it = pointerTypes.lower_bound(pointer);
+  if(it != pointerTypes.end() && it->first == pointer)
+    return it->second;
+
+  uint32_t id = MakeId();
+  SPIRVOperation decl(spv::OpTypePointer, {id, (uint32_t)pointer.storage, pointer.baseId});
+  AddType(decl);
+
+  pointerTypes.insert(it, std::make_pair(pointer, id));
+
+  return id;
+}

renderdoc/driver/shaders/spirv/spirv_editor.h

@@ -25,11 +25,13 @@
 #pragma once
 
 #include <stdint.h>
+#include <map>
 #include <string>
 #include <vector>
 #include "3rdparty/glslang/SPIRV/spirv.hpp"
 
 class SPIRVOperation;
+class SPIRVEditor;
 
 class SPIRVIterator
 {
@@ -45,7 +47,6 @@ public:
     return *this;
   }
   // utility functions
-  std::vector<uint32_t>::iterator it() { return words->begin() + offset; }
   operator bool() const { return words != NULL && offset < words->size(); }
   uint32_t &operator*() { return cur(); }
   const uint32_t &operator*() const { return cur(); }
@@ -56,11 +57,13 @@ public:
 private:
   inline uint32_t &cur() { return words->at(offset); }
   inline const uint32_t &cur() const { return words->at(offset); }
-  // we make this a friend so it can poke directly into words when it wants to edit
+  // we add some friend classes to poke directly into words when it wants to edit
   friend class SPIRVOperation;
-
-  std::vector<uint32_t> *words = NULL;
+  friend class SPIRVEditor;
+  std::vector<uint32_t>::iterator it() { return words->begin() + offset; }
+  std::vector<uint32_t>::const_iterator it() const { return words->cbegin() + offset; }
   size_t offset = 0;
+  std::vector<uint32_t> *words = NULL;
 };
 
 class SPIRVOperation
@@ -68,15 +71,17 @@ class SPIRVOperation
 public:
   // constructor of a synthetic operation, from an operation & subsequent words, calculates the
   // length then constructs the first word with opcode + length.
-  template <typename WordContainer>
-  SPIRVOperation(spv::Op op, const WordContainer &data)
+  SPIRVOperation(spv::Op op, std::initializer_list<uint32_t> data)
   {
-    auto a = std::begin(data);
-    auto b = std::end(data);
-    size_t count = 1 + (b - a);
-    words.resize(data.size());
-    words[0] = MakeHeader(op, data.size());
-    words.insert(words.begin() + 1, a, b);
+    words.push_back(MakeHeader(op, data.size() + 1));
+    words.insert(words.begin() + 1, data.begin(), data.end());
+
+    iter = SPIRVIterator(words, 0);
+  }
+  SPIRVOperation(spv::Op op, std::vector<uint32_t> data)
+  {
+    words.push_back(MakeHeader(op, data.size() + 1));
+    words.insert(words.begin() + 1, data.begin(), data.end());
 
     iter = SPIRVIterator(words, 0);
   }
@@ -100,6 +105,10 @@ public:
   }
 
 private:
+  friend class SPIRVEditor;
+
+  std::vector<uint32_t>::const_iterator begin() const { return iter.it(); }
+  std::vector<uint32_t>::const_iterator end() const { return iter.it() + size(); }
   inline uint32_t MakeHeader(spv::Op op, size_t WordCount)
   {
     return (uint32_t(op) & spv::OpCodeMask) | (uint16_t(WordCount) << spv::WordCountShift);
@@ -117,21 +126,151 @@ struct SPIRVEntry
 {
   uint32_t id;
   std::string name;
+  SPIRVIterator entryPoint;
+  SPIRVIterator function;
+  SPIRVIterator blocks;
+};
+
+struct SPIRVFunction
+{
+  uint32_t id;
   SPIRVIterator iter;
 };
 
+struct SPIRVScalar
+{
+  constexpr SPIRVScalar(spv::Op t, uint32_t w, bool s) : type(t), width(w), signedness(s) {}
+  SPIRVScalar(SPIRVIterator op);
+
+  spv::Op type;
+  uint32_t width;
+  bool signedness;
+
+  bool operator<(const SPIRVScalar &o) const
+  {
+    if(type != o.type)
+      return type < o.type;
+    if(signedness != o.signedness)
+      return signedness < o.signedness;
+    return width < o.width;
+  }
+
+  bool operator!=(const SPIRVScalar &o) const { return !operator==(o); }
+  bool operator==(const SPIRVScalar &o) const
+  {
+    return type == o.type && width == o.width && signedness == o.signedness;
+  }
+
+  inline SPIRVOperation decl() const
+  {
+    if(type == spv::Op::OpTypeBool)
+      return SPIRVOperation(type, {0});
+    else if(type == spv::Op::OpTypeFloat)
+      return SPIRVOperation(type, {0, width});
+    else if(type == spv::Op::OpTypeInt)
+      return SPIRVOperation(type, {0, width, signedness ? 1U : 0U});
+    else
+      return SPIRVOperation(spv::Op::OpNop, {0});
+  }
+};
+
+// helper to create SPIRVScalar objects for known types
+template <typename T>
+inline constexpr SPIRVScalar scalar();
+
+#define SCALAR_TYPE(ctype, op, width, sign)    \
+  template <>                                  \
+  inline constexpr SPIRVScalar scalar<ctype>() \
+  {                                            \
+    return SPIRVScalar(op, width, sign);       \
+  }
+
+SCALAR_TYPE(bool, spv::Op::OpTypeBool, 0, false);
+SCALAR_TYPE(uint16_t, spv::Op::OpTypeInt, 16, false);
+SCALAR_TYPE(uint32_t, spv::Op::OpTypeInt, 32, false);
+SCALAR_TYPE(int16_t, spv::Op::OpTypeInt, 16, true);
+SCALAR_TYPE(int32_t, spv::Op::OpTypeInt, 32, true);
+SCALAR_TYPE(float, spv::Op::OpTypeFloat, 32, false);
+SCALAR_TYPE(double, spv::Op::OpTypeFloat, 64, false);
+
+struct SPIRVVector
+{
+  SPIRVVector(const SPIRVScalar &s, uint32_t c) : scalar(s), count(c) {}
+  SPIRVScalar scalar;
+  uint32_t count;
+
+  bool operator<(const SPIRVVector &o) const
+  {
+    if(scalar != o.scalar)
+      return scalar < o.scalar;
+    return count < o.count;
+  }
+
+  bool operator!=(const SPIRVVector &o) const { return !operator==(o); }
+  bool operator==(const SPIRVVector &o) const { return scalar == o.scalar && count == o.count; }
+};
+
+struct SPIRVMatrix
+{
+  SPIRVMatrix(const SPIRVVector &v, uint32_t c) : vector(v), count(c) {}
+  SPIRVVector vector;
+  uint32_t count;
+
+  bool operator<(const SPIRVMatrix &o) const
+  {
+    if(vector != o.vector)
+      return vector < o.vector;
+    return count < o.count;
+  }
+
+  bool operator!=(const SPIRVMatrix &o) const { return !operator==(o); }
+  bool operator==(const SPIRVMatrix &o) const { return vector == o.vector && count == o.count; }
+};
+
+struct SPIRVPointer
+{
+  SPIRVPointer(uint32_t b, spv::StorageClass s) : baseId(b), storage(s) {}
+  uint32_t baseId;
+  spv::StorageClass storage;
+
+  bool operator<(const SPIRVPointer &o) const
+  {
+    if(baseId != o.baseId)
+      return baseId < o.baseId;
+    return storage < o.storage;
+  }
+
+  bool operator!=(const SPIRVPointer &o) const { return !operator==(o); }
+  bool operator==(const SPIRVPointer &o) const
+  {
+    return baseId == o.baseId && storage == o.storage;
+  }
+};
+
 class SPIRVEditor
 {
 public:
-  SPIRVEditor(std::vector<uint32_t> &spirvWords);
+  SPIRVEditor(const std::vector<uint32_t> &spirvWords);
+
+  // gets the modified SPIR-V with any modifications applied.
+  // This doesn't happen "live" because inserting any new ops invalidates all subsequent iterators
+  std::vector<uint32_t> GetWords();
 
-  std::vector<uint32_t> &GetWords() { return spirv; }
   uint32_t MakeId();
+
+  void SetName(uint32_t id, const char *name);
   void AddDecoration(const SPIRVOperation &op);
   void AddType(const SPIRVOperation &op);
   void AddVariable(const SPIRVOperation &op);
   void AddFunction(const SPIRVOperation *ops, size_t count);
 
+  // fetches the id of this type. If it exists already the old ID will be returned, otherwise it
+  // will be declared and the new ID returned
+  uint32_t DeclareType(const SPIRVScalar &scalar);
+  uint32_t DeclareType(const SPIRVVector &vector);
+  uint32_t DeclareType(const SPIRVMatrix &matrix);
+  uint32_t DeclareType(const SPIRVPointer &pointer);
+
   // simple properties that are public.
   struct
   {
@@ -142,18 +281,31 @@ public:
   spv::SourceLanguage sourceLang = spv::SourceLanguageUnknown;
   uint32_t sourceVer = 0;
 
+  // accessors to structs/vectors of data
   const std::vector<SPIRVEntry> &GetEntries() { return entries; }
+  const std::vector<SPIRVFunction> &GetFunctions() { return functions; }
 private:
   SPIRVIterator idBound;
 
-  struct
+  struct LogicalSection
   {
-    SPIRVIterator decoration;
-    SPIRVIterator typeVar;
-    SPIRVIterator function;
-  } globaliters;
+    SPIRVIterator iter;
+    std::vector<uint32_t> additions;
+  };
+
+  LogicalSection debugSection;
+  LogicalSection decorationSection;
+  LogicalSection typeVarSection;
+
+  std::vector<SPIRVIterator> ids;
 
   std::vector<SPIRVEntry> entries;
+  std::vector<SPIRVFunction> functions;
 
-  std::vector<uint32_t> &spirv;
+  std::map<SPIRVScalar, uint32_t> scalarTypes;
+  std::map<SPIRVVector, uint32_t> vectorTypes;
+  std::map<SPIRVMatrix, uint32_t> matrixTypes;
+  std::map<SPIRVPointer, uint32_t> pointerTypes;
+
+  std::vector<uint32_t> spirv;
 };