vulkan: Support desktop GLSL shaders, mat_spec shader inputs via UBOs

panda/src/shaderpipeline/shaderCompilerGlslang.cxx

@@ -379,7 +379,7 @@ compile_now(ShaderModule::Stage stage, std::istream &in,
       );
     }
   }
-  shader.setEnvClient(glslang::EShClient::EShClientVulkan, glslang::EShTargetVulkan_1_0);
+  shader.setEnvClient(glslang::EShClient::EShClientOpenGL, glslang::EShTargetOpenGL_450);
   shader.setEnvTarget(glslang::EShTargetSpv, glslang::EShTargetSpv_1_0);
 
   // This will squelch the warnings about missing bindings and locations, since

panda/src/vulkandisplay/config_vulkandisplay.cxx

@@ -48,6 +48,15 @@ ConfigVariableInt64 vulkan_memory_page_size
           "several pages will be allocated at application start due to the "
           "need for pages with varying memory types."));
 
+ConfigVariableInt64 vulkan_global_uniform_buffer_size
+("vulkan-global-uniform-buffer-size", 64 * 1024,
+ PRC_DESC("This value indicates how large the uniform buffer should be that is "
+          "allocated to contain all of the global uniforms values used by all "
+          "shaders in a give frame.  To optimize this value, enable "
+          "\"notify-level-vulkandisplay debug\" and look for the highest value "
+          "in the message \"Used at most # bytes of global uniform buffer.\" "
+          "in the most complex scene, then add a generous safety margin."));
+
 #define VK_ERROR_INVALID_SHADER_NV -1000012000
 
 /**

panda/src/vulkandisplay/config_vulkandisplay.h

@@ -23,6 +23,7 @@ NotifyCategoryDecl(vulkandisplay, EXPCL_VULKANDISPLAY, EXPTP_VULKANDISPLAY);
 
 extern ConfigVariableInt vulkan_color_palette_size;
 extern ConfigVariableInt64 vulkan_memory_page_size;
+extern ConfigVariableInt64 vulkan_global_uniform_buffer_size;
 
 extern EXPCL_VULKANDISPLAY void init_libvulkandisplay();
 extern "C" EXPCL_VULKANDISPLAY int get_pipe_type_p3vulkandisplay();

panda/src/vulkandisplay/vulkanGraphicsStateGuardian.I

@@ -11,44 +11,6 @@
  * @date 2016-02-16
  */
 
-/**
- * Copy constructor for DescriptorSetKey.
- */
-INLINE VulkanGraphicsStateGuardian::DescriptorSetKey::
-DescriptorSetKey(const DescriptorSetKey &copy) :
-  _tex_attrib(copy._tex_attrib),
-  _shader_attrib(copy._shader_attrib) {
-}
-
-/**
- * Move constructor for DescriptorSetKey, defined for efficiency (so that we
- * don't have to increase and decrease the reference count when moving this
- * object).
- */
-INLINE VulkanGraphicsStateGuardian::DescriptorSetKey::
-DescriptorSetKey(DescriptorSetKey &&from) noexcept :
-  _tex_attrib(std::move(from._tex_attrib)),
-  _shader_attrib(std::move(from._shader_attrib)) {
-}
-
-/**
- * Copy assignment operator for DescriptorSetKey.
- */
-INLINE void VulkanGraphicsStateGuardian::DescriptorSetKey::
-operator = (const DescriptorSetKey &copy) {
-  _tex_attrib = copy._tex_attrib;
-  _shader_attrib = copy._shader_attrib;
-}
-
-/**
- * Move assignment operator for DescriptorSetKey.
- */
-INLINE void VulkanGraphicsStateGuardian::DescriptorSetKey::
-operator = (DescriptorSetKey &&from) noexcept {
-  _tex_attrib = std::move(from._tex_attrib);
-  _shader_attrib = std::move(from._shader_attrib);
-}
-
 /**
  * Returns true if these two DescriptorSetKey objects are identical.
  */

panda/src/vulkandisplay/vulkanGraphicsStateGuardian.cxx

@@ -30,16 +30,14 @@
 #include "transparencyAttrib.h"
 
 static const std::string default_vshader =
-  "#version 430\n"
+  "#version 330\n"
   "in vec4 p3d_Vertex;\n"
   "in vec4 p3d_Color;\n"
   "in vec2 p3d_MultiTexCoord0;\n"
   "out vec2 texcoord;\n"
   "out vec4 color;\n"
-  "layout(push_constant, std140) uniform p3d_PushConstants {\n"
-  "  uniform mat4 p3d_ModelViewProjectionMatrix;\n"
-  "  uniform vec4 p3d_ColorScale;\n"
-  "};\n"
+  "uniform mat4 p3d_ModelViewProjectionMatrix;\n"
+  "uniform vec4 p3d_ColorScale;\n"
   "void main(void) {\n"
   "  gl_Position = p3d_ModelViewProjectionMatrix * p3d_Vertex;\n"
   "  texcoord = p3d_MultiTexCoord0;\n"
@@ -47,7 +45,7 @@ static const std::string default_vshader =
   "}\n";
 
 static const std::string default_fshader =
