source/blender/gpu/vulkan/vk_backend.cc

@@ -6,6 +6,8 @@
  * \ingroup gpu
  */
 
+#include "GHOST_C-api.h"
+
 #include "gpu_capabilities_private.hh"
 #include "gpu_platform_private.hh"
 
@@ -121,7 +123,19 @@ void VKBackend::compute_dispatch_indirect(StorageBuf *indirect_buf)
 
 Context *VKBackend::context_alloc(void *ghost_window, void *ghost_context)
 {
-  return new VKContext(ghost_window, ghost_context);
+  if (ghost_window) {
+    BLI_assert(ghost_context == nullptr);
+    ghost_context = GHOST_GetDrawingContext((GHOST_WindowHandle)ghost_window);
+  }
+
+  BLI_assert(ghost_context != nullptr);
+  if (!device_.is_initialized()) {
+    device_.init(ghost_context);
+  }
+
+  VKContext *context = new VKContext(ghost_window, ghost_context);
+  device_.context_register(*context);
+  return context;
 }
 
 Batch *VKBackend::batch_alloc()

source/blender/gpu/vulkan/vk_context.cc

@@ -23,14 +23,7 @@ namespace blender::gpu {
 VKContext::VKContext(void *ghost_window, void *ghost_context)
 {
   ghost_window_ = ghost_window;
-  if (ghost_window) {
-    ghost_context = GHOST_GetDrawingContext((GHOST_WindowHandle)ghost_window);
-  }
   ghost_context_ = ghost_context;
-  VKDevice &device = VKBackend::get().device_;
-  if (!device.is_initialized()) {
-    device.init(ghost_context);
-  }
 
   state_manager = new VKStateManager();
   imm = new VKImmediate();
@@ -43,6 +36,8 @@ VKContext::VKContext(void *ghost_window, void *ghost_context)
 
 VKContext::~VKContext()
 {
+  VKBackend::get().device_.context_unregister(*this);
+
   delete imm;
   imm = nullptr;
 }
@@ -129,13 +124,14 @@ void VKContext::memory_statistics_get(int * /*total_mem*/, int * /*free_mem*/) {
 /* -------------------------------------------------------------------- */
 /** \name State manager
  * \{ */
-
+#if 0
 const VKStateManager &VKContext::state_manager_get() const
 {
   return *static_cast<const VKStateManager *>(state_manager);
 }
+#endif
 
-VKStateManager &VKContext::state_manager_get()
+VKStateManager &VKContext::state_manager_get() const
 {
   return *static_cast<VKStateManager *>(state_manager);
 }

source/blender/gpu/vulkan/vk_context.hh

@@ -68,8 +68,13 @@ class VKContext : public Context, NonCopyable {
     return command_buffer_;
   }
 
-  const VKStateManager &state_manager_get() const;
-  VKStateManager &state_manager_get();
+  // const VKStateManager &state_manager_get() const;
+  VKStateManager &state_manager_get() const;
 };
 
+BLI_INLINE bool operator==(const VKContext &a, const VKContext &b)
+{
+  return static_cast<const void *>(&a) == static_cast<const void *>(&b);
+}
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_device.cc

@@ -8,7 +8,12 @@
 
 #include "vk_device.hh"
 #include "vk_backend.hh"
+#include "vk_context.hh"
 #include "vk_memory.hh"
+#include "vk_state_manager.hh"
+#include "vk_storage_buffer.hh"
+#include "vk_texture.hh"
+#include "vk_vertex_buffer.hh"
 
 #include "GHOST_C-api.h"
 
@@ -178,4 +183,50 @@ std::string VKDevice::driver_version() const
 
 /** \} */
 
+/* -------------------------------------------------------------------- */
+/** \name Resource management
+ * \{ */
+
+void VKDevice::context_register(VKContext &context)
+{
+  contexts_.append(std::reference_wrapper(context));
+}
+
+void VKDevice::context_unregister(VKContext &context)
+{
+  contexts_.remove(contexts_.first_index_of(std::reference_wrapper(context)));
+}
+
+void VKDevice::unbind(VKUniformBuffer &uniform_buffer) const
+{
+  for (const VKContext &context : contexts_) {
+    context.state_manager_get().uniform_buffer_unbind(&uniform_buffer);
+  }
+}
+
+void VKDevice::unbind(VKTexture &texture) const
+{
+  for (const VKContext &context : contexts_) {
+    context.state_manager_get().image_unbind(wrap(&texture));
+    context.state_manager_get().texture_unbind(wrap(&texture));
+  }
+}
+
+void VKDevice::unbind(VKStorageBuffer &storage_buffer) const
+{
+  for (VKContext &context : contexts_) {
+    context.state_manager_get().storage_buffer_unbind(&storage_buffer);
+  }
+}
+
+void VKDevice::unbind(VKVertexBuffer &vertex_buffer) const
+{
+  for (VKContext &context : contexts_) {
+    context.state_manager_get().texel_buffer_unbind(&vertex_buffer);
+    context.state_manager_get().storage_buffer_unbind(&vertex_buffer);
+  }
+}
+
+/** \} */
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_device.hh

@@ -9,6 +9,7 @@
 #pragma once
 
 #include "BLI_utility_mixins.hh"
+#include "BLI_vector.hh"
 
 #include "vk_common.hh"
 #include "vk_debug.hh"
@@ -36,6 +37,17 @@ class VKDevice : public NonCopyable {
   uint32_t vk_queue_family_ = 0;
   VkQueue vk_queue_ = VK_NULL_HANDLE;
 
+  /**
+   * Available Contexts for this device.
+   *
+   * Device keeps track of each contexts. When buffers/images are freed they need to be removed
+   * from all contexts state managers.
+   *
+   * The contexts inside this list aren't owned by the VKDevice. Caller of `GPU_context_create`
+   * holds the ownership.
+   */
+  Vector<std::reference_wrapper<VKContext>> contexts_;
+
   /** Allocator used for texture and buffers and other resources. */
   VmaAllocator mem_allocator_ = VK_NULL_HANDLE;
   VKDescriptorPools descriptor_pools_;
@@ -114,6 +126,35 @@ class VKDevice : public NonCopyable {
     return workarounds_;
   }
 
+  /* -------------------------------------------------------------------- */
+  /** \name Resource management
+   * \{ */
+
+  void context_register(VKContext &context);
+  void context_unregister(VKContext &context);
+
+  /**
+   * Unbind given uniform buffer from registered contexts.
+   */
+  void unbind(VKUniformBuffer &uniform_buffer) const;
+
+  /**
+   * Unbind given texture from registered contexts.
+   */
+  void unbind(VKTexture &texture) const;
+
+  /**
+   * Unbind given storage buffer from registered contexts.
+   */
+  void unbind(VKStorageBuffer &storage_buffer) const;
+
+  /**
+   * Unbind given vertex_buffer from registered contexts.
+   */
+  void unbind(VKVertexBuffer &vertex_buffer) const;
+
+  /** \} */
+
  private:
   void init_physical_device_properties();
   void init_debug_callbacks();

source/blender/gpu/vulkan/vk_state_manager.cc

@@ -10,6 +10,7 @@
 #include "vk_context.hh"
 #include "vk_pipeline.hh"
 #include "vk_shader.hh"
+#include "vk_storage_buffer.hh"
 #include "vk_texture.hh"
 #include "vk_vertex_buffer.hh"
 
@@ -27,6 +28,8 @@ VKStateManager::VKStateManager()
   texture_bindings_.fill(TextureBinding());
   uniform_buffer_bindings_ = Array<UniformBufferBinding>(max_bindings);
   uniform_buffer_bindings_.fill(UniformBufferBinding());
+  storage_buffer_bindings_ = Array<StorageBufferBinding>(max_bindings);
+  storage_buffer_bindings_.fill(StorageBufferBinding());
 }
 
 void VKStateManager::apply_state()
@@ -54,7 +57,8 @@ void VKStateManager::apply_bindings()
         texture_bindings_[binding].texture->bind(binding, sampler_);
       }
       else if (texture_bindings_[binding].vertex_buffer) {
-        texture_bindings_[binding].vertex_buffer->bind(binding);
+        texture_bindings_[binding].vertex_buffer->bind(
+            binding, shader::ShaderCreateInfo::Resource::BindType::SAMPLER);
       }
     }
 
