4 files changed, 18 insertions, 7 deletions

diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index 168ffa7e9..ca59042ff 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -2,6 +2,7 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <vector>
 
 #include <boost/container/small_vector.hpp>
@@ -27,6 +28,9 @@ ComputePipeline::ComputePipeline(const Device& device_, DescriptorPool& descript
                                  vk::ShaderModule spv_module_)
     : device{device_}, update_descriptor_queue{update_descriptor_queue_}, info{info_},
       spv_module(std::move(spv_module_)) {
+    std::copy_n(info.constant_buffer_used_sizes.begin(), uniform_buffer_sizes.size(),
+                uniform_buffer_sizes.begin());
+
     auto func{[this, &descriptor_pool] {
         DescriptorLayoutBuilder builder{device.GetLogical()};
         builder.Add(info, VK_SHADER_STAGE_COMPUTE_BIT);
@@ -75,7 +79,7 @@ void ComputePipeline::Configure(Tegra::Engines::KeplerCompute& kepler_compute,
                                 BufferCache& buffer_cache, TextureCache& texture_cache) {
     update_descriptor_queue.Acquire();
 
-    buffer_cache.SetEnabledComputeUniformBuffers(info.constant_buffer_mask);
+    buffer_cache.SetComputeUniformBufferState(info.constant_buffer_mask, &uniform_buffer_sizes);
     buffer_cache.UnbindComputeStorageBuffers();
     size_t ssbo_index{};
     for (const auto& desc : info.storage_buffers_descriptors) {
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.h b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
index a560e382e..a6043866d 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
@@ -44,6 +44,8 @@ private:
     VKUpdateDescriptorQueue& update_descriptor_queue;
     Shader::Info info;
 
+    VideoCommon::ComputeUniformBufferSizes uniform_buffer_sizes{};
+
     vk::ShaderModule spv_module;
     vk::DescriptorSetLayout descriptor_set_layout;
     DescriptorAllocator descriptor_allocator;
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
index d381109d6..627ca0158 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.cpp
@@ -218,10 +218,14 @@ GraphicsPipeline::GraphicsPipeline(Tegra::Engines::Maxwell3D& maxwell3d_,
       update_descriptor_queue{update_descriptor_queue_}, spv_modules{std::move(stages)} {
     std::ranges::transform(infos, stage_infos.begin(),
                            [](const Shader::Info* info) { return info ? *info : Shader::Info{}; });
-    std::ranges::transform(infos, enabled_uniform_buffers.begin(), [](const Shader::Info* info) {
-        return info ? info->constant_buffer_mask : 0;
-    });
-
+    for (size_t stage = 0; stage < NUM_STAGES; ++stage) {
+        const Shader::Info* const info{infos[stage]};
+        if (!info) {
+            continue;
+        }
+        enabled_uniform_buffer_masks[stage] = info->constant_buffer_mask;
+        std::ranges::copy(info->constant_buffer_used_sizes, uniform_buffer_sizes[stage].begin());
+    }
     auto func{[this, &render_pass_cache, &descriptor_pool] {
         DescriptorLayoutBuilder builder{MakeBuilder(device, stage_infos)};
         descriptor_set_layout = builder.CreateDescriptorSetLayout();
@@ -262,7 +266,7 @@ void GraphicsPipeline::ConfigureImpl(bool is_indexed) {
 
     texture_cache.SynchronizeGraphicsDescriptors();
 
-    buffer_cache.SetEnabledUniformBuffers(enabled_uniform_buffers);
+    buffer_cache.SetUniformBuffersState(enabled_uniform_buffer_masks, &uniform_buffer_sizes);
 
     const auto& regs{maxwell3d.regs};
     const bool via_header_index{regs.sampler_index == Maxwell::SamplerIndex::ViaHeaderIndex};
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
index 4068a0edc..8c81c28a8 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
@@ -130,7 +130,8 @@ private:
     std::array<vk::ShaderModule, NUM_STAGES> spv_modules;
 
     std::array<Shader::Info, NUM_STAGES> stage_infos;
-    std::array<u32, 5> enabled_uniform_buffers{};
+    std::array<u32, 5> enabled_uniform_buffer_masks{};
+    VideoCommon::UniformBufferSizes uniform_buffer_sizes{};
 
     vk::DescriptorSetLayout descriptor_set_layout;
     DescriptorAllocator descriptor_allocator;