Vulkan: Resource Submission Tracking

In Vulkan multiple commands can be in flight simultaneously.

These commands can share resources like descriptor sets or push
constants. When between commands these resources are updated
a new version of the resources should be created.

When a resource is updated it should check the submission id of the
command buffer. If this is different than last known by the resources,
the previous resources should be freed.
If the submission id is the same than previously it has to create a
new version of the resource to not intervene with other commands that
uses the resource before the update.
When the resource wasn't updated between multiple usages in the same
submission id it could reuse the previous resource.

This PR introduces a `ResourceTracker` and a `SubmissionTracker`.
A submission tracker can check if the command buffer is submitted.
In this case all resources of the resource tracker should be freed.
Unmodified resources in the same submission can be shared.

A resource tracker will keep track of all resources that are in
flight. After the resources are used (submission + execution) have
finished the resources can be cleared.

Pull Request: blender/blender#105183

source/blender/gpu/CMakeLists.txt

@@ -216,6 +216,7 @@ set(VULKAN_SRC
   vulkan/vk_pixel_buffer.cc
   vulkan/vk_push_constants.cc
   vulkan/vk_query.cc
+  vulkan/vk_resource_tracker.cc
   vulkan/vk_shader.cc
   vulkan/vk_shader_interface.cc
   vulkan/vk_shader_log.cc
@@ -243,6 +244,7 @@ set(VULKAN_SRC
   vulkan/vk_pixel_buffer.hh
   vulkan/vk_push_constants.hh
   vulkan/vk_query.hh
+  vulkan/vk_resource_tracker.hh
   vulkan/vk_shader.hh
   vulkan/vk_shader_interface.hh
   vulkan/vk_shader_log.hh

source/blender/gpu/tests/push_constants_test.cc

@@ -78,6 +78,11 @@ struct Shader {
   GPUShader *shader = nullptr;
   Vector<CallData> call_datas;
 
+  Shader()
+  {
+    call_datas.reserve(10);
+  }
+
   ~Shader()
   {
     if (shader != nullptr) {
@@ -117,7 +122,9 @@ struct Shader {
 
   void dispatch()
   {
-    GPU_compute_dispatch(shader, 1, 1, 1);
+    /* Dispatching 1000000 times to add some stress to the GPU. Without it tests may succeed when
+     * using too simple shaders. */
+    GPU_compute_dispatch(shader, 1000, 1000, 1);
   }
 };
 
@@ -177,11 +184,7 @@ static void test_push_constants_512bytes()
 }
 GPU_TEST(push_constants_512bytes)
 
-#if 0
 /* Schedule multiple simultaneously. */
-/* These test have been disabled for now as this will to be solved in a separate PR.
- * - `DescriptorSets` may not be altered, when they are in the command queue or being executed.
- */
 static void test_push_constants_multiple()
 {
   do_push_constants_test("gpu_push_constants_test", 10);
@@ -205,6 +208,5 @@ static void test_push_constants_multiple_512bytes()
   do_push_constants_test("gpu_push_constants_512bytes_test", 10);
 }
 GPU_TEST(push_constants_multiple_512bytes)
-#endif
 
 }  // namespace blender::gpu::tests

source/blender/gpu/vulkan/vk_backend.cc

@@ -66,13 +66,14 @@ void VKBackend::compute_dispatch(int groups_x_len, int groups_y_len, int groups_
   VKShader *shader = static_cast<VKShader *>(context.shader);
   VKCommandBuffer &command_buffer = context.command_buffer_get();
   VKPipeline &pipeline = shader->pipeline_get();
-  VKDescriptorSet &descriptor_set = pipeline.descriptor_set_get();
+  VKDescriptorSetTracker &descriptor_set = pipeline.descriptor_set_get();
   VKPushConstants &push_constants = pipeline.push_constants_get();
 
   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);
+  descriptor_set.update(context);
+  command_buffer.bind(*descriptor_set.active_descriptor_set(),
+                      shader->vk_pipeline_layout_get(),
+                      VK_PIPELINE_BIND_POINT_COMPUTE);
   command_buffer.dispatch(groups_x_len, groups_y_len, groups_z_len);
 }
 

source/blender/gpu/vulkan/vk_buffer.cc

@@ -6,6 +6,7 @@
  */
 
 #include "vk_buffer.hh"
+#include "vk_context.hh"
 
 namespace blender::gpu {
 

source/blender/gpu/vulkan/vk_buffer.hh

@@ -10,9 +10,9 @@
 #include "gpu_context_private.hh"
 
 #include "vk_common.hh"
-#include "vk_context.hh"
 
 namespace blender::gpu {
+class VKContext;
 
 /**
  * Class for handing vulkan buffers (allocation/updating/binding).

source/blender/gpu/vulkan/vk_command_buffer.cc

@@ -32,6 +32,7 @@ void VKCommandBuffer::init(const VkDevice vk_device,
   vk_device_ = vk_device;
   vk_queue_ = vk_queue;
   vk_command_buffer_ = vk_command_buffer;
+  submission_id_.reset();
 
   if (vk_fence_ == VK_NULL_HANDLE) {
     VK_ALLOCATION_CALLBACKS;
@@ -160,6 +161,7 @@ void VKCommandBuffer::submit_encoded_commands()
   submit_info.pCommandBuffers = &vk_command_buffer_;
 
   vkQueueSubmit(vk_queue_, 1, &submit_info, vk_fence_);
+  submission_id_.next();
 }
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_command_buffer.hh

@@ -8,6 +8,7 @@
 #pragma once
 
 #include "vk_common.hh"
+#include "vk_resource_tracker.hh"
 
 #include "BLI_utility_mixins.hh"
 
@@ -27,6 +28,7 @@ class VKCommandBuffer : NonCopyable, NonMovable {
 
   /** Owning handles */
   VkFence vk_fence_ = VK_NULL_HANDLE;
+  VKSubmissionID submission_id_;
 
  public:
   virtual ~VKCommandBuffer();
@@ -59,6 +61,11 @@ class VKCommandBuffer : NonCopyable, NonMovable {
    */
   void submit();
 
+  const VKSubmissionID &submission_id_get() const
+  {
+    return submission_id_;
+  }
+
  private:
   void encode_recorded_commands();
   void submit_encoded_commands();

source/blender/gpu/vulkan/vk_descriptor_pools.cc

@@ -80,7 +80,8 @@ bool VKDescriptorPools::is_last_pool_active()
   return active_pool_index_ == pools_.size() - 1;
 }
 
-VKDescriptorSet VKDescriptorPools::allocate(const VkDescriptorSetLayout &descriptor_set_layout)
+std::unique_ptr<VKDescriptorSet> VKDescriptorPools::allocate(
+    const VkDescriptorSetLayout &descriptor_set_layout)
 {
   VkDescriptorSetAllocateInfo allocate_info = {};
   VkDescriptorPool pool = active_pool_get();
@@ -102,7 +103,7 @@ VKDescriptorSet VKDescriptorPools::allocate(const VkDescriptorSetLayout &descrip
     return allocate(descriptor_set_layout);
   }
 
-  return VKDescriptorSet(pool, vk_descriptor_set);
+  return std::make_unique<VKDescriptorSet>(pool, vk_descriptor_set);
 }
 
 void VKDescriptorPools::free(VKDescriptorSet &descriptor_set)

source/blender/gpu/vulkan/vk_descriptor_pools.hh

@@ -47,7 +47,7 @@ class VKDescriptorPools {
 
   void init(const VkDevice vk_device);
 
-  VKDescriptorSet allocate(const VkDescriptorSetLayout &descriptor_set_layout);
+  std::unique_ptr<VKDescriptorSet> allocate(const VkDescriptorSetLayout &descriptor_set_layout);
   void free(VKDescriptorSet &descriptor_set);
 
   /**

source/blender/gpu/vulkan/vk_descriptor_set.cc

@@ -7,6 +7,7 @@
 
 #include "vk_descriptor_set.hh"
 #include "vk_index_buffer.hh"
+#include "vk_shader.hh"
 #include "vk_storage_buffer.hh"
 #include "vk_texture.hh"
 #include "vk_uniform_buffer.hh"
@@ -17,9 +18,7 @@
 namespace blender::gpu {
 
 VKDescriptorSet::VKDescriptorSet(VKDescriptorSet &&other)
-    : vk_descriptor_pool_(other.vk_descriptor_pool_),
-      vk_descriptor_set_(other.vk_descriptor_set_),
-      bindings_(std::move(other.bindings_))
+    : vk_descriptor_pool_(other.vk_descriptor_pool_), vk_descriptor_set_(other.vk_descriptor_set_)
 {
   other.mark_freed();
 }
@@ -40,7 +39,8 @@ void VKDescriptorSet::mark_freed()
   vk_descriptor_pool_ = VK_NULL_HANDLE;
 }
 
-void VKDescriptorSet::bind(VKStorageBuffer &buffer, const Location location)
+void VKDescriptorSetTracker::bind(VKStorageBuffer &buffer,
+                                  const VKDescriptorSet::Location location)
 {
   Binding &binding = ensure_location(location);
   binding.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
@@ -48,7 +48,8 @@ void VKDescriptorSet::bind(VKStorageBuffer &buffer, const Location location)
   binding.buffer_size = buffer.size_in_bytes();
 }
 
-void VKDescriptorSet::bind_as_ssbo(VKVertexBuffer &buffer, const Location location)
+void VKDescriptorSetTracker::bind_as_ssbo(VKVertexBuffer &buffer,
+                                          const VKDescriptorSet::Location location)
 {
   Binding &binding = ensure_location(location);
   binding.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
@@ -56,7 +57,8 @@ void VKDescriptorSet::bind_as_ssbo(VKVertexBuffer &buffer, const Location locati
   binding.buffer_size = buffer.size_used_get();
 }
 
-void VKDescriptorSet::bind(VKUniformBuffer &buffer, const Location location)
+void VKDescriptorSetTracker::bind(VKUniformBuffer &buffer,
+                                  const VKDescriptorSet::Location location)
 {
   Binding &binding = ensure_location(location);
   binding.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
@@ -64,7 +66,8 @@ void VKDescriptorSet::bind(VKUniformBuffer &buffer, const Location location)
   binding.buffer_size = buffer.size_in_bytes();
 }
 
-void VKDescriptorSet::bind_as_ssbo(VKIndexBuffer &buffer, const Location location)
+void VKDescriptorSetTracker::bind_as_ssbo(VKIndexBuffer &buffer,
+                                          const VKDescriptorSet::Location location)
 {
   Binding &binding = ensure_location(location);
   binding.type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
@@ -72,14 +75,16 @@ void VKDescriptorSet::bind_as_ssbo(VKIndexBuffer &buffer, const Location locatio
   binding.buffer_size = buffer.size_get();
 }
 
-void VKDescriptorSet::image_bind(VKTexture &texture, const Location location)
+void VKDescriptorSetTracker::image_bind(VKTexture &texture,
+                                        const VKDescriptorSet::Location location)
 {
   Binding &binding = ensure_location(location);
   binding.type = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
   binding.vk_image_view = texture.vk_image_view_handle();
 }
 
-VKDescriptorSet::Binding &VKDescriptorSet::ensure_location(const Location location)
+VKDescriptorSetTracker::Binding &VKDescriptorSetTracker::ensure_location(
+    const VKDescriptorSet::Location location)
 {
   for (Binding &binding : bindings_) {
     if (binding.location == location) {
@@ -93,8 +98,12 @@ VKDescriptorSet::Binding &VKDescriptorSet::ensure_location(const Location locati
   return bindings_.last();
 }
 
-void VKDescriptorSet::update(VkDevice vk_device)
+void VKDescriptorSetTracker::update(VKContext &context)
 {
+  tracked_resource_for(context, !bindings_.is_empty());
+  std::unique_ptr<VKDescriptorSet> &descriptor_set = active_descriptor_set();
+  VkDescriptorSet vk_descriptor_set = descriptor_set->vk_handle();
+
   Vector<VkDescriptorBufferInfo> buffer_infos;
   Vector<VkWriteDescriptorSet> descriptor_writes;
 
@@ -109,7 +118,7 @@ void VKDescriptorSet::update(VkDevice vk_device)
 
     VkWriteDescriptorSet write_descriptor = {};
     write_descriptor.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
-    write_descriptor.dstSet = vk_descriptor_set_;
+    write_descriptor.dstSet = vk_descriptor_set;
     write_descriptor.dstBinding = binding.location;
     write_descriptor.descriptorCount = 1;
     write_descriptor.descriptorType = binding.type;
@@ -129,7 +138,7 @@ void VKDescriptorSet::update(VkDevice vk_device)
 
     VkWriteDescriptorSet write_descriptor = {};
     write_descriptor.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
-    write_descriptor.dstSet = vk_descriptor_set_;
+    write_descriptor.dstSet = vk_descriptor_set;
     write_descriptor.dstBinding = binding.location;
     write_descriptor.descriptorCount = 1;
     write_descriptor.descriptorType = binding.type;
@@ -141,10 +150,16 @@ void VKDescriptorSet::update(VkDevice vk_device)
                  "Not all changes have been converted to a write descriptor. Check "
                  "`Binding::is_buffer` and `Binding::is_image`.");
 
+  VkDevice vk_device = context.device_get();
   vkUpdateDescriptorSets(
       vk_device, descriptor_writes.size(), descriptor_writes.data(), 0, nullptr);
 
   bindings_.clear();
 }
 
+std::unique_ptr<VKDescriptorSet> VKDescriptorSetTracker::create_resource(VKContext &context)
+{
+  return context.descriptor_pools_get().allocate(layout_);
+}
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_descriptor_set.hh

@@ -12,7 +12,10 @@
 
 #include "gpu_shader_private.hh"
 
+#include "vk_buffer.hh"
 #include "vk_common.hh"
+#include "vk_resource_tracker.hh"
+#include "vk_uniform_buffer.hh"
 
 namespace blender::gpu {
 class VKIndexBuffer;
@@ -21,6 +24,7 @@ class VKStorageBuffer;
 class VKTexture;
 class VKUniformBuffer;
 class VKVertexBuffer;
+class VKDescriptorSetTracker;
 
 /**
  * In vulkan shader resources (images and buffers) are grouped in descriptor sets.
@@ -31,7 +35,6 @@ class VKVertexBuffer;
  * to use 2 descriptor sets per shader. One for each #blender::gpu::shader::Frequency.
  */
 class VKDescriptorSet : NonCopyable {
-  struct Binding;
 
  public:
   /**
@@ -69,42 +72,13 @@ class VKDescriptorSet : NonCopyable {
       return binding;
     }
 
-    friend struct Binding;
+    friend struct VKDescriptorSetTracker;
     friend class VKShaderInterface;
   };
 
- private:
-  struct Binding {
-    Location location;
-    VkDescriptorType type;
-
-    VkBuffer vk_buffer = VK_NULL_HANDLE;
-    VkDeviceSize buffer_size = 0;
-
-    VkImageView vk_image_view = VK_NULL_HANDLE;
-
-    Binding()
-    {
-      location.binding = 0;
-    }
-
-    bool is_buffer() const
-    {
-      return ELEM(type, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER);
-    }
-
-    bool is_image() const
-    {
-      return ELEM(type, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE);
-    }
-  };
-
   VkDescriptorPool vk_descriptor_pool_ = VK_NULL_HANDLE;
   VkDescriptorSet vk_descriptor_set_ = VK_NULL_HANDLE;
 
-  /** A list of bindings that needs to be updated. */
-  Vector<Binding> bindings_;
-
  public:
   VKDescriptorSet() = default;
   VKDescriptorSet(VkDescriptorPool vk_descriptor_pool, VkDescriptorSet vk_descriptor_set)
@@ -131,22 +105,73 @@ class VKDescriptorSet : NonCopyable {
   {
     return vk_descriptor_pool_;
   }
+  void mark_freed();
+};
 
-  void bind_as_ssbo(VKVertexBuffer &buffer, Location location);
-  void bind_as_ssbo(VKIndexBuffer &buffer, Location location);
-  void bind(VKStorageBuffer &buffer, Location location);
-  void bind(VKUniformBuffer &buffer, Location location);
-  void image_bind(VKTexture &texture, Location location);
+class VKDescriptorSetTracker : protected VKResourceTracker<VKDescriptorSet> {
+  friend class VKDescriptorSet;
+
+ public:
+  struct Binding {
+    VKDescriptorSet::Location location;
+    VkDescriptorType type;
+
+    VkBuffer vk_buffer = VK_NULL_HANDLE;
+    VkDeviceSize buffer_size = 0;
+
+    VkImageView vk_image_view = VK_NULL_HANDLE;
+
+    Binding()
+    {
+      location.binding = 0;
+    }
+
+    bool is_buffer() const
+    {
+      return ELEM(type, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER);
+    }
+
+    bool is_image() const
+    {
+      return ELEM(type, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE);
+    }
+  };
+
+ private:
+  /** A list of bindings that needs to be updated.*/
+  Vector<Binding> bindings_;
+  VkDescriptorSetLayout layout_;
+
+ public:
+  VKDescriptorSetTracker()
+  {
+  }
+
+  VKDescriptorSetTracker(VkDescriptorSetLayout layout) : layout_(layout)
+  {
+  }
+
+  void bind_as_ssbo(VKVertexBuffer &buffer, VKDescriptorSet::Location location);
+  void bind_as_ssbo(VKIndexBuffer &buffer, VKDescriptorSet::Location location);
+  void bind(VKStorageBuffer &buffer, VKDescriptorSet::Location location);
+  void bind(VKUniformBuffer &buffer, VKDescriptorSet::Location location);
+  void image_bind(VKTexture &texture, VKDescriptorSet::Location location);
 
   /**
    * Update the descriptor set on the device.
    */
-  void update(VkDevice vk_device);
+  void update(VKContext &context);
 
-  void mark_freed();
+  std::unique_ptr<VKDescriptorSet> &active_descriptor_set()
+  {
+    return active_resource();
+  }
+
+ protected:
+  std::unique_ptr<VKDescriptorSet> create_resource(VKContext &context) override;
 
  private:
-  Binding &ensure_location(Location location);
+  Binding &ensure_location(VKDescriptorSet::Location location);
 };
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_pipeline.cc

@@ -12,10 +12,10 @@
 namespace blender::gpu {
 
 VKPipeline::VKPipeline(VkPipeline vk_pipeline,
-                       VKDescriptorSet &&descriptor_set,
+                       VkDescriptorSetLayout vk_descriptor_set_layout,
                        VKPushConstants &&push_constants)
     : vk_pipeline_(vk_pipeline),
-      descriptor_set_(std::move(descriptor_set)),
+      descriptor_set_(vk_descriptor_set_layout),
       push_constants_(std::move(push_constants))
 {
 }
@@ -56,9 +56,8 @@ VKPipeline VKPipeline::create_compute_pipeline(
     return VKPipeline();
   }
 
-  VKDescriptorSet descriptor_set = context.descriptor_pools_get().allocate(descriptor_set_layout);
   VKPushConstants push_constants(&push_constants_layout);
-  return VKPipeline(vk_pipeline, std::move(descriptor_set), std::move(push_constants));
+  return VKPipeline(vk_pipeline, descriptor_set_layout, std::move(push_constants));
 }
 
 VkPipeline VKPipeline::vk_handle() const

source/blender/gpu/vulkan/vk_pipeline.hh

@@ -21,7 +21,7 @@ class VKContext;
 
 class VKPipeline : NonCopyable {
   VkPipeline vk_pipeline_ = VK_NULL_HANDLE;
-  VKDescriptorSet descriptor_set_;
+  VKDescriptorSetTracker descriptor_set_;
   VKPushConstants push_constants_;
 
  public:
@@ -29,7 +29,7 @@ class VKPipeline : NonCopyable {
 
   virtual ~VKPipeline();
   VKPipeline(VkPipeline vk_pipeline,
-             VKDescriptorSet &&vk_descriptor_set,
+             VkDescriptorSetLayout vk_descriptor_set_layout,
              VKPushConstants &&push_constants);
   VKPipeline &operator=(VKPipeline &&other)
   {
@@ -46,7 +46,7 @@ class VKPipeline : NonCopyable {
                                             VkPipelineLayout &pipeline_layouts,
                                             const VKPushConstants::Layout &push_constants_layout);
 
-  VKDescriptorSet &descriptor_set_get()
+  VKDescriptorSetTracker &descriptor_set_get()
   {
     return descriptor_set_;
   }

source/blender/gpu/vulkan/vk_pixel_buffer.cc

@@ -7,6 +7,8 @@
 
 #include "vk_pixel_buffer.hh"
 
+#include "vk_context.hh"
+
 namespace blender::gpu {
 
 VKPixelBuffer::VKPixelBuffer(int64_t size) : PixelBuffer(size)

source/blender/gpu/vulkan/vk_push_constants.cc

@@ -7,6 +7,7 @@
 
 #include "vk_push_constants.hh"
 #include "vk_backend.hh"
+#include "vk_context.hh"
 #include "vk_memory_layout.hh"
 #include "vk_shader.hh"
 #include "vk_shader_interface.hh"
@@ -103,24 +104,12 @@ VKPushConstants::VKPushConstants() = default;
 VKPushConstants::VKPushConstants(const Layout *layout) : layout_(layout)
 {
   data_ = MEM_mallocN(layout->size_in_bytes(), __func__);
-  switch (layout_->storage_type_get()) {
-    case StorageType::UNIFORM_BUFFER:
-      uniform_buffer_ = new VKUniformBuffer(layout_->size_in_bytes(), __func__);
-      break;
-
-    case StorageType::PUSH_CONSTANTS:
-    case StorageType::NONE:
-      break;
-  }
 }
 
 VKPushConstants::VKPushConstants(VKPushConstants &&other) : layout_(other.layout_)
 {
   data_ = other.data_;
   other.data_ = nullptr;
-
-  uniform_buffer_ = other.uniform_buffer_;
-  other.uniform_buffer_ = nullptr;
 }
 
 VKPushConstants::~VKPushConstants()
@@ -129,9 +118,6 @@ VKPushConstants::~VKPushConstants()
     MEM_freeN(data_);
     data_ = nullptr;
   }
-
-  delete uniform_buffer_;
-  uniform_buffer_ = nullptr;
 }
 
 VKPushConstants &VKPushConstants::operator=(VKPushConstants &&other)
@@ -141,9 +127,6 @@ VKPushConstants &VKPushConstants::operator=(VKPushConstants &&other)
   data_ = other.data_;
   other.data_ = nullptr;
 
-  uniform_buffer_ = other.uniform_buffer_;
-  other.uniform_buffer_ = nullptr;
-
   return *this;
 }
 
@@ -155,7 +138,7 @@ void VKPushConstants::update(VKContext &context)
   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();
+  VKDescriptorSetTracker &descriptor_set = pipeline.descriptor_set_get();
 
   switch (layout_get().storage_type_get()) {
     case VKPushConstants::StorageType::NONE:
@@ -167,7 +150,7 @@ void VKPushConstants::update(VKContext &context)
 
     case VKPushConstants::StorageType::UNIFORM_BUFFER:
       update_uniform_buffer();
-      descriptor_set.bind(uniform_buffer_get(), layout_get().descriptor_set_location_get());
+      descriptor_set.bind(*uniform_buffer_get(), layout_get().descriptor_set_location_get());
       break;
   }
 }
@@ -175,16 +158,22 @@ void VKPushConstants::update(VKContext &context)
 void VKPushConstants::update_uniform_buffer()
 {
   BLI_assert(layout_->storage_type_get() == StorageType::UNIFORM_BUFFER);
-  BLI_assert(uniform_buffer_ != nullptr);
   BLI_assert(data_ != nullptr);
-  uniform_buffer_->update(data_);
+  VKContext &context = *VKContext::get();
+  std::unique_ptr<VKUniformBuffer> &uniform_buffer = tracked_resource_for(context, is_dirty_);
+  uniform_buffer->update(data_);
+  is_dirty_ = false;
 }
 
-VKUniformBuffer &VKPushConstants::uniform_buffer_get()
+std::unique_ptr<VKUniformBuffer> &VKPushConstants::uniform_buffer_get()
 {
   BLI_assert(layout_->storage_type_get() == StorageType::UNIFORM_BUFFER);
-  BLI_assert(uniform_buffer_ != nullptr);
-  return *uniform_buffer_;
+  return active_resource();
+}
+
+std::unique_ptr<VKUniformBuffer> VKPushConstants::create_resource(VKContext & /*context*/)
+{
+  return std::make_unique<VKUniformBuffer>(layout_->size_in_bytes(), __func__);
 }
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_push_constants.hh

@@ -43,7 +43,7 @@ class VKContext;
  * It should also keep track of the submissions in order to reuse the allocated
  * data.
  */
-class VKPushConstants : NonCopyable {
+class VKPushConstants : VKResourceTracker<VKUniformBuffer> {
  public:
   /** Different methods to store push constants. */
   enum class StorageType {
@@ -151,7 +151,7 @@ class VKPushConstants : NonCopyable {
  private:
   const Layout *layout_ = nullptr;
   void *data_ = nullptr;
-  VKUniformBuffer *uniform_buffer_ = nullptr;
+  bool is_dirty_ = false;
 
  public:
   VKPushConstants();
@@ -171,6 +171,11 @@ class VKPushConstants : NonCopyable {
     return *layout_;
   }
 
+  /**
+   * Part of Resource Tracking API is called when new resource is needed.
+   */
+  std::unique_ptr<VKUniformBuffer> create_resource(VKContext &context) override;
+
   /**
    * Get the reference to the active data.
    *
@@ -209,6 +214,7 @@ class VKPushConstants : NonCopyable {
                          layout_->size_in_bytes(),
                      "Tried to write outside the push constant allocated memory.");
       memcpy(dst, input_data, comp_len * array_size * sizeof(T));
+      is_dirty_ = true;
       return;
     }
 
@@ -222,6 +228,7 @@ class VKPushConstants : NonCopyable {
       src += comp_len;
       dst += 4;
     }
+    is_dirty_ = true;
   }
 
   /**
@@ -243,7 +250,7 @@ class VKPushConstants : NonCopyable {
    *
    * Only valid when storage type = StorageType::UNIFORM_BUFFER.
    */
-  VKUniformBuffer &uniform_buffer_get();
+  std::unique_ptr<VKUniformBuffer> &uniform_buffer_get();
 };
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_resource_tracker.cc

@@ -0,0 +1,23 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#include "vk_resource_tracker.hh"
+#include "vk_context.hh"
+
+namespace blender::gpu {
+bool VKSubmissionTracker::is_changed(VKContext &context)
+{
+  VKCommandBuffer &command_buffer = context.command_buffer_get();
+  const VKSubmissionID &current_id = command_buffer.submission_id_get();
+  if (last_known_id_ != current_id) {
+    last_known_id_ = current_id;
+    return true;
+  }
+  return false;
+}
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_resource_tracker.hh

@@ -0,0 +1,177 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#pragma once
+
+#include "BLI_utility_mixins.hh"
+#include "BLI_vector.hh"
+
+#include "vk_common.hh"
+
+namespace blender::gpu {
+class VKContext;
+class VKCommandBuffer;
+
+/**
+ * In vulkan multiple commands can be in flight simultaneously.
+ *
+ * These commands can share the same resources like descriptor sets
+ * or push constants. When between commands these resources are updated
+ * a new version of these resources should be created.
+ *
+ * When a resource is updated it should check the submission id of the
+ * command buffer. If it is different, then the resource can be reused.
+ * If the submission id is the same a new version of the resource to now
+ * intervene with other commands that uses the resource.
+ *
+ * VKSubmissionID is the identifier to keep track if a new submission is
+ * being recorded.
+ */
+struct VKSubmissionID {
+ private:
+  int64_t id_ = -1;
+
+ public:
+  VKSubmissionID() = default;
+
+ private:
+  /**
+   * Reset the submission id.
+   *
+   * This should only be called during initialization of the command buffer.
+   * As it leads to undesired behavior after resources are already tracking
+   * the submission id.
+   */
+  void reset()
+  {
+    id_ = 0;
+  }
+
+  /**
+   * Change the submission id.
+   *
+   * Is called when submitting a command buffer to the queue. In this case resource
+   * known that the next time it is used that it can free its sub resources used by
+   * the previous submission.
+   */
+  void next()
+  {
+    id_++;
+  }
+
+ public:
+  const VKSubmissionID &operator=(const VKSubmissionID &other)
+  {
+    id_ = other.id_;
+    return *this;
+  }
+
+  bool operator==(const VKSubmissionID &other)
+  {
+    return id_ == other.id_;
+  }
+
+  bool operator!=(const VKSubmissionID &other)
+  {
+    return id_ != other.id_;
+  }
+
+  friend class VKCommandBuffer;
+};
+
+/**
+ * Submission tracker keeps track of the last known submission id of the
+ * command buffer.
+ */
+class VKSubmissionTracker {
+  VKSubmissionID last_known_id_;
+
+ public:
+  /**
+   * Check if the submission_id has changed since the last time it was called
+   * on this VKSubmissionTracker.
+   */
+  bool is_changed(VKContext &context);
+};
+
+/**
+ * VKResourceTracker will keep track of resources.
+ */
+template<typename Resource> class VKResourceTracker : NonCopyable {
+  VKSubmissionTracker submission_tracker_;
+  Vector<std::unique_ptr<Resource>> tracked_resources_;
+
+ protected:
+  VKResourceTracker<Resource>() = default;
+  VKResourceTracker<Resource>(VKResourceTracker<Resource> &&other)
+      : submission_tracker_(other.submission_tracker_),
+        tracked_resources_(std::move(other.tracked_resources_))
+  {
+  }
+
+  VKResourceTracker<Resource> &operator=(VKResourceTracker<Resource> &&other)
+  {
+    submission_tracker_ = other.submission_tracker_;
+    tracked_resources_ = std::move(other.tracked_resources_);
+    return *this;
+  }
+
+  virtual ~VKResourceTracker()
+  {
+    free_tracked_resources();
+  }
+
+  /**
+   * Get a resource what can be used by the resource tracker.
+   *
+   * When a different submission was detected all previous resources
+   * will be freed and a new resource will be returned.
+   *
+   * When still in the same submission and we need to update the resource
+   * (is_dirty=true) then a new resource will be returned. Otherwise
+   * the previous used resource will be used.
+   *
+   * When no resources exists, a new resource will be created.
+   *
+   * The resource given back is owned by this resource tracker. And
+   * the resource should not be stored outside this class as it might
+   * be destroyed when the next submission is detected.
+   */
+  std::unique_ptr<Resource> &tracked_resource_for(VKContext &context, const bool is_dirty)
+  {
+    if (submission_tracker_.is_changed(context)) {
+      free_tracked_resources();
+      tracked_resources_.append(create_resource(context));
+    }
+    else if (is_dirty || tracked_resources_.is_empty()) {
+      tracked_resources_.append(create_resource(context));
+    }
+    return active_resource();
+  }
+
+  /**
+   * Callback to create a new resource. Can be called by the `tracked_resource_for` method.
+   */
+  virtual std::unique_ptr<Resource> create_resource(VKContext &context) = 0;
+
+  /**
+   * Return the active resource of the tracker.
+   */
+  std::unique_ptr<Resource> &active_resource()
+  {
+    BLI_assert(!tracked_resources_.is_empty());
+    return tracked_resources_.last();
+  }
+
+ private:
+  void free_tracked_resources()
+  {
+    tracked_resources_.clear();
+  }
+};
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_uniform_buffer.hh

@@ -7,13 +7,15 @@
 
 #pragma once
 
+#include "BLI_utility_mixins.hh"
+
 #include "gpu_uniform_buffer_private.hh"
 
 #include "vk_buffer.hh"
 
 namespace blender::gpu {
 
-class VKUniformBuffer : public UniformBuf {
+class VKUniformBuffer : public UniformBuf, NonCopyable {
   VKBuffer buffer_;
 
  public: