shader: Work around bug with struct uniforms used by more than 1 stage

panda/src/gobj/shader.cxx

@@ -1713,7 +1713,8 @@ do_load_source(ShaderModule::Stage stage, const std::string &source, BamCacheRec
 bool Shader::
 link() {
   // Go through all the modules to fetch the parameters.
-  pmap<CPT_InternalName, const ShaderModule::Variable *> parameters;
+  pmap<CPT_InternalName, const ShaderModule::Variable *> parameters_by_name;
+  pvector<const ShaderModule::Variable *> parameters;
   BitArray used_locations;
 
   for (COWPT(ShaderModule) &cow_module : _modules) {
@@ -1721,7 +1722,7 @@ link() {
     pmap<int, int> remap;
 
     for (const ShaderModule::Variable &var : module->_parameters) {
-      const auto result = parameters.insert({var.name, &var});
+      const auto result = parameters_by_name.insert({var.name, &var});
       const auto &it = result.first;
 
       if (!result.second) {
@@ -1730,16 +1731,23 @@ link() {
         const ShaderModule::Variable &other = *(it->second);
         if (other.type != var.type) {
           shader_cat.error()
-            << "Parameter " << var.name << " in module " << *module
+            << "Parameter " << *var.name << " in module " << *module
             << " is declared in another stage with a mismatching type!\n";
           return false;
         }
 
-        if (it->second->get_location() != var.get_location()) {
-          // Different location; need to remap this.
-          remap[var.get_location()] = it->second->get_location();
+        // Aggregate types don't seem to work properly when sharing uniforms
+        // between shader stages.  Needs revisiting.
+        if (!var.type->is_aggregate_type()) {
+          if (it->second->get_location() != var.get_location()) {
+            // Different location; need to remap this.
+            remap[var.get_location()] = it->second->get_location();
+          }
+          continue;
         }
-      } else if (var.has_location()) {
+      }
+
+      if (var.has_location()) {
         // Check whether the locations occupied by this variable are already in
         // use by another stage.
         int num_locations = var.type->get_num_parameter_locations();
@@ -1759,6 +1767,7 @@ link() {
           used_locations.set_range(var.get_location(), num_locations);
         }
       }
+      parameters.push_back(&var);
     }
 
     if (!remap.empty()) {
@@ -1781,10 +1790,8 @@ link() {
 
   // Now bind all of the parameters.
   bool success = true;
-  for (const auto &pair : parameters) {
-    const ShaderModule::Variable &var = *pair.second;
-
-    if (!bind_parameter(var.name, var.type, var.get_location())) {
+  for (const ShaderModule::Variable *var : parameters) {
+    if (!bind_parameter(var->name, var->type, var->get_location())) {
       success = false;
     }
   }

panda/src/gobj/shaderType.h

@@ -68,6 +68,7 @@ PUBLISHED:
   static const ShaderType::Sampler *sampler_type;
 
 public:
+  virtual bool is_aggregate_type() const { return false; }
   virtual bool as_scalar_type(ScalarType &type,
                               uint32_t &num_elements,
                               uint32_t &num_rows,
@@ -233,6 +234,7 @@ public:
 
   virtual int get_num_parameter_locations() const override;
 
+  bool is_aggregate_type() const override { return true; }
   const Struct *as_struct() const override { return this; }
 
 PUBLISHED:
@@ -278,6 +280,7 @@ public:
 
   virtual int get_num_parameter_locations() const override;
 
+  bool is_aggregate_type() const override { return true; }
   const Array *as_array() const override { return this; }
 
 PUBLISHED:

panda/src/shaderpipeline/shaderModuleSpirV.I

@@ -91,6 +91,34 @@ begin() {
   return iterator(&_words[5]);
 }
 
+/**
+ * Returns an iterator to the beginning of the annotations block.
+ */
+INLINE ShaderModuleSpirV::InstructionIterator ShaderModuleSpirV::InstructionStream::
+begin_annotations() {
+  iterator it;
+  for (it = begin(); it != end(); ++it) {
+    SpvOp opcode = (*it).opcode;
+    if (opcode != SpvOpNop &&
+        opcode != SpvOpCapability &&
+        opcode != SpvOpExtension &&
+        opcode != SpvOpExtInstImport &&
+        opcode != SpvOpMemoryModel &&
+        opcode != SpvOpEntryPoint &&
+        opcode != SpvOpExecutionMode &&
+        opcode != SpvOpString &&
+        opcode != SpvOpSourceExtension &&
+        opcode != SpvOpSource &&
+        opcode != SpvOpSourceContinued &&
+        opcode != SpvOpName &&
+        opcode != SpvOpMemberName &&
+        opcode != SpvOpModuleProcessed) {
+      break;
+    }
+  }
+  return it;
+}
+
 /**
  * Returns an iterator past the end of the instruction stream.
  */

panda/src/shaderpipeline/shaderModuleSpirV.cxx

@@ -43,7 +43,7 @@ ShaderModuleSpirV(Stage stage, const uint32_t *words, size_t size) :
   for (uint32_t id = 0; id < defs.size(); ++id) {
     Definition &def = defs[id];
     if (def._dtype == DT_type && def._name == "$Global") {
-      unwrap_uniform_block(defs, id);
+      flatten_struct(defs, id);
     }
   }
 
@@ -480,35 +480,9 @@ assign_locations(Definitions &defs) {
     return;
   }
 
-  // Find the end of the annotation block, so that we know where to insert the
-  // new locations.
-  InstructionIterator it;
-  for (it = _instructions.begin(); it != _instructions.end(); ++it) {
-    SpvOp opcode = (*it).opcode;
-    if (opcode != SpvOpNop &&
-        opcode != SpvOpCapability &&
-        opcode != SpvOpExtension &&
-        opcode != SpvOpExtInstImport &&
-        opcode != SpvOpMemoryModel &&
-        opcode != SpvOpEntryPoint &&
-        opcode != SpvOpExecutionMode &&
-        opcode != SpvOpString &&
-        opcode != SpvOpSourceExtension &&
-        opcode != SpvOpSource &&
-        opcode != SpvOpSourceContinued &&
-        opcode != SpvOpName &&
-        opcode != SpvOpMemberName &&
-        opcode != SpvOpModuleProcessed &&
-        opcode != SpvOpDecorate &&
-        opcode != SpvOpMemberDecorate &&
-        opcode != SpvOpGroupDecorate &&
-        opcode != SpvOpGroupMemberDecorate &&
-        opcode != SpvOpDecorationGroup) {
-      break;
-    }
-  }
-
-  // Now insert decorations for every unassigned variable.
+  // Insert decorations for every unassigned variable at the beginning of the
+  // annotations block.
+  InstructionIterator it = _instructions.begin_annotations();
   for (uint32_t id = 0; id < defs.size(); ++id) {
     Definition &def = defs[id];
     if (def._dtype == DT_variable &&
@@ -613,11 +587,11 @@ remap_locations(SpvStorageClass storage_class, const pmap<int, int> &locations)
 }
 
 /**
- * Converts the variables in the uniform block with the given ID to regular
+ * Converts the members of the struct type with the given ID to regular
  * variables.
  */
 void ShaderModuleSpirV::
-unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
+flatten_struct(Definitions &defs, uint32_t type_id) {
   const ShaderType::Struct *struct_type;
   DCAST_INTO_V(struct_type, defs[type_id]._type);
 
@@ -645,6 +619,7 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
     case SpvOpTypeStruct:
       // Delete the struct definition itself.
       if (op.nargs >= 1 && op.args[0] == type_id) {
+        defs[type_id].clear();
         it = _instructions.erase(it);
         continue;
       }
@@ -654,13 +629,7 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
       if (op.nargs >= 3 && op.args[2] == type_id) {
         // Remember this pointer.
         deleted_ids.insert(op.args[0]);
-
-        //if ((SpvStorageClass)op.args[1] == SpvStorageClassUniform) {
-        //  // Change storage class to UniformConstant.
-        //  op.args[1] = SpvStorageClassUniformConstant;
-        //  defs[op.args[0]]._storage_class = SpvStorageClassUniformConstant;
-        //}
-
+        defs[op.args[0]].clear();
         it = _instructions.erase(it);
         continue;
       }
@@ -670,9 +639,14 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
       if (op.nargs >= 3 && deleted_ids.count(op.args[0])) {
         // Delete this variable entirely, and replace it instead with individual
         // variable definitions for all its members.
-        deleted_ids.insert(op.args[1]);
+        uint32_t struct_var_id = op.args[1];
+        int struct_location = defs[struct_var_id]._location;
+        deleted_ids.insert(struct_var_id);
         it = _instructions.erase(it);
 
+        std::string struct_var_name = std::move(defs[struct_var_id]._name);
+        defs[struct_var_id].clear();
+
         for (size_t mi = 0; mi < struct_type->get_num_members(); ++mi) {
           const ShaderType::Struct::Member &member = struct_type->get_member(mi);
 
@@ -710,7 +684,16 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
           ++it;
 
           defs.push_back(Definition());
-          defs[variable_id]._name = member.name;
+          if (struct_var_name.empty()) {
+            defs[variable_id]._name = member.name;
+          } else {
+            defs[variable_id]._name = struct_var_name + "." + member.name;
+          }
+          if (struct_location >= 0) {
+            // Assign decorations to the individual members.
+            int location = struct_location + mi;
+            defs[variable_id]._location = location;
+          }
           defs[variable_id].set_variable(member.type, SpvStorageClassUniformConstant);
 
           member_ids[mi] = variable_id;
@@ -737,7 +720,8 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
       break;
 
     case SpvOpLoad:
-      // Shouldn't be loading the struct directly.
+      // If this triggers, the struct is being loaded into another variable,
+      // which means we can't unwrap this (for now).
       nassertv(!deleted_ids.count(op.args[2]));
 
       if (deleted_access_chains.count(op.args[2])) {
@@ -761,6 +745,16 @@ unwrap_uniform_block(Definitions &defs, uint32_t type_id) {
     ++it;
   }
 
+  // Insert decorations for the individual members.
+  it = _instructions.begin_annotations();
+  for (uint32_t var_id : member_ids) {
+    int location = defs[var_id]._location;
+    if (location >= 0) {
+      it = _instructions.insert(it,
+        SpvOpDecorate, {var_id, SpvDecorationLocation, (uint32_t)location});
+    }
+  }
+
   // Go over it again now that we know the deleted IDs, to remove any
   // decorations on them.
   if (deleted_ids.empty()) {
@@ -937,3 +931,17 @@ set_constant(const ShaderType *type, const uint32_t *words, uint32_t nwords) {
     _constant = 0;
   }
 }
+
+/**
+ * Clears this definition, in case it has just been removed.
+ */
+void ShaderModuleSpirV::Definition::
+clear() {
+  _dtype = DT_none;
+  _name.clear();
+  _type = nullptr;
+  _location = -1;
+  _builtin = SpvBuiltInMax;
+  _constant = 0;
+  _member_names.clear();
+}

panda/src/shaderpipeline/shaderModuleSpirV.h

@@ -77,6 +77,7 @@ protected:
     InstructionStream strip() const;
 
     INLINE iterator begin();
+    INLINE iterator begin_annotations();
     INLINE iterator end();
     INLINE iterator insert(iterator &it, SpvOp opcode, std::initializer_list<uint32_t > args);
     INLINE iterator insert(iterator &it, SpvOp opcode, const uint32_t *args, uint16_t nargs);
@@ -125,6 +126,8 @@ protected:
     void set_type_pointer(SpvStorageClass storage_class, const ShaderType *type);
     void set_variable(const ShaderType *type, SpvStorageClass storage_class);
     void set_constant(const ShaderType *type, const uint32_t *words, uint32_t nwords);
+
+    void clear();
   };
   typedef pvector<Definition> Definitions;
 
@@ -136,9 +139,11 @@ private:
   void assign_locations(Definitions &defs);
   void remap_locations(SpvStorageClass storage_class, const pmap<int, int> &locations);
 
-  void unwrap_uniform_block(Definitions &defs, uint32_t type_id);
+  void flatten_struct(Definitions &defs, uint32_t type_id);
   void strip();
 
+  int _index;
+
 public:
   static TypeHandle get_class_type() {
     return _type_handle;