shaderpipeline: Fix various issues with uniform block layouts

panda/src/gobj/shaderType.I

@@ -134,34 +134,6 @@ get_member(size_t i) const {
   return _members[i];
 }
 
-/**
- * Adds a member to this struct.
- */
-INLINE void ShaderType::Struct::
-add_member(const ShaderType *type, std::string name) {
-  Member member;
-  member.type = type;
-  member.name = std::move(name);
-  member.offset = _members.empty() ? 0 : _members.back().offset + _members.back().type->get_size_bytes();
-  _members.push_back(std::move(member));
-}
-
-/**
- * Adds a member to this struct with a given offset.
- */
-INLINE void ShaderType::Struct::
-add_member(const ShaderType *type, std::string name, uint32_t offset) {
-  pvector<Member>::iterator it = _members.begin();
-  while (it != _members.end() && it->offset < offset) {
-    ++it;
-  }
-  Member member;
-  member.type = type;
-  member.name = std::move(name);
-  member.offset = offset;
-  _members.insert(it, std::move(member));
-}
-
 /**
  * Constructs an array type from a base type and number of elements.
  */

panda/src/gobj/shaderType.cxx

@@ -141,6 +141,14 @@ compare_to_impl(const ShaderType &other) const {
        - (_scalar_type < other_scalar._scalar_type);
 }
 
+/**
+ * Returns the alignment in bytes of this type in memory, if applicable.
+ */
+int ShaderType::Scalar::
+get_align_bytes() const {
+  return (_scalar_type == ST_double) ? 8 : 4;
+}
+
 /**
  * Returns true if this type contains the given scalar type.
  */
@@ -185,6 +193,15 @@ compare_to_impl(const ShaderType &other) const {
        - (_num_components < other_vector._num_components);
 }
 
+/**
+ * Returns the alignment in bytes of this type in memory, if applicable.
+ */
+int ShaderType::Vector::
+get_align_bytes() const {
+  int component_align = (_scalar_type == ST_double) ? 8 : 4;
+  return component_align * ((_num_components == 3) ? 4 : _num_components);
+}
+
 /**
  * Returns true if this type contains the given scalar type.
  */
@@ -232,6 +249,16 @@ compare_to_impl(const ShaderType &other) const {
        - (_num_columns < other_matrix._num_columns);
 }
 
+/**
+ * Returns the alignment in bytes of this type in memory, if applicable.
+ */
+int ShaderType::Matrix::
+get_align_bytes() const {
+  //TODO: needs to be checked
+  int row_align = (_scalar_type == ST_double) ? 32 : 16;
+  return row_align * _num_rows;
+}
+
 /**
  * Returns the number of in/out locations taken up by in/out variables having
  * this type.
@@ -241,6 +268,36 @@ get_num_interface_locations() const {
   return _num_rows;
 }
 
+/**
+ * Adds a member to this struct.
+ */
+void ShaderType::Struct::
+add_member(const ShaderType *type, std::string name) {
+  Member member;
+  member.type = type;
+  member.name = std::move(name);
+  member.offset = _members.empty() ? 0 : _members.back().offset + _members.back().type->get_size_bytes();
+  int alignment = type->get_align_bytes();
+  member.offset += alignment - ((member.offset + (alignment - 1)) % alignment) - 1;
+  _members.push_back(std::move(member));
+}
+
+/**
+ * Adds a member to this struct with a given offset.
+ */
+void ShaderType::Struct::
+add_member(const ShaderType *type, std::string name, uint32_t offset) {
+  pvector<Member>::iterator it = _members.begin();
+  while (it != _members.end() && it->offset < offset) {
+    ++it;
+  }
+  Member member;
+  member.type = type;
+  member.name = std::move(name);
+  member.offset = offset;
+  _members.insert(it, std::move(member));
+}
+
 /**
  * Returns true if this type contains the given scalar type.
  */
@@ -295,6 +352,18 @@ compare_to_impl(const ShaderType &other) const {
   return 0;
 }
 