-  "#version 430\n"
+  "#version 330\n"
   "in vec2 texcoord;\n"
   "in vec4 color;\n"
   "out vec4 p3d_FragColor;\n"
@@ -84,6 +82,9 @@ VulkanGraphicsStateGuardian(GraphicsEngine *engine, VulkanGraphicsPipe *pipe,
   _default_sc(nullptr),
   _total_allocated(0)
 {
+  const VkPhysicalDeviceLimits &limits = pipe->_gpu_properties.limits;
+  const VkPhysicalDeviceFeatures &features = pipe->_gpu_features;
+
   const char *const layers[] = {
     "VK_LAYER_LUNARG_standard_validation",
   };
@@ -231,6 +232,26 @@ VulkanGraphicsStateGuardian(GraphicsEngine *engine, VulkanGraphicsPipe *pipe,
   _color_palette[LColorf(1, 1, 1, 1)] = 1;
   _next_palette_index = 2;
 
+  // Create a push constant layout based on the available space.
+  {
+    ShaderType::Struct struct_type;
+    struct_type.add_member(ShaderType::register_type(ShaderType::Matrix(ShaderType::ST_float, 4, 4)), "p3d_ModelViewProjectionMatrix");
+    struct_type.add_member(ShaderType::register_type(ShaderType::Vector(ShaderType::ST_float, 4)), "p3d_ColorScale");
+    _push_constant_block_type = ShaderType::register_type(std::move(struct_type));
+  }
+
+  // Create a uniform buffer that we'll use for everything.
+  VkDeviceSize uniform_buffer_size = vulkan_global_uniform_buffer_size;
+  if (!create_buffer(uniform_buffer_size, _uniform_buffer, _uniform_buffer_memory,
+                     VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
+                     VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)) {
+    vulkandisplay_cat.error()
+      << "Failed to create uniform buffer buffer.\n";
+    return;
+  }
+  _uniform_buffer_size = uniform_buffer_size;
+  _uniform_buffer_offset_alignment = limits.minUniformBufferOffsetAlignment;
+
   // Load the default shader.  Temporary hack.
   static PT(Shader) default_shader;
   if (default_shader.is_null()) {
@@ -244,8 +265,6 @@ VulkanGraphicsStateGuardian(GraphicsEngine *engine, VulkanGraphicsPipe *pipe,
   }
 
   // Fill in the features supported by this physical device.
-  const VkPhysicalDeviceLimits &limits = pipe->_gpu_properties.limits;
-  const VkPhysicalDeviceFeatures &features = pipe->_gpu_features;
   _is_hardware = (pipe->_gpu_properties.deviceType != VK_PHYSICAL_DEVICE_TYPE_CPU);
 
   _max_vertices_per_array = std::max((uint32_t)0x7fffffff, limits.maxDrawIndexedIndexValue);
@@ -376,6 +395,7 @@ VulkanGraphicsStateGuardian::
   }
 
   // Remove the things we created in the constructor, in reverse order.
+  vkDestroyBuffer(_device, _uniform_buffer, nullptr);
   vkDestroyBuffer(_device, _color_vertex_buffer, nullptr);
   vkDestroyDescriptorPool(_device, _descriptor_pool, nullptr);
   vkDestroyDescriptorSetLayout(_device, _descriptor_set_layout, nullptr);
@@ -387,6 +407,7 @@ VulkanGraphicsStateGuardian::
   _memory_pages.clear();
 
   vkDestroyDevice(_device, nullptr);
+  _device = VK_NULL_HANDLE;
 }
 
 /**
@@ -1354,32 +1375,60 @@ prepare_shader(Shader *shader) {
 
   VulkanShaderContext *sc = new VulkanShaderContext(shader);
 
-  for (COWPT(ShaderModule) &cow_module : shader->_modules) {
-    CPT(ShaderModule) module = cow_module.get_read_pointer();
-    const ShaderModuleSpirV *spv_module = DCAST(ShaderModuleSpirV, module.p());
-    nassertd(spv_module != nullptr) continue;
-
-    VkShaderModuleCreateInfo module_info;
-    module_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
-    module_info.pNext = nullptr;
-    module_info.flags = 0;
-    module_info.codeSize = spv_module->get_data_size() * 4;
-    module_info.pCode = (const uint32_t *)spv_module->get_data();
-
-    VkResult err;
-    err = vkCreateShaderModule(_device, &module_info, nullptr, &sc->_modules[(size_t)spv_module->get_stage()]);
-    if (err) {
-      vulkan_error(err, "Failed to load shader module");
-      delete sc;
-      return nullptr;
-    }
+  if (!sc->create_modules(_device, _push_constant_block_type) ||
+      !sc->make_pipeline_layout(_device)) {
+    delete sc;
+    return nullptr;
   }
 
-  if (!sc->make_pipeline_layout(_device)) {
+  VkResult err;
+
+  // Create a descriptor set for the UBOs.
+  //TODO: actually figure out where this should live.
+  VkDescriptorSetAllocateInfo alloc_info;
+  alloc_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
+  alloc_info.pNext = nullptr;
+  alloc_info.descriptorPool = _descriptor_pool;
+  alloc_info.descriptorSetCount = 1;
+  alloc_info.pSetLayouts = &sc->_descriptor_set_layouts[1];
+  err = vkAllocateDescriptorSets(_device, &alloc_info, &sc->_uniform_descriptor_set);
+  if (err) {
+    vulkan_error(err, "Failed to allocate descriptor set");
     delete sc;
     return nullptr;
   }
 
+  // We set the offsets to 0, since we use dynamic offsets.
+  size_t count = 0;
+  VkDescriptorBufferInfo buffer_info[2];
+  if (sc->_mat_block_size > 0) {
+    buffer_info[count].buffer = _uniform_buffer;
+    buffer_info[count].offset = 0;
+    buffer_info[count].range = sc->_mat_block_size;
+    ++count;
+  }
+  if (sc->_ptr_block_size > 0) {
+    buffer_info[count].buffer = _uniform_buffer;
+    buffer_info[count].offset = 0;
+    buffer_info[count].range = sc->_ptr_block_size;
+    ++count;
+  }
+
+  VkWriteDescriptorSet write[2];
+  for (size_t i = 0; i < count; ++i) {
+    write[i].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+    write[i].pNext = nullptr;
+    write[i].dstSet = sc->_uniform_descriptor_set;
+    write[i].dstBinding = i;
+    write[i].dstArrayElement = 0;
+    write[i].descriptorCount = 1;
+    write[i].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
+    write[i].pImageInfo = nullptr;
+    write[i].pBufferInfo = &buffer_info[i];
+    write[i].pTexelBufferView = nullptr;
+  }
+  vkUpdateDescriptorSets(_device, count, write, 0, nullptr);
+
   return sc;
 }
 
@@ -1413,7 +1462,8 @@ release_shader(ShaderContext *context) {
   }
 
   vkDestroyPipelineLayout(_device, sc->_pipeline_layout, nullptr);
-  vkDestroyDescriptorSetLayout(_device, sc->_descriptor_set_layout, nullptr);
+  vkDestroyDescriptorSetLayout(_device, sc->_descriptor_set_layouts[0], nullptr);
+  vkDestroyDescriptorSetLayout(_device, sc->_descriptor_set_layouts[1], nullptr);
 
   delete sc;
 }
@@ -1678,26 +1728,36 @@ set_state_and_transform(const RenderState *state,
   // This does not actually set the state just yet, because we can't make a
   // pipeline state without knowing the vertex format.
   _state_rs = state;
+  _target_rs = state;
+  _internal_transform = trans;
 
   VulkanShaderContext *sc = _default_sc;
   const ShaderAttrib *sa;
   if (state->get_attrib(sa) && sa->has_shader()) {
     Shader *shader = (Shader *)sa->get_shader();
+    nassertv(shader != nullptr);
     DCAST_INTO_V(sc, shader->prepare_now(get_prepared_objects(), this));
   }
 
   _current_shader = sc;
   nassertv(sc != nullptr);
 
-  // Put the modelview projection matrix in the push constants.
-  CPT(TransformState) combined = _projection_mat->compose(trans);
-  LMatrix4f matrix = LCAST(float, combined->get_mat());
-  vkCmdPushConstants(_cmd, sc->_pipeline_layout, VK_SHADER_STAGE_VERTEX_BIT, 0, 64, matrix.get_data());
+  // Put the modelview projection matrix and color scale in the push constants.
+  if (sc->_projection_mat_stage_mask != 0) {
+    CPT(TransformState) combined = _projection_mat->compose(trans);
+    LMatrix4f matrix = LCAST(float, combined->get_mat());
+    vkCmdPushConstants(_cmd, sc->_pipeline_layout, sc->_push_constant_stage_mask, 0, 64, matrix.get_data());
+  }
 
-  const ColorScaleAttrib *color_scale_attrib;
-  state->get_attrib_def(color_scale_attrib);
-  LColorf color = LCAST(float, color_scale_attrib->get_scale());
-  vkCmdPushConstants(_cmd, sc->_pipeline_layout, VK_SHADER_STAGE_VERTEX_BIT, 64, 16, color.get_data());
+  if (sc->_color_scale_stage_mask != 0) {
+    const ColorScaleAttrib *color_scale_attrib;
+    state->get_attrib_def(color_scale_attrib);
+    LColorf color = LCAST(float, color_scale_attrib->get_scale());
+    vkCmdPushConstants(_cmd, sc->_pipeline_layout, sc->_push_constant_stage_mask, 64, 16, color.get_data());
+  }
+
+  //TODO: properly compute altered field.
+  sc->update_uniform_buffers(this, _cmd, ~0);
 
   const TextureAttrib *tex_attrib;
   state->get_attrib_def(tex_attrib);
@@ -1909,6 +1969,20 @@ begin_frame(Thread *current_thread) {
   _pending_delete_buffers.clear();
   _pending_free.clear();
 
+  // Recycle the global uniform buffer.  It's probably not worth it to expect
+  // values to stick around between frames, because the vast majority of global
+  // uniforms will change more frequently than that anyway.
+  if (vulkandisplay_cat.is_debug()) {
+    static VkDeviceSize max_used = 0;
+    if (_uniform_buffer_offset > max_used) {
+      max_used = _uniform_buffer_offset;
+      vulkandisplay_cat.debug()
+        << "Used at most " << max_used << " of " << _uniform_buffer_size
+        << " bytes of global uniform buffer.\n";
+    }
+  }
+  _uniform_buffer_offset = 0;
+
   // Reset the fence to unsignaled status.
   err = vkResetFences(_device, 1, &_fence);
   nassertr(!err, false);
@@ -2713,7 +2787,7 @@ make_pipeline(VulkanShaderContext *sc, const RenderState *state,
     int array_index;
     const GeomVertexColumn *column;
 
-    attrib_desc[i].location = spec._id._seqno;
+    attrib_desc[i].location = spec._id._location;
 
     if (!format->get_array_info(spec._name, array_index, column) ||
         (spec._name == InternalName::get_color() &&
@@ -2732,25 +2806,7 @@ make_pipeline(VulkanShaderContext *sc, const RenderState *state,
           attrib_desc[i].offset = 16;
         } else {
           LColorf color = LCAST(float, color_attr->get_color());
-
-          ColorPaletteIndices::const_iterator it = _color_palette.find(color);
-          if (it != _color_palette.end()) {
-            attrib_desc[i].offset = it->second * 16;
-          } else {
-            // Not yet in the palette.  Write an entry.
-            if (_next_palette_index >= vulkan_color_palette_size) {
-              attrib_desc[i].offset = 1;
-              vulkandisplay_cat.error()
-                << "Ran out of color palette entries.  Increase "
-                   "vulkan-color-palette-size value in Config.prc.\n";
-            } else {
-              uint32_t offset = _next_palette_index * 16;
-              attrib_desc[i].offset = offset;
-              _color_palette[color] = _next_palette_index++;
-              vkCmdUpdateBuffer(_transfer_cmd, _color_vertex_buffer, offset, 16,
-                                (const uint32_t *)color.get_data());
-            }
-          }
+          attrib_desc[i].offset = get_color_palette_offset(color);
         }
       } else {
         attrib_desc[i].offset = 0;
@@ -3211,6 +3267,56 @@ make_descriptor_set(const RenderState *state) {
   return ds;
 }
 
+/**
+ * Reserves space in the global uniform buffer.
+ */
+VkDeviceSize VulkanGraphicsStateGuardian::
+update_dynamic_uniform_buffer(void *data, VkDeviceSize size) {
+  VkDeviceSize offset = _uniform_buffer_offset;
+  VkDeviceSize align = _uniform_buffer_offset_alignment;
+  offset = offset - 1 - (offset - 1) % align + align;
+
+  //TODO: fail more gracefully.  Create a new buffer on the fly and manage
+  // multiple buffers?  Or submit work and insert a fence, then replace the
+  // buffer with a new one?
+  if (offset + size > _uniform_buffer_size) {
+    vulkandisplay_cat.error()
+      << "Ran out of space in the global uniform buffer.  Increase "
+         "vulkan-global-uniform-buffer-size in Config.prc.\n";
+    abort();
+    return offset;
+  }
+
+  vkCmdUpdateBuffer(_transfer_cmd, _uniform_buffer, offset, size, data);
+  _uniform_buffer_offset = offset + size;
+  return offset;
+}
+
+/**
+ * Returns the offset of the given color in the color palette.
+ */
+uint32_t VulkanGraphicsStateGuardian::
+get_color_palette_offset(const LColor &color) {
+  ColorPaletteIndices::const_iterator it = _color_palette.find(color);
+  if (it != _color_palette.end()) {
+    return it->second * 16;
+  }
+
+  // Not yet in the palette.  Write an entry.
+  if (_next_palette_index >= vulkan_color_palette_size) {
+    vulkandisplay_cat.error()
+      << "Ran out of color palette entries.  Increase "
+         "vulkan-color-palette-size value in Config.prc.\n";
+    return 1;
+  }
+
+  uint32_t offset = _next_palette_index * 16;
+  _color_palette[color] = _next_palette_index++;
+  vkCmdUpdateBuffer(_transfer_cmd, _color_vertex_buffer, offset, 16,
+                    (const uint32_t *)color.get_data());
+  return offset;
+}
+
 /**
  * Returns a VkFormat suitable for the given texture.
  */

panda/src/vulkandisplay/vulkanGraphicsStateGuardian.h

@@ -145,28 +145,12 @@ public:
                            VkPrimitiveTopology topology);
   VkPipeline make_compute_pipeline(VulkanShaderContext *sc);
 
-  /**
-   * Stores whatever is used to key a cached descriptor set into the
-   * descriptor set map.
-   */
-  struct DescriptorSetKey {
-    INLINE DescriptorSetKey() = default;
-    INLINE DescriptorSetKey(const DescriptorSetKey &copy);
-    INLINE DescriptorSetKey(DescriptorSetKey &&from) noexcept;
-
-    INLINE void operator = (const DescriptorSetKey &copy);
-    INLINE void operator = (DescriptorSetKey &&from) noexcept;
-
-    INLINE bool operator ==(const DescriptorSetKey &other) const;
-    INLINE bool operator < (const DescriptorSetKey &other) const;
-
-    CPT(TextureAttrib) _tex_attrib;
-    CPT(ShaderAttrib) _shader_attrib;
-  };
-
   VkDescriptorSet get_descriptor_set(const RenderState *state);
   VkDescriptorSet make_descriptor_set(const RenderState *state);
 
+  VkDeviceSize update_dynamic_uniform_buffer(void *data, VkDeviceSize size);
+  uint32_t get_color_palette_offset(const LColor &color);
+
   VkFormat get_image_format(const Texture *texture) const;
   static bool lookup_image_format(VkFormat vk_format, Texture::Format &format,
                                   Texture::ComponentType &type);
@@ -189,9 +173,18 @@ private:
   VkDescriptorPool _descriptor_pool;
   VulkanShaderContext *_default_sc;
   VulkanShaderContext *_current_shader;
+  const ShaderType::Struct *_push_constant_block_type = nullptr;
   CPT(GeomVertexFormat) _format;
   PT(Texture) _white_texture;
 
+  // Single large uniform buffer used for everything in a frame.
+  VkBuffer _uniform_buffer;
+  VulkanMemoryBlock _uniform_buffer_memory;
+  VkDeviceSize _uniform_buffer_size = 0;
+  VkDeviceSize _uniform_buffer_offset = 0;
+  VkDeviceSize _uniform_buffer_offset_alignment;
+  VkDescriptorSet _uniform_descriptor_set;
+
   // Stores current framebuffer info.
   VkRenderPass _render_pass;
   VulkanTextureContext *_fb_color_tc;
@@ -209,6 +202,19 @@ private:
   typedef pmap<LColorf, uint32_t> ColorPaletteIndices;
   ColorPaletteIndices _color_palette;
 
+  /**
+   * Stores whatever is used to key a cached descriptor set into the
+   * descriptor set map.
+   */
+  struct DescriptorSetKey {
+    INLINE bool operator ==(const DescriptorSetKey &other) const;
+    INLINE bool operator < (const DescriptorSetKey &other) const;
+
+    CPT(TextureAttrib) _tex_attrib;
+    CPT(ShaderAttrib) _shader_attrib;
+  };
+
+  //TODO: we need a garbage collection mechanism for this.
   typedef pmap<DescriptorSetKey, VkDescriptorSet> DescriptorSetMap;
   DescriptorSetMap _descriptor_set_map;
 

panda/src/vulkandisplay/vulkanGraphicsWindow.cxx

@@ -394,6 +394,7 @@ close_window() {
 
     // Wait until the queue is done with any commands that might use the swap
     // chain, then destroy it.
+    nassertv(vkgsg->_device != VK_NULL_HANDLE);
     vkQueueWaitIdle(vkgsg->_queue);
     destroy_swapchain();
 

panda/src/vulkandisplay/vulkanShaderContext.I

@@ -20,6 +20,16 @@ INLINE VulkanShaderContext::
 VulkanShaderContext(Shader *shader) :
   ShaderContext(shader),
   _modules{VK_NULL_HANDLE} {
+
+  _mat_part_cache = new LMatrix4[_shader->cp_get_mat_cache_size()];
+}
+
+/**
+ *
+ */
+INLINE VulkanShaderContext::
+~VulkanShaderContext()  {
+  delete _mat_part_cache;
 }
 
 /**

panda/src/vulkandisplay/vulkanShaderContext.cxx

@@ -15,55 +15,234 @@
 
 TypeHandle VulkanShaderContext::_type_handle;
 
+/**
+ * Creates the shader modules.
+ */
+bool VulkanShaderContext::
+create_modules(VkDevice device, const ShaderType::Struct *push_constant_block_type) {
+  nassertr(push_constant_block_type != nullptr, false);
+
+  _push_constant_block_type = push_constant_block_type;
+
+  // Compose a struct type for all the mat inputs, also gathering ones that
+  // should go into a separate push constant block.  This will become a new
+  // uniform block in the shader that replaces the regular uniforms.
+  pvector<int> mat_struct_locations;
+  pvector<int> push_constant_locations = {-1, -1};
+
+  if (!_shader->_mat_spec.empty()) {
+    ShaderType::Struct struct_type;
+
+    for (const Shader::ShaderMatSpec &spec : _shader->_mat_spec) {
+      if (spec._id._location >= 0) {
+        if (spec._func == Shader::SMF_compose &&
+            spec._piece == Shader::SMP_whole &&
+            spec._part[0] == Shader::SMO_model_to_apiview &&
+            spec._part[1] == Shader::SMO_apiview_to_apiclip) {
+          // This is p3d_ModelViewProjectionMatrix, which is a push constant.
+          push_constant_locations[0] = spec._id._location;
+          _projection_mat_stage_mask |= spec._id._stage_mask;
+          _push_constant_stage_mask |= spec._id._stage_mask;
+        }
+        else if (spec._func == Shader::SMF_first &&
+                 (spec._piece == Shader::SMP_row3 || spec._piece == Shader::SMP_row3x3) &&
+                 spec._part[0] == Shader::SMO_attr_colorscale) {
+          // This is p3d_ColorScale or equivalent, which is a push constant.
+          push_constant_locations[1] = spec._id._location;
+          _color_scale_stage_mask |= spec._id._stage_mask;
+          _push_constant_stage_mask |= spec._id._stage_mask;
+        }
+        else {
+          // Other inputs are done via UBOs.
+          struct_type.add_member(spec._id._type, spec._id._name->get_name());
+          mat_struct_locations.push_back(spec._id._location);
+          _mat_deps |= spec._dep;
+          _mat_block_stage_mask |= spec._id._stage_mask;
+          _mat_spec.push_back(spec);
+        }
+      }
+    }
+
+    if (struct_type.get_num_members() > 0) {
+      _mat_block_type = ShaderType::register_type(std::move(struct_type));
+      _mat_block_size = _mat_block_type->get_size_bytes();
+    }
+  }
+
+  // Compose a struct type for all the ptr inputs.
+  pvector<int> ptr_struct_locations;
+  if (!_shader->_ptr_spec.empty()) {
+    ShaderType::Struct struct_type;
+
+    for (const Shader::ShaderPtrSpec &spec : _shader->_ptr_spec) {
+      if (spec._id._location >= 0) {
+        struct_type.add_member(spec._id._type, spec._id._name->get_name());
+        ptr_struct_locations.push_back(spec._id._location);
+        _ptr_block_stage_mask |= spec._id._stage_mask;
+      }
+    }
+
+    if (struct_type.get_num_members() > 0) {
+      _ptr_block_type = ShaderType::register_type(std::move(struct_type));
+      _ptr_block_size = _ptr_block_type->get_size_bytes();
+    }
+  }
+
+  for (COWPT(ShaderModule) &cow_module : _shader->_modules) {
+    CPT(ShaderModule) module = cow_module.get_read_pointer();
+
+    const ShaderModuleSpirV *spv_module = DCAST(ShaderModuleSpirV, module.p());
+    nassertd(spv_module != nullptr) continue;
+
+    // Make a clean copy, so we can do some transformations on it.
+    ShaderModuleSpirV::InstructionStream instructions = spv_module->_instructions;
+    ShaderModuleSpirV::InstructionWriter writer(instructions);
+
+    size_t count = 0;
+    if (_mat_block_size > 0) {
+      writer.make_block(_mat_block_type, mat_struct_locations, spv::StorageClassUniform, count++, 1);
+    }
+    if (_ptr_block_size > 0) {
+      writer.make_block(_ptr_block_type, ptr_struct_locations, spv::StorageClassUniform, count++, 1);
+    }
+    if (_push_constant_stage_mask & (1 << (int)module->get_stage())) {
+      writer.make_block(push_constant_block_type, push_constant_locations, spv::StorageClassPushConstant, 0, 0);
+    }
+
+    // Change OpenGL conventions to Vulkan conventions.
+    for (ShaderModuleSpirV::Instruction op : instructions) {
+      if (op.opcode == spv::OpExecutionMode) {
+        if (op.nargs >= 2 && (spv::ExecutionMode)op.args[1] == spv::ExecutionModeOriginLowerLeft) {
+          op.args[1] = spv::ExecutionModeOriginUpperLeft;
+        }
+      }
+      else if (op.opcode == spv::OpDecorate) {
+        if (op.nargs >= 3 && op.args[1] == spv::DecorationBuiltIn) {
+          switch ((spv::BuiltIn)op.args[2]) {
+          case spv::BuiltInVertexId:
+            op.args[2] = spv::BuiltInVertexIndex;
+            break;
+          case spv::BuiltInInstanceId:
+            op.args[2] = spv::BuiltInInstanceIndex;
+            break;
+          default:
+            break;
+          }
+        }
+      }
+    }
+
+    VkShaderModuleCreateInfo module_info;
+    module_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
+    module_info.pNext = nullptr;
+    module_info.flags = 0;
+    module_info.codeSize = instructions.get_data_size() * 4;
+    module_info.pCode = (const uint32_t *)instructions.get_data();
+
+    VkResult err;
+    err = vkCreateShaderModule(device, &module_info, nullptr, &_modules[(size_t)spv_module->get_stage()]);
+    if (err) {
+      vulkan_error(err, "Failed to create shader modules");
+      return false;
+    }
+  }
+
+  return true;
+}
+
 /**
  *
  */
 bool VulkanShaderContext::
 make_pipeline_layout(VkDevice device) {
-  // Create a descriptor set layout.
-  uint32_t binding_count = (uint32_t)_shader->_tex_spec.size();
-  VkDescriptorSetLayoutBinding *ds_bindings = (VkDescriptorSetLayoutBinding *)
-    alloca(sizeof(VkDescriptorSetLayoutBinding) * binding_count);
+  // Create a layout containing two sets: one for textures, one for uniforms.
+  uint32_t tex_binding_count = (uint32_t)_shader->_tex_spec.size();
+  VkDescriptorSetLayoutBinding *ds_bindings[2] = {
+    (VkDescriptorSetLayoutBinding *)alloca(sizeof(VkDescriptorSetLayoutBinding) * tex_binding_count),
+    (VkDescriptorSetLayoutBinding *)alloca(sizeof(VkDescriptorSetLayoutBinding) * 2),
+  };
 
   uint32_t index = 0;
-  for (Shader::ShaderTexSpec &spec : _shader->_tex_spec) {
-    VkDescriptorSetLayoutBinding &binding = ds_bindings[index];
-    binding.binding = spec._id._seqno;
+  for (const Shader::ShaderTexSpec &spec : _shader->_tex_spec) {
+    VkDescriptorSetLayoutBinding &binding = ds_bindings[0][index];
+    binding.binding = index;
     binding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
     binding.descriptorCount = 1;
-    binding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;
+    binding.stageFlags = spec._id._stage_mask;
     binding.pImmutableSamplers = nullptr;
     ++index;
   }
 
-  VkDescriptorSetLayoutCreateInfo set_info;
-  set_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
-  set_info.pNext = nullptr;
-  set_info.flags = 0;
-  set_info.bindingCount = binding_count;
-  set_info.pBindings = ds_bindings;
+  {
+    VkDescriptorSetLayoutCreateInfo set_info;
+    set_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
+    set_info.pNext = nullptr;
+    set_info.flags = 0;
+    set_info.bindingCount = tex_binding_count;
+    set_info.pBindings = ds_bindings[0];
 
-  VkResult
-  err = vkCreateDescriptorSetLayout(device, &set_info, nullptr, &_descriptor_set_layout);
-  if (err) {
-    vulkan_error(err, "Failed to create descriptor set layout");
-    return false;
+    VkResult
+    err = vkCreateDescriptorSetLayout(device, &set_info, nullptr, &_descriptor_set_layouts[0]);
+    if (err) {
+      vulkan_error(err, "Failed to create descriptor set layout");
+      return false;
+    }
   }
 
-  VkPushConstantRange ranges[1];
-  ranges[0].stageFlags = VK_SHADER_STAGE_VERTEX_BIT;
-  ranges[0].offset = 0;
-  ranges[0].size = 64 + 16;
+  // Same for the UBOs.
+  size_t count = 0;
+  if (_mat_block_size > 0) {
+    VkDescriptorSetLayoutBinding &binding = ds_bindings[1][count];
+    binding.binding = count++;
+    binding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
+    binding.descriptorCount = 1;
+    binding.stageFlags = _mat_block_stage_mask;
+    binding.pImmutableSamplers = nullptr;
+  }
+  if (_ptr_block_size > 0) {
+    VkDescriptorSetLayoutBinding &binding = ds_bindings[1][count];
+    binding.binding = count++;
+    binding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
+    binding.descriptorCount = 1;
+    binding.stageFlags = _ptr_block_stage_mask;
+    binding.pImmutableSamplers = nullptr;
+  }
+
+  {
+    VkDescriptorSetLayoutCreateInfo set_info;
+    set_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
+    set_info.pNext = nullptr;
+    set_info.flags = 0;
+    set_info.bindingCount = count;
+    set_info.pBindings = ds_bindings[1];
+
+    VkResult
+    err = vkCreateDescriptorSetLayout(device, &set_info, nullptr, &_descriptor_set_layouts[1]);
+    if (err) {
+      vulkan_error(err, "Failed to create descriptor set layout");
+      return false;
+    }
+  }
+
+  // Create the push constant range.  Because we have declared the same block
+  // in all stages that use any push constant, we only need to define one range.
+  // I'm happy with that, because I've gone down the route of trying to sort out
+  // the ranges exactly in the way that Vulkan wants it, and it's not fun.
+  VkPushConstantRange range;
+  range.stageFlags = _push_constant_stage_mask;
+  range.offset = 0;
+  range.size = _push_constant_stage_mask ? _push_constant_block_type->get_size_bytes() : 0;
 
   VkPipelineLayoutCreateInfo layout_info;
   layout_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
   layout_info.pNext = nullptr;
   layout_info.flags = 0;
-  layout_info.setLayoutCount = 1;
-  layout_info.pSetLayouts = &_descriptor_set_layout;
-  layout_info.pushConstantRangeCount = 1;
-  layout_info.pPushConstantRanges = ranges;
+  layout_info.setLayoutCount = 2;
+  layout_info.pSetLayouts = _descriptor_set_layouts;
+  layout_info.pushConstantRangeCount = _push_constant_stage_mask ? 1 : 0;
+  layout_info.pPushConstantRanges = &range;
 
+  VkResult
   err = vkCreatePipelineLayout(device, &layout_info, nullptr, &_pipeline_layout);
   if (err) {
     vulkan_error(err, "Failed to create pipeline layout");
@@ -73,6 +252,92 @@ make_pipeline_layout(VkDevice device) {
   return true;
 }
 
+/**
+ * Updates the ShaderMatSpec uniforms.
+ */
+void VulkanShaderContext::
+update_uniform_buffers(VulkanGraphicsStateGuardian *gsg, VkCommandBuffer cmd, int altered) {
+  if (_mat_block_size == 0 && _ptr_block_size == 0) {
+    return;
+  }
+
+  size_t count = 0;
+  if (altered & _mat_deps) {
+    gsg->update_shader_matrix_cache(_shader, _mat_part_cache, altered);
+
+    void *ptr = alloca(_mat_block_size);
+
+    size_t i = 0;
+    for (Shader::ShaderMatSpec &spec : _mat_spec) {
+      const LMatrix4 *val = gsg->fetch_specified_value(spec, _mat_part_cache, altered);
+      if (!val) continue;
+  #ifndef STDFLOAT_DOUBLE
+      // In this case, the data is already single-precision.
+      const PN_float32 *data = val->get_data();
+  #else
+      // In this case, we have to convert it.
+      LMatrix4f valf = LCAST(PN_float32, *val);
+      const PN_float32 *data = valf.get_data();
+  #endif
+
+      uint32_t offset = _mat_block_type->get_member(i++).offset;
+      char *dest = (char *)ptr + offset;
+
+      switch (spec._piece) {
+      case Shader::SMP_whole:
+        memcpy(dest, data, 64);
+        break;
+      case Shader::SMP_row0:
+        memcpy(dest, data + 0, 16);
+        break;
+      case Shader::SMP_row1:
+        memcpy(dest, data + 4, 16);
+        break;
+      case Shader::SMP_row2:
+        memcpy(dest, data + 8, 16);
+        break;
+      case Shader::SMP_row3:
+        memcpy(dest, data + 12, 16);
+        break;
+      case Shader::SMP_row3x1:
+        memcpy(dest, data + 12, 4);
+        break;
+      case Shader::SMP_row3x2:
+        memcpy(dest, data + 12, 8);
+        break;
+      case Shader::SMP_row3x3:
+        memcpy(dest, data + 12, 12);
+        break;
+      case Shader::SMP_cell15:
+        memcpy(dest, data + 15, 4);
+        break;
+      case Shader::SMP_cell14:
+        memcpy(dest, data + 14, 4);
+        break;
+      case Shader::SMP_cell13:
+        memcpy(dest, data + 13, 4);
+        break;
+      default:
+        //TODO: transpositions and such
+        assert(false);
+        break;
+      }
+    }
+
+    _uniform_offsets[count] = gsg->update_dynamic_uniform_buffer(ptr, _mat_block_size);
+  }
+
+  if (_mat_block_size != 0) {
+    ++count;
+  }
+
+  //TODO: ptr inputs
+  _uniform_offsets[count] = 0;
+
+  vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, _pipeline_layout,
+                          1, 1, &_uniform_descriptor_set, count, _uniform_offsets);
+}
+
 /**
  * Returns a VkPipeline for the given RenderState+GeomVertexFormat combination.
  */

panda/src/vulkandisplay/vulkanShaderContext.h

@@ -15,6 +15,7 @@
 #define VULKANSHADERCONTEXT_H
 
 #include "config_vulkandisplay.h"
+#include "shaderModuleSpirV.h"
 
 /**
  * Manages a set of Vulkan shader modules.
@@ -22,11 +23,15 @@
 class EXPCL_VULKANDISPLAY VulkanShaderContext : public ShaderContext {
 public:
   INLINE VulkanShaderContext(Shader *shader);
-  ~VulkanShaderContext() {};
+  INLINE ~VulkanShaderContext();
 
   ALLOC_DELETED_CHAIN(VulkanShaderContext);
 
+  bool create_modules(VkDevice device, const ShaderType::Struct *push_constant_block_type);
   bool make_pipeline_layout(VkDevice device);
+
+  void update_uniform_buffers(VulkanGraphicsStateGuardian *gsg, VkCommandBuffer cmd, int altered);
+
   VkPipeline get_pipeline(VulkanGraphicsStateGuardian *gsg,
                           const RenderState *state,
                           const GeomVertexFormat *format,
@@ -35,9 +40,30 @@ public:
 
 private:
   VkShaderModule _modules[(size_t)Shader::Stage::compute + 1];
-  VkDescriptorSetLayout _descriptor_set_layout;
+  VkDescriptorSetLayout _descriptor_set_layouts[2];
   VkPipelineLayout _pipeline_layout;
 
+  // Describe the two UBOs and push constant range we create.
+  const ShaderType::Struct *_mat_block_type = nullptr;
+  const ShaderType::Struct *_ptr_block_type = nullptr;
+  VkDeviceSize _mat_block_size = 0;
+  VkDeviceSize _ptr_block_size = 0;
+  int _mat_block_stage_mask = 0;
+  int _ptr_block_stage_mask = 0;
+  int _mat_deps = Shader::SSD_NONE;
+  LMatrix4 *_mat_part_cache = nullptr;
+  pvector<Shader::ShaderMatSpec> _mat_spec;
+
+  VkDescriptorSet _uniform_descriptor_set;
+  uint32_t _uniform_offsets[2];
+
+  // These are for the push constants; maybe in the future we'll replace this
+  // with a more generic and flexible system.
+  const ShaderType::Struct *_push_constant_block_type = nullptr;
+  int _push_constant_stage_mask = 0;
+  int _projection_mat_stage_mask = 0;
+  int _color_scale_stage_mask = 0;
+
   /**
    * Stores whatever is used to key a cached pipeline into the pipeline map.
    * This allows us to map Panda states to Vulkan pipelines effectively.