@@ -64,6 +68,17 @@ void VKStateManager::apply_bindings()
             binding, shader::ShaderCreateInfo::Resource::BindType::UNIFORM_BUFFER);
       }
     }
+
+    for (int binding : IndexRange(storage_buffer_bindings_.size())) {
+      if (storage_buffer_bindings_[binding].storage_buffer) {
+        storage_buffer_bindings_[binding].storage_buffer->bind(
+            binding, shader::ShaderCreateInfo::Resource::BindType::STORAGE_BUFFER);
+      }
+      else if (storage_buffer_bindings_[binding].vertex_buffer) {
+        storage_buffer_bindings_[binding].vertex_buffer->bind(
+            binding, shader::ShaderCreateInfo::Resource::BindType::STORAGE_BUFFER);
+      }
+    }
   }
 }
 
@@ -162,6 +177,36 @@ void VKStateManager::texel_buffer_unbind(VKVertexBuffer *vertex_buffer)
   }
 }
 
+void VKStateManager::storage_buffer_bind(VKStorageBuffer *storage_buffer, int slot)
+{
+  storage_buffer_bindings_[slot].storage_buffer = storage_buffer;
+  storage_buffer_bindings_[slot].vertex_buffer = nullptr;
+}
+
+void VKStateManager::storage_buffer_bind(VKVertexBuffer *vertex_buffer, int slot)
+{
+  storage_buffer_bindings_[slot].storage_buffer = nullptr;
+  storage_buffer_bindings_[slot].vertex_buffer = vertex_buffer;
+}
+
+void VKStateManager::storage_buffer_unbind(VKStorageBuffer *storage_buffer)
+{
+  for (StorageBufferBinding &binding : storage_buffer_bindings_) {
+    if (binding.storage_buffer == storage_buffer) {
+      binding.storage_buffer = nullptr;
+    }
+  }
+}
+
+void VKStateManager::storage_buffer_unbind(VKVertexBuffer *vertex_buffer)
+{
+  for (StorageBufferBinding &binding : storage_buffer_bindings_) {
+    if (binding.vertex_buffer == vertex_buffer) {
+      binding.vertex_buffer = nullptr;
+    }
+  }
+}
+
 void VKStateManager::texture_unpack_row_length_set(uint len)
 {
   texture_unpack_row_length_ = len;

source/blender/gpu/vulkan/vk_state_manager.hh

@@ -18,6 +18,7 @@ namespace blender::gpu {
 class VKTexture;
 class VKUniformBuffer;
 class VKVertexBuffer;
+class VKStorageBuffer;
 
 class VKStateManager : public StateManager {
   /* Dummy sampler for now. */
@@ -36,9 +37,14 @@ class VKStateManager : public StateManager {
   struct UniformBufferBinding {
     VKUniformBuffer *buffer = nullptr;
   };
+  struct StorageBufferBinding {
+    VKStorageBuffer *storage_buffer = nullptr;
+    VKVertexBuffer *vertex_buffer = nullptr;
+  };
   Array<ImageBinding> image_bindings_;
   Array<TextureBinding> texture_bindings_;
   Array<UniformBufferBinding> uniform_buffer_bindings_;
+  Array<StorageBufferBinding> storage_buffer_bindings_;
 
  public:
   VKStateManager();
@@ -65,6 +71,11 @@ class VKStateManager : public StateManager {
   void texel_buffer_bind(VKVertexBuffer *vertex_buffer, int slot);
   void texel_buffer_unbind(VKVertexBuffer *vertex_buffer);
 
+  void storage_buffer_bind(VKStorageBuffer *storage_buffer, int slot);
+  void storage_buffer_unbind(VKStorageBuffer *storage_buffer);
+  void storage_buffer_bind(VKVertexBuffer *vertex_buffer, int slot);
+  void storage_buffer_unbind(VKVertexBuffer *vertex_buffer);
+
   void texture_unpack_row_length_set(uint len) override;
 
   /**

source/blender/gpu/vulkan/vk_storage_buffer.cc

@@ -7,12 +7,23 @@
  */
 #include "vk_shader.hh"
 #include "vk_shader_interface.hh"
+#include "vk_state_manager.hh"
 #include "vk_vertex_buffer.hh"
 
 #include "vk_storage_buffer.hh"
 
 namespace blender::gpu {
 
+VKStorageBuffer::VKStorageBuffer(int size, GPUUsageType usage, const char *name)
+    : StorageBuf(size, name), usage_(usage)
+{
+}
+
+VKStorageBuffer::~VKStorageBuffer()
+{
+  unbind();
+}
+
 void VKStorageBuffer::update(const void *data)
 {
   ensure_allocated();
@@ -38,18 +49,27 @@ void VKStorageBuffer::allocate()
 
 void VKStorageBuffer::bind(int slot)
 {
-  ensure_allocated();
   VKContext &context = *VKContext::get();
-  VKShader *shader = static_cast<VKShader *>(context.shader);
-  const VKShaderInterface &shader_interface = shader->interface_get();
-  const std::optional<VKDescriptorSet::Location> location =
-      shader_interface.descriptor_set_location(
-          shader::ShaderCreateInfo::Resource::BindType::STORAGE_BUFFER, slot);
-  BLI_assert_msg(location, "Locations to SSBOs should always exist.");
-  shader->pipeline_get().descriptor_set_get().bind(*this, *location);
+  context.state_manager_get().storage_buffer_bind(this, slot);
 }
 
-void VKStorageBuffer::unbind() {}
+void VKStorageBuffer::bind(int slot, shader::ShaderCreateInfo::Resource::BindType bind_type)
+{
+  VKContext &context = *VKContext::get();
+  VKShader *shader = static_cast<VKShader *>(context.shader);
+  ensure_allocated();
+  const VKShaderInterface &shader_interface = shader->interface_get();
+  const std::optional<VKDescriptorSet::Location> location =
+      shader_interface.descriptor_set_location(bind_type, slot);
+  if (location) {
+    shader->pipeline_get().descriptor_set_get().bind(*this, *location);
+  }
+}
+
+void VKStorageBuffer::unbind()
+{
+  VKBackend::get().device_get().unbind(*this);
+}
 
 void VKStorageBuffer::clear(uint32_t clear_value)
 {

source/blender/gpu/vulkan/vk_storage_buffer.hh

@@ -23,13 +23,12 @@ class VKStorageBuffer : public StorageBuf {
   VKBuffer buffer_;
 
  public:
-  VKStorageBuffer(int size, GPUUsageType usage, const char *name)
-      : StorageBuf(size, name), usage_(usage)
-  {
-  }
+  VKStorageBuffer(int size, GPUUsageType usage, const char *name);
+  virtual ~VKStorageBuffer();
 
   void update(const void *data) override;
   void bind(int slot) override;
+  void bind(int slot, shader::ShaderCreateInfo::Resource::BindType bind_type);
   void unbind() override;
   void clear(uint32_t clear_value) override;
   void copy_sub(VertBuf *src, uint dst_offset, uint src_offset, uint copy_size) override;
@@ -51,7 +50,7 @@ class VKStorageBuffer : public StorageBuf {
   void allocate();
 };
 
-static inline VKStorageBuffer *unwrap(StorageBuf *storage_buffer)
+BLI_INLINE VKStorageBuffer *unwrap(StorageBuf *storage_buffer)
 {
   return static_cast<VKStorageBuffer *>(storage_buffer);
 }

source/blender/gpu/vulkan/vk_texture.hh

@@ -149,9 +149,14 @@ class VKTexture : public Texture {
   /** \} */
 };
 
-static inline VKTexture *unwrap(Texture *tex)
+BLI_INLINE VKTexture *unwrap(Texture *tex)
 {
   return static_cast<VKTexture *>(tex);
 }
 
+BLI_INLINE Texture *wrap(VKTexture *texture)
+{
+  return static_cast<Texture *>(texture);
+}
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_uniform_buffer.cc

@@ -73,10 +73,7 @@ void VKUniformBuffer::bind_as_ssbo(int slot)
 
 void VKUniformBuffer::unbind()
 {
-  VKContext *context = VKContext::get();
-  if (context) {
-    context->state_manager_get().uniform_buffer_unbind(this);
-  }
+  VKBackend::get().device_get().unbind(*this);
 }
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_uniform_buffer.hh

@@ -44,4 +44,9 @@ class VKUniformBuffer : public UniformBuf, NonCopyable {
   void allocate();
 };
 
+BLI_INLINE UniformBuf *wrap(VKUniformBuffer *uniform_buffer)
+{
+  return static_cast<UniformBuf *>(uniform_buffer);
+}
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_vertex_buffer.cc

@@ -24,18 +24,9 @@ VKVertexBuffer::~VKVertexBuffer()
 
 void VKVertexBuffer::bind_as_ssbo(uint binding)
 {
-  if (!buffer_.is_allocated()) {
-    allocate();
-  }
-
   VKContext &context = *VKContext::get();
-  VKShader *shader = static_cast<VKShader *>(context.shader);
-  const VKShaderInterface &shader_interface = shader->interface_get();
-  const std::optional<VKDescriptorSet::Location> location =
-      shader_interface.descriptor_set_location(
-          shader::ShaderCreateInfo::Resource::BindType::STORAGE_BUFFER, binding);
-  BLI_assert_msg(location, "Locations to SSBOs should always exist.");
-  shader->pipeline_get().descriptor_set_get().bind_as_ssbo(*this, *location);
+  VKStateManager &state_manager = context.state_manager_get();
+  state_manager.storage_buffer_bind(this, binding);
 }
 
 void VKVertexBuffer::bind_as_texture(uint binding)
@@ -43,24 +34,24 @@ void VKVertexBuffer::bind_as_texture(uint binding)
   VKContext &context = *VKContext::get();
   VKStateManager &state_manager = context.state_manager_get();
   state_manager.texel_buffer_bind(this, binding);
-  should_unbind_ = true;
 }
 
-void VKVertexBuffer::bind(uint binding)
+void VKVertexBuffer::bind(uint binding, shader::ShaderCreateInfo::Resource::BindType bind_type)
 {
   VKContext &context = *VKContext::get();
   VKShader *shader = static_cast<VKShader *>(context.shader);
   const VKShaderInterface &shader_interface = shader->interface_get();
   const std::optional<VKDescriptorSet::Location> location =
-      shader_interface.descriptor_set_location(
-          shader::ShaderCreateInfo::Resource::BindType::SAMPLER, binding);
+      shader_interface.descriptor_set_location(bind_type, binding);
   if (!location) {
     return;
   }
 
   upload_data();
 
-  if (vk_buffer_view_ == VK_NULL_HANDLE) {
+  if (bind_type == shader::ShaderCreateInfo::Resource::BindType::SAMPLER &&
+      vk_buffer_view_ == VK_NULL_HANDLE)
+  {
     VkBufferViewCreateInfo buffer_view_info = {};
     eGPUTextureFormat texture_format = to_texture_format(&format);
 
@@ -75,7 +66,14 @@ void VKVertexBuffer::bind(uint binding)
         device.device_get(), &buffer_view_info, vk_allocation_callbacks, &vk_buffer_view_);
   }
 
-  shader->pipeline_get().descriptor_set_get().bind(*this, *location);
+  /* TODO: Check if we can move this check inside the descriptor set. */
+  VKDescriptorSetTracker &descriptor_set = shader->pipeline_get().descriptor_set_get();
+  if (bind_type == shader::ShaderCreateInfo::Resource::BindType::SAMPLER) {
+    descriptor_set.bind(*this, *location);
+  }
+  else {
+    descriptor_set.bind_as_ssbo(*this, *location);
+  }
 }
 
 void VKVertexBuffer::wrap_handle(uint64_t /*handle*/)
@@ -119,14 +117,11 @@ void VKVertexBuffer::resize_data()
 
 void VKVertexBuffer::release_data()
 {
-  VKContext *context = VKContext::get();
-  if (should_unbind_ && context) {
-    context->state_manager_get().texel_buffer_unbind(this);
-  }
+  const VKDevice &device = VKBackend::get().device_get();
+  device.unbind(*this);
 
   if (vk_buffer_view_ != VK_NULL_HANDLE) {
     VK_ALLOCATION_CALLBACKS;
-    const VKDevice &device = VKBackend::get().device_get();
     vkDestroyBufferView(device.device_get(), vk_buffer_view_, vk_allocation_callbacks);
     vk_buffer_view_ = VK_NULL_HANDLE;
   }

source/blender/gpu/vulkan/vk_vertex_buffer.hh

@@ -16,7 +16,6 @@ namespace blender::gpu {
 
 class VKVertexBuffer : public VertBuf {
   VKBuffer buffer_;
-  bool should_unbind_ = false;
   /** When a vertex buffer is used as a UNIFORM_TEXEL_BUFFER the buffer requires a buffer view. */
   VkBufferView vk_buffer_view_ = VK_NULL_HANDLE;
 
@@ -25,7 +24,7 @@ class VKVertexBuffer : public VertBuf {
 
   void bind_as_ssbo(uint binding) override;
   void bind_as_texture(uint binding) override;
-  void bind(uint binding);
+  void bind(uint binding, shader::ShaderCreateInfo::Resource::BindType bind_type);
   void wrap_handle(uint64_t handle) override;
 
   void update_sub(uint start, uint len, const void *data) override;
@@ -58,7 +57,7 @@ class VKVertexBuffer : public VertBuf {
   friend class VKTexture;
 };
 
-static inline VKVertexBuffer *unwrap(VertBuf *vertex_buffer)
+BLI_INLINE VKVertexBuffer *unwrap(VertBuf *vertex_buffer)
 {
   return static_cast<VKVertexBuffer *>(vertex_buffer);
 }