+/**
+ * Returns the alignment in bytes of this type in memory, if applicable.
+ */
+int ShaderType::Struct::
+get_align_bytes() const {
+  int align = 16;
+  for (const Member &member : _members) {
+    align = std::max(align, member.type->get_align_bytes());
+  }
+  return align;
+}
+
 /**
  * Returns the size in bytes of this type in memory, if applicable.  Opaque
  * types will return -1.
@@ -376,13 +445,30 @@ compare_to_impl(const ShaderType &other) const {
        - (_num_elements < other_array._num_elements);
 }
 
+/**
+ * Returns the array stride in bytes.
+ */
+int ShaderType::Array::
+get_stride_bytes() const {
+  int element_size = _element_type->get_size_bytes();
+  return (element_size + 15) & ~15;
+}
+
+/**
+ * Returns the alignment in bytes of this type in memory, if applicable.
+ */
+int ShaderType::Array::
+get_align_bytes() const {
+  return get_stride_bytes();
+}
+
 /**
  * Returns the size in bytes of this type in memory, if applicable.  Opaque
  * types will return -1.
  */
 int ShaderType::Array::
 get_size_bytes() const {
-  return _element_type->get_size_bytes() * _num_elements;
+  return get_stride_bytes() * _num_elements;
 }
 
 /**

panda/src/gobj/shaderType.h

@@ -34,6 +34,7 @@ public:
 
   virtual void output(std::ostream &out) const=0;
 
+  virtual int get_align_bytes() const { return 1; }
   virtual int get_size_bytes() const;
   virtual int get_num_interface_locations() const { return 1; }
   virtual int get_num_parameter_locations() const { return 1; }
@@ -129,6 +130,8 @@ public:
 private:
   virtual int compare_to_impl(const ShaderType &other) const override;
 
+  virtual int get_align_bytes() const override;
+
   const ScalarType _scalar_type;
 
 public:
@@ -167,6 +170,8 @@ public:
 private:
   virtual int compare_to_impl(const ShaderType &other) const override;
 
+  virtual int get_align_bytes() const override;
+
   const ScalarType _scalar_type;
   const uint32_t _num_components;
 
@@ -208,6 +213,8 @@ public:
 private:
   virtual int compare_to_impl(const ShaderType &other) const override;
 
+  virtual int get_align_bytes() const override;
+
   const ScalarType _scalar_type;
   const uint32_t _num_rows;
   const uint32_t _num_columns;
@@ -235,12 +242,13 @@ public:
 
   INLINE size_t get_num_members() const;
   INLINE const Member &get_member(size_t i) const;
-  INLINE void add_member(const ShaderType *type, std::string name);
-  INLINE void add_member(const ShaderType *type, std::string name, uint32_t offset);
+  void add_member(const ShaderType *type, std::string name);
+  void add_member(const ShaderType *type, std::string name, uint32_t offset);
 
   virtual void output(std::ostream &out) const override;
   virtual int compare_to_impl(const ShaderType &other) const override;
 
+  virtual int get_align_bytes() const override;
   virtual int get_size_bytes() const override;
   virtual int get_num_interface_locations() const override;
   virtual int get_num_parameter_locations() const override;
@@ -292,6 +300,8 @@ public:
   virtual void output(std::ostream &out) const override;
   virtual int compare_to_impl(const ShaderType &other) const override;
 
+  int get_stride_bytes() const;
+  virtual int get_align_bytes() const override;
   virtual int get_size_bytes() const override;
   virtual int get_num_interface_locations() const override;
   virtual int get_num_parameter_locations() const override;

panda/src/shaderpipeline/shaderModuleSpirV.cxx

@@ -1344,6 +1344,17 @@ r_annotate_struct_layout(InstructionIterator &it, uint32_t type_id) {
     const ShaderType *base_type = member.type;
     while (const ShaderType::Array *array_type = base_type->as_array()) {
       base_type = array_type->get_element_type();
+
+      // Also make sure there's an ArrayStride decoration for this array.
+      uint32_t array_type_id = _type_map[array_type];
+      Definition &array_def = _defs[array_type_id];
+
+      if (array_def._array_stride == 0) {
+        array_def._array_stride = array_type->get_stride_bytes();
+        it = _instructions.insert(it, spv::OpDecorate,
+          {array_type_id, spv::DecorationArrayStride, array_def._array_stride});
+        ++it;
+      }
     }
 
     if (const ShaderType::Matrix *matrix_type = base_type->as_matrix()) {
@@ -1617,6 +1628,10 @@ parse_instruction(const Instruction &op) {
       _defs[op.args[0]]._location = op.args[2];
       break;
 
+    case spv::DecorationArrayStride:
+      _defs[op.args[0]]._array_stride = op.args[2];
+      break;
+
     default:
       break;
     }

panda/src/shaderpipeline/shaderModuleSpirV.h

@@ -147,6 +147,7 @@ public:
     spv::BuiltIn _builtin = spv::BuiltInMax;
     uint32_t _constant = 0;
     uint32_t _type_id = 0;
+    uint32_t _array_stride = 0;
     MemberDefinitions _members;
     bool _used = false;