4 files changed, 56 insertions, 26 deletions

diff --git a/src/video_core/buffer_cache/buffer_cache.h b/src/video_core/buffer_cache/buffer_cache.h
index 6701aab82..29746f61d 100644
--- a/src/video_core/buffer_cache/buffer_cache.h
+++ b/src/video_core/buffer_cache/buffer_cache.h
@@ -154,7 +154,7 @@ public:
     void UnbindGraphicsTextureBuffers(size_t stage);
 
     void BindGraphicsTextureBuffer(size_t stage, size_t tbo_index, GPUVAddr gpu_addr, u32 size,
-                                   PixelFormat format);
+                                   PixelFormat format, bool is_written);
 
     void UnbindComputeStorageBuffers();
 
@@ -163,8 +163,8 @@ public:
 
     void UnbindComputeTextureBuffers();
 
-    void BindComputeTextureBuffer(size_t tbo_index, GPUVAddr gpu_addr, u32 size,
-                                  PixelFormat format);
+    void BindComputeTextureBuffer(size_t tbo_index, GPUVAddr gpu_addr, u32 size, PixelFormat format,
+                                  bool is_written);
 
     void FlushCachedWrites();
 
@@ -393,7 +393,9 @@ private:
     u32 written_compute_storage_buffers = 0;
 
     std::array<u32, NUM_STAGES> enabled_texture_buffers{};
+    std::array<u32, NUM_STAGES> written_texture_buffers{};
     u32 enabled_compute_texture_buffers = 0;
+    u32 written_compute_texture_buffers = 0;
 
     std::array<u32, NUM_STAGES> fast_bound_uniform_buffers{};
 
@@ -700,12 +702,14 @@ void BufferCache<P>::BindGraphicsStorageBuffer(size_t stage, size_t ssbo_index,
 template <class P>
 void BufferCache<P>::UnbindGraphicsTextureBuffers(size_t stage) {
     enabled_texture_buffers[stage] = 0;
+    written_texture_buffers[stage] = 0;
 }
 
 template <class P>
 void BufferCache<P>::BindGraphicsTextureBuffer(size_t stage, size_t tbo_index, GPUVAddr gpu_addr,
-                                               u32 size, PixelFormat format) {
+                                               u32 size, PixelFormat format, bool is_written) {
     enabled_texture_buffers[stage] |= 1U << tbo_index;
+    written_texture_buffers[stage] |= (is_written ? 1U : 0U) << tbo_index;
     texture_buffers[stage][tbo_index] = GetTextureBufferBinding(gpu_addr, size, format);
 }
 
@@ -732,12 +736,14 @@ void BufferCache<P>::BindComputeStorageBuffer(size_t ssbo_index, u32 cbuf_index,
 template <class P>
 void BufferCache<P>::UnbindComputeTextureBuffers() {
     enabled_compute_texture_buffers = 0;
+    written_compute_texture_buffers = 0;
 }
 
 template <class P>
 void BufferCache<P>::BindComputeTextureBuffer(size_t tbo_index, GPUVAddr gpu_addr, u32 size,
-                                              PixelFormat format) {
+                                              PixelFormat format, bool is_written) {
     enabled_compute_texture_buffers |= 1U << tbo_index;
+    written_compute_texture_buffers |= (is_written ? 1U : 0U) << tbo_index;
     compute_texture_buffers[tbo_index] = GetTextureBufferBinding(gpu_addr, size, format);
 }
 
@@ -1274,6 +1280,10 @@ void BufferCache<P>::UpdateTextureBuffers(size_t stage) {
     ForEachEnabledBit(enabled_texture_buffers[stage], [&](u32 index) {
         Binding& binding = texture_buffers[stage][index];
         binding.buffer_id = FindBuffer(binding.cpu_addr, binding.size);
+        // Mark buffer as written if needed
+        if (((written_texture_buffers[stage] >> index) & 1) != 0) {
+            MarkWrittenBuffer(binding.buffer_id, binding.cpu_addr, binding.size);
+        }
     });
 }
 
@@ -1343,6 +1353,10 @@ void BufferCache<P>::UpdateComputeTextureBuffers() {
     ForEachEnabledBit(enabled_compute_texture_buffers, [&](u32 index) {
         Binding& binding = compute_texture_buffers[index];
         binding.buffer_id = FindBuffer(binding.cpu_addr, binding.size);
+        // Mark as written if needed
+        if (((written_compute_texture_buffers >> index) & 1) != 0) {
+            MarkWrittenBuffer(binding.buffer_id, binding.cpu_addr, binding.size);
+        }
     });
 }
 
diff --git a/src/video_core/renderer_vulkan/pipeline_helper.h b/src/video_core/renderer_vulkan/pipeline_helper.h
index e2167dc4b..aaf9a735e 100644
--- a/src/video_core/renderer_vulkan/pipeline_helper.h
+++ b/src/video_core/renderer_vulkan/pipeline_helper.h
@@ -88,6 +88,7 @@ public:
         Add(VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, stage, info.constant_buffer_descriptors.size());
         Add(VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, stage, info.storage_buffers_descriptors.size());
         Add(VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, stage, info.texture_buffer_descriptors.size());
+        Add(VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, stage, info.image_buffer_descriptors.size());
         Add(VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, stage, info.texture_descriptors.size());
         Add(VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, stage, info.image_descriptors.size());
     }
@@ -126,6 +127,7 @@ inline void PushImageDescriptors(const Shader::Info& info, const VkSampler*& sam
                                  const ImageId*& image_view_ids, TextureCache& texture_cache,
                                  VKUpdateDescriptorQueue& update_descriptor_queue) {
     image_view_ids += info.texture_buffer_descriptors.size();
+    image_view_ids += info.image_buffer_descriptors.size();
     for (const auto& desc : info.texture_descriptors) {
         const VkSampler sampler{*(samplers++)};
         ImageView& image_view{texture_cache.GetImageView(*(image_view_ids++))};
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index 3d690f335..3c907ec5a 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -97,10 +97,12 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
         const u32 raw_handle{gpu_memory.Read<u32>(addr)};
         return TextureHandle(raw_handle, via_header_index);
     }};
-    for (const auto& desc : info.texture_buffer_descriptors) {
+    const auto add_image{[&](const auto& desc) {
         const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
         image_view_indices.push_back(handle.image);
-    }
+    }};
+    std::ranges::for_each(info.texture_buffer_descriptors, add_image);
+    std::ranges::for_each(info.image_buffer_descriptors, add_image);
     for (const auto& desc : info.texture_descriptors) {
         const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
         image_view_indices.push_back(handle.image);
@@ -108,24 +110,29 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
         Sampler* const sampler = texture_cache.GetComputeSampler(handle.sampler);
         samplers.push_back(sampler->Handle());
     }
-    for (const auto& desc : info.image_descriptors) {
-        const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
-        image_view_indices.push_back(handle.image);
-    }
+    std::ranges::for_each(info.image_descriptors, add_image);
+
     const std::span indices_span(image_view_indices.data(), image_view_indices.size());
     texture_cache.FillComputeImageViews(indices_span, image_view_ids);
 
     buffer_cache.UnbindComputeTextureBuffers();
     ImageId* texture_buffer_ids{image_view_ids.data()};
     size_t index{};
-    for (const auto& desc : info.texture_buffer_descriptors) {
+    const auto add_buffer{[&](const auto& desc) {
         ASSERT(desc.count == 1);
+        bool is_written{false};
+        if constexpr (std::is_same_v<decltype(desc), const Shader::ImageBufferDescriptor&>) {
+            is_written = desc.is_written;
+        }
         ImageView& image_view = texture_cache.GetImageView(*texture_buffer_ids);
         buffer_cache.BindComputeTextureBuffer(index, image_view.GpuAddr(), image_view.BufferSize(),
-                                              image_view.format);
+                                              image_view.format, is_written);
         ++texture_buffer_ids;
         ++index;
-    }
+    }};
+    std::ranges::for_each(info.texture_buffer_descriptors, add_buffer);
+    std::ranges::for_each(info.image_buffer_descriptors, add_buffer);
+
     buffer_cache.UpdateComputeBuffers();
     buffer_cache.BindHostComputeBuffers();
 
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index 23c01f24e..84720a6f9 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -175,10 +175,12 @@ void GraphicsPipeline::Configure(bool is_indexed) {
             const u32 raw_handle{gpu_memory.Read<u32>(addr)};
             return TextureHandle(raw_handle, via_header_index);
         }};
