source/blender/gpu/vulkan/vk_backend.cc

@@ -69,22 +69,7 @@ void VKBackend::compute_dispatch(int groups_x_len, int groups_y_len, int groups_
   VKDescriptorSet &descriptor_set = pipeline.descriptor_set_get();
   VKPushConstants &push_constants = pipeline.push_constants_get();
 
-  /* Update push constants based on their storage type.*/
-  switch (push_constants.layout_get().storage_type_get()) {
-    case VKPushConstants::StorageType::NONE:
-      break;
-
-    case VKPushConstants::StorageType::PUSH_CONSTANTS:
-      command_buffer.push_constants(
-          push_constants, shader->vk_pipeline_layout_get(), VK_SHADER_STAGE_ALL);
-      break;
-
-    case VKPushConstants::StorageType::UNIFORM_BUFFER:
-      push_constants.update_uniform_buffer();
-      descriptor_set.bind(push_constants.uniform_buffer_get(),
-                          push_constants.layout_get().descriptor_set_location_get());
-      break;
-  }
+  push_constants.update(context);
   descriptor_set.update(context.device_get());
   command_buffer.bind(
       descriptor_set, shader->vk_pipeline_layout_get(), VK_PIPELINE_BIND_POINT_COMPUTE);

source/blender/gpu/vulkan/vk_push_constants.cc

@@ -8,6 +8,7 @@
 #include "vk_push_constants.hh"
 #include "vk_backend.hh"
 #include "vk_memory_layout.hh"
+#include "vk_shader.hh"
 #include "vk_shader_interface.hh"
 #include "vk_storage_buffer.hh"
 #include "vk_uniform_buffer.hh"
@@ -146,6 +147,31 @@ VKPushConstants &VKPushConstants::operator=(VKPushConstants &&other)
   return *this;
 }
 
+void VKPushConstants::update(VKContext &context)
+{
+  VKShader *shader = static_cast<VKShader *>(context.shader);
+  VKCommandBuffer &command_buffer = context.command_buffer_get();
+  VKPipeline &pipeline = shader->pipeline_get();
+  BLI_assert_msg(&pipeline.push_constants_get() == this,
+                 "Invalid state detected. Push constants doesn't belong to the active shader of "
+                 "the given context.");
+  VKDescriptorSet &descriptor_set = pipeline.descriptor_set_get();
+
+  switch (layout_get().storage_type_get()) {
+    case VKPushConstants::StorageType::NONE:
+      break;
+
+    case VKPushConstants::StorageType::PUSH_CONSTANTS:
+      command_buffer.push_constants(*this, shader->vk_pipeline_layout_get(), VK_SHADER_STAGE_ALL);
+      break;
+
+    case VKPushConstants::StorageType::UNIFORM_BUFFER:
+      update_uniform_buffer();
+      descriptor_set.bind(uniform_buffer_get(), layout_get().descriptor_set_location_get());
+      break;
+  }
+}
+
 void VKPushConstants::update_uniform_buffer()
 {
   BLI_assert(layout_->storage_type_get() == StorageType::UNIFORM_BUFFER);

source/blender/gpu/vulkan/vk_push_constants.hh

@@ -29,6 +29,7 @@
 namespace blender::gpu {
 class VKShaderInterface;
 class VKUniformBuffer;
+class VKContext;
 
 /**
  * Container to store push constants in a buffer.
@@ -170,21 +171,6 @@ class VKPushConstants : NonCopyable {
     return *layout_;
   }
 
-  /**
-   * When storage type = StorageType::UNIFORM_BUFFER use this method to update the uniform
-   * buffer.
-   *
-   * It must be called just before adding a draw/compute command to the command queue.
-   */
-  void update_uniform_buffer();
-
-  /**
-   * Get a reference to the uniform buffer.
-   *
-   * Only valid when storage type = StorageType::UNIFORM_BUFFER.
-   */
-  VKUniformBuffer &uniform_buffer_get();
-
   /**
    * Get the reference to the active data.
    *
@@ -237,6 +223,27 @@ class VKPushConstants : NonCopyable {
       dst += 4;
     }
   }
+
+  /**
+   * Update the GPU resources with the latest push constants.
+   */
+  void update(VKContext &context);
+
+ private:
+  /**
+   * When storage type = StorageType::UNIFORM_BUFFER use this method to update the uniform
+   * buffer.
+   *
+   * It must be called just before adding a draw/compute command to the command queue.
+   */
+  void update_uniform_buffer();
+
+  /**
+   * Get a reference to the uniform buffer.
+   *
+   * Only valid when storage type = StorageType::UNIFORM_BUFFER.
+   */
+  VKUniformBuffer &uniform_buffer_get();
 };
 
 }  // namespace blender::gpu