-        for (const auto& desc : info.texture_buffer_descriptors) {
+        const auto add_image{[&](const auto& desc) {
             const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
             image_view_indices.push_back(handle.image);
-        }
+        }};
+        std::ranges::for_each(info.texture_buffer_descriptors, add_image);
+        std::ranges::for_each(info.image_buffer_descriptors, add_image);
         for (const auto& desc : info.texture_descriptors) {
             const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
             image_view_indices.push_back(handle.image);
@@ -186,28 +188,33 @@ void GraphicsPipeline::Configure(bool is_indexed) {
             Sampler* const sampler{texture_cache.GetGraphicsSampler(handle.sampler)};
             samplers.push_back(sampler->Handle());
         }
-        for (const auto& desc : info.image_descriptors) {
-            const TextureHandle handle{read_handle(desc.cbuf_index, desc.cbuf_offset)};
-            image_view_indices.push_back(handle.image);
-        }
+        std::ranges::for_each(info.image_descriptors, add_image);
     }
     const std::span indices_span(image_view_indices.data(), image_view_indices.size());
     texture_cache.FillGraphicsImageViews(indices_span, image_view_ids);
 
     ImageId* texture_buffer_index{image_view_ids.data()};
     for (size_t stage = 0; stage < Maxwell::MaxShaderStage; ++stage) {
-        const Shader::Info& info{stage_infos[stage]};
-        buffer_cache.UnbindGraphicsTextureBuffers(stage);
         size_t index{};
-        for (const auto& desc : info.texture_buffer_descriptors) {
+        const auto add_buffer{[&](const auto& desc) {
             ASSERT(desc.count == 1);
-            ImageView& image_view = texture_cache.GetImageView(*texture_buffer_index);
+            bool is_written{false};
+            if constexpr (std::is_same_v<decltype(desc), const Shader::ImageBufferDescriptor&>) {
+                is_written = desc.is_written;
+            }
+            ImageView& image_view{texture_cache.GetImageView(*texture_buffer_index)};
             buffer_cache.BindGraphicsTextureBuffer(stage, index, image_view.GpuAddr(),
-                                                   image_view.BufferSize(), image_view.format);
+                                                   image_view.BufferSize(), image_view.format,
+                                                   is_written);
             ++index;
             ++texture_buffer_index;
-        }
+        }};
+        const Shader::Info& info{stage_infos[stage]};
+        buffer_cache.UnbindGraphicsTextureBuffers(stage);
+        std::ranges::for_each(info.texture_buffer_descriptors, add_buffer);
+        std::ranges::for_each(info.image_buffer_descriptors, add_buffer);
         texture_buffer_index += info.texture_descriptors.size();
+        texture_buffer_index += info.image_descriptors.size();
     }
     buffer_cache.UpdateGraphicsBuffers(is_indexed);