11 files changed, 78 insertions, 38 deletions

diff --git a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
index 6a5243c9f..fb2031fc8 100644
--- a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
+++ b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
@@ -560,32 +560,45 @@ void VisitUsages(Info& info, IR::Inst& inst) {
     case IR::Opcode::GetCbufU32:
     case IR::Opcode::GetCbufF32:
     case IR::Opcode::GetCbufU32x2: {
-        if (const IR::Value index{inst.Arg(0)}; index.IsImmediate()) {
-            AddConstantBufferDescriptor(info, index.U32(), 1);
-        } else {
+        const IR::Value index{inst.Arg(0)};
+        const IR::Value offset{inst.Arg(1)};
+        if (!index.IsImmediate()) {
             throw NotImplementedException("Constant buffer with non-immediate index");
         }
+        AddConstantBufferDescriptor(info, index.U32(), 1);
+        u32 element_size{};
         switch (inst.GetOpcode()) {
         case IR::Opcode::GetCbufU8:
         case IR::Opcode::GetCbufS8:
             info.used_constant_buffer_types |= IR::Type::U8;
+            element_size = 1;
             break;
         case IR::Opcode::GetCbufU16:
         case IR::Opcode::GetCbufS16:
             info.used_constant_buffer_types |= IR::Type::U16;
+            element_size = 2;
             break;
         case IR::Opcode::GetCbufU32:
             info.used_constant_buffer_types |= IR::Type::U32;
+            element_size = 4;
             break;
         case IR::Opcode::GetCbufF32:
             info.used_constant_buffer_types |= IR::Type::F32;
+            element_size = 4;
             break;
         case IR::Opcode::GetCbufU32x2:
             info.used_constant_buffer_types |= IR::Type::U32x2;
+            element_size = 8;
             break;
         default:
             break;
         }
+        u32& size{info.constant_buffer_used_sizes[index.U32()]};
+        if (offset.IsImmediate()) {
+            size = std::max(size, offset.U32() + element_size);
+        } else {
+            size = 0x10'000;
+        }
         break;
     }
     case IR::Opcode::BindlessImageSampleImplicitLod:
diff --git a/src/shader_recompiler/shader_info.h b/src/shader_recompiler/shader_info.h
index d5b2ca7bc..32f8a50ea 100644
--- a/src/shader_recompiler/shader_info.h
+++ b/src/shader_recompiler/shader_info.h
@@ -197,6 +197,7 @@ struct Info {
     IR::Type used_storage_buffer_types{};
 
     u32 constant_buffer_mask{};
+    std::array<u32, MAX_CBUFS> constant_buffer_used_sizes{};
     u32 nvn_buffer_base{};
     std::bitset<16> nvn_buffer_used{};
 
diff --git a/src/video_core/buffer_cache/buffer_cache.h b/src/video_core/buffer_cache/buffer_cache.h
index d6b9eb99f..ec64f2293 100644
--- a/src/video_core/buffer_cache/buffer_cache.h
+++ b/src/video_core/buffer_cache/buffer_cache.h
@@ -44,6 +44,7 @@ MICROPROFILE_DECLARE(GPU_DownloadMemory);
 using BufferId = SlotId;
 
 using VideoCore::Surface::PixelFormat;
+using namespace Common::Literals;
 
 constexpr u32 NUM_VERTEX_BUFFERS = 32;
 constexpr u32 NUM_TRANSFORM_FEEDBACK_BUFFERS = 4;
@@ -53,7 +54,8 @@ constexpr u32 NUM_STORAGE_BUFFERS = 16;
 constexpr u32 NUM_TEXTURE_BUFFERS = 16;
 constexpr u32 NUM_STAGES = 5;
 
-using namespace Common::Literals;
+using UniformBufferSizes = std::array<std::array<u32, NUM_GRAPHICS_UNIFORM_BUFFERS>, NUM_STAGES>;
+using ComputeUniformBufferSizes = std::array<u32, NUM_COMPUTE_UNIFORM_BUFFERS>;
 
 template <typename P>
 class BufferCache {
@@ -142,9 +144,10 @@ public:
 
     void BindHostComputeBuffers();
 
-    void SetEnabledUniformBuffers(const std::array<u32, NUM_STAGES>& mask);
+    void SetUniformBuffersState(const std::array<u32, NUM_STAGES>& mask,
+                                const UniformBufferSizes* sizes);
 
-    void SetEnabledComputeUniformBuffers(u32 enabled);
+    void SetComputeUniformBufferState(u32 mask, const ComputeUniformBufferSizes* sizes);
 
     void UnbindGraphicsStorageBuffers(size_t stage);
 
@@ -384,8 +387,11 @@ private:
     std::array<Binding, NUM_STORAGE_BUFFERS> compute_storage_buffers;
     std::array<TextureBufferBinding, NUM_TEXTURE_BUFFERS> compute_texture_buffers;
 
-    std::array<u32, NUM_STAGES> enabled_uniform_buffers{};
-    u32 enabled_compute_uniform_buffers = 0;
+    std::array<u32, NUM_STAGES> enabled_uniform_buffer_masks{};
+    u32 enabled_compute_uniform_buffer_mask = 0;
+
+    const UniformBufferSizes* uniform_buffer_sizes{};
+    const ComputeUniformBufferSizes* compute_uniform_buffer_sizes{};
 
     std::array<u32, NUM_STAGES> enabled_storage_buffers{};
     std::array<u32, NUM_STAGES> written_storage_buffers{};
@@ -670,18 +676,22 @@ void BufferCache<P>::BindHostComputeBuffers() {
 }
 
 template <class P>
-void BufferCache<P>::SetEnabledUniformBuffers(const std::array<u32, NUM_STAGES>& mask) {
+void BufferCache<P>::SetUniformBuffersState(const std::array<u32, NUM_STAGES>& mask,
+                                            const UniformBufferSizes* sizes) {
     if constexpr (HAS_PERSISTENT_UNIFORM_BUFFER_BINDINGS) {
-        if (enabled_uniform_buffers != mask) {
+        if (enabled_uniform_buffer_masks != mask) {
             dirty_uniform_buffers.fill(~u32{0});
         }
     }
-    enabled_uniform_buffers = mask;
+    enabled_uniform_buffer_masks = mask;
+    uniform_buffer_sizes = sizes;
 }
 
 template <class P>
-void BufferCache<P>::SetEnabledComputeUniformBuffers(u32 enabled) {
-    enabled_compute_uniform_buffers = enabled;
+void BufferCache<P>::SetComputeUniformBufferState(u32 mask,
+                                                  const ComputeUniformBufferSizes* sizes) {
+    enabled_compute_uniform_buffer_mask = mask;
+    compute_uniform_buffer_sizes = sizes;
 }
 
 template <class P>
@@ -984,7 +994,7 @@ void BufferCache<P>::BindHostGraphicsUniformBuffers(size_t stage) {
         dirty = std::exchange(dirty_uniform_buffers[stage], 0);
     }
     u32 binding_index = 0;
-    ForEachEnabledBit(enabled_uniform_buffers[stage], [&](u32 index) {
+    ForEachEnabledBit(enabled_uniform_buffer_masks[stage], [&](u32 index) {
         const bool needs_bind = ((dirty >> index) & 1) != 0;
         BindHostGraphicsUniformBuffer(stage, index, binding_index, needs_bind);
         if constexpr (NEEDS_BIND_UNIFORM_INDEX) {
@@ -998,7 +1008,7 @@ void BufferCache<P>::BindHostGraphicsUniformBuffer(size_t stage, u32 index, u32
                                                    bool needs_bind) {
     const Binding& binding = uniform_buffers[stage][index];
     const VAddr cpu_addr = binding.cpu_addr;
-    const u32 size = binding.size;
+    const u32 size = std::min(binding.size, (*uniform_buffer_sizes)[stage][index]);
     Buffer& buffer = slot_buffers[binding.buffer_id];
     TouchBuffer(buffer);
     const bool use_fast_buffer = binding.buffer_id != NULL_BUFFER_ID &&
@@ -1113,11 +1123,11 @@ void BufferCache<P>::BindHostComputeUniformBuffers() {
         dirty_uniform_buffers.fill(~u32{0});
     }
     u32 binding_index = 0;
-    ForEachEnabledBit(enabled_compute_uniform_buffers, [&](u32 index) {
+    ForEachEnabledBit(enabled_compute_uniform_buffer_mask, [&](u32 index) {
         const Binding& binding = compute_uniform_buffers[index];
         Buffer& buffer = slot_buffers[binding.buffer_id];
         TouchBuffer(buffer);
-        const u32 size = binding.size;
+        const u32 size = std::min(binding.size, (*compute_uniform_buffer_sizes)[index]);
         SynchronizeBuffer(buffer, binding.cpu_addr, size);
 
         const u32 offset = buffer.Offset(binding.cpu_addr);
@@ -1261,7 +1271,7 @@ void BufferCache<P>::UpdateVertexBuffer(u32 index) {
 
 template <class P>
 void BufferCache<P>::UpdateUniformBuffers(size_t stage) {
-    ForEachEnabledBit(enabled_uniform_buffers[stage], [&](u32 index) {
+    ForEachEnabledBit(enabled_uniform_buffer_masks[stage], [&](u32 index) {
         Binding& binding = uniform_buffers[stage][index];
         if (binding.buffer_id) {
             // Already updated
@@ -1334,7 +1344,7 @@ void BufferCache<P>::UpdateTransformFeedbackBuffer(u32 index) {
 
 template <class P>
 void BufferCache<P>::UpdateComputeUniformBuffers() {
-    ForEachEnabledBit(enabled_compute_uniform_buffers, [&](u32 index) {
+    ForEachEnabledBit(enabled_compute_uniform_buffer_mask, [&](u32 index) {
         Binding& binding = compute_uniform_buffers[index];
         binding = NULL_BINDING;
         const auto& launch_desc = kepler_compute.launch_description;
diff --git a/src/video_core/renderer_opengl/gl_compute_pipeline.cpp b/src/video_core/renderer_opengl/gl_compute_pipeline.cpp
index 5cf5f97a9..61b6fe4b7 100644
--- a/src/video_core/renderer_opengl/gl_compute_pipeline.cpp
+++ b/src/video_core/renderer_opengl/gl_compute_pipeline.cpp
@@ -43,6 +43,8 @@ ComputePipeline::ComputePipeline(const Device& device, TextureCache& texture_cac
     : texture_cache{texture_cache_}, buffer_cache{buffer_cache_}, gpu_memory{gpu_memory_},
       kepler_compute{kepler_compute_}, program_manager{program_manager_}, info{info_},
       source_program{std::move(source_program_)}, assembly_program{std::move(assembly_program_)} {
+    std::copy_n(info.constant_buffer_used_sizes.begin(), uniform_buffer_sizes.size(),
+                uniform_buffer_sizes.begin());
 
     num_texture_buffers = AccumulateCount(info.texture_buffer_descriptors);
     num_image_buffers = AccumulateCount(info.image_buffer_descriptors);
@@ -63,7 +65,7 @@ ComputePipeline::ComputePipeline(const Device& device, TextureCache& texture_cac
 }
 
 void ComputePipeline::Configure() {
-    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_opengl/gl_compute_pipeline.h b/src/video_core/renderer_opengl/gl_compute_pipeline.h
index dd6b62ef2..b5dfb65e9 100644
--- a/src/video_core/renderer_opengl/gl_compute_pipeline.h
+++ b/src/video_core/renderer_opengl/gl_compute_pipeline.h
@@ -72,6 +72,7 @@ private:
     Shader::Info info;
     OGLProgram source_program;
     OGLAssemblyProgram assembly_program;
+    VideoCommon::ComputeUniformBufferSizes uniform_buffer_sizes{};
 
     u32 num_texture_buffers{};
     u32 num_image_buffers{};
diff --git a/src/video_core/renderer_opengl/gl_graphics_pipeline.cpp b/src/video_core/renderer_opengl/gl_graphics_pipeline.cpp
index 976897067..a5d65fdca 100644
--- a/src/video_core/renderer_opengl/gl_graphics_pipeline.cpp
+++ b/src/video_core/renderer_opengl/gl_graphics_pipeline.cpp
@@ -60,6 +60,14 @@ std::pair<GLint, GLint> TransformFeedbackEnum(u8 location) {
     UNIMPLEMENTED_MSG("index={}", index);
     return {GL_POSITION, 0};
 }
+
+struct Spec {
+    static constexpr std::array<bool, 5> enabled_stages{true, true, true, true, true};
+    static constexpr bool has_storage_buffers = true;
+    static constexpr bool has_texture_buffers = true;
+    static constexpr bool has_image_buffers = true;
+    static constexpr bool has_images = true;
+};
 } // Anonymous namespace
 
 size_t GraphicsPipelineKey::Hash() const noexcept {
@@ -100,7 +108,8 @@ GraphicsPipeline::GraphicsPipeline(const Device& device, TextureCache& texture_c
             base_uniform_bindings[stage + 1] += AccumulateCount(info.constant_buffer_descriptors);
             base_storage_bindings[stage + 1] += AccumulateCount(info.storage_buffers_descriptors);
         }
-        enabled_uniform_buffers[stage] = info.constant_buffer_mask;
+        enabled_uniform_buffer_masks[stage] = info.constant_buffer_mask;
+        std::ranges::copy(info.constant_buffer_used_sizes, uniform_buffer_sizes[stage].begin());
 
         const u32 num_tex_buffer_bindings{AccumulateCount(info.texture_buffer_descriptors)};
         num_texture_buffers[stage] += num_tex_buffer_bindings;
@@ -130,14 +139,6 @@ GraphicsPipeline::GraphicsPipeline(const Device& device, TextureCache& texture_c
     }
 }
 
-struct Spec {
-    static constexpr std::array<bool, 5> enabled_stages{true, true, true, true, true};
-    static constexpr bool has_storage_buffers = true;
-    static constexpr bool has_texture_buffers = true;
-    static constexpr bool has_image_buffers = true;
-    static constexpr bool has_images = true;
-};
-
 void GraphicsPipeline::Configure(bool is_indexed) {
     std::array<ImageId, MAX_TEXTURES + MAX_IMAGES> image_view_ids;
     std::array<u32, MAX_TEXTURES + MAX_IMAGES> image_view_indices;
@@ -147,7 +148,7 @@ void GraphicsPipeline::Configure(bool is_indexed) {
 
     texture_cache.SynchronizeGraphicsDescriptors();
 
-    buffer_cache.SetEnabledUniformBuffers(enabled_uniform_buffers);
+    buffer_cache.SetUniformBuffersState(enabled_uniform_buffer_masks, &uniform_buffer_sizes);
     buffer_cache.runtime.SetBaseUniformBindings(base_uniform_bindings);
     buffer_cache.runtime.SetBaseStorageBindings(base_storage_bindings);
     buffer_cache.runtime.SetEnableStorageBuffers(use_storage_buffers);
diff --git a/src/video_core/renderer_opengl/gl_graphics_pipeline.h b/src/video_core/renderer_opengl/gl_graphics_pipeline.h
index bf33ce604..508fad5bb 100644
--- a/src/video_core/renderer_opengl/gl_graphics_pipeline.h
+++ b/src/video_core/renderer_opengl/gl_graphics_pipeline.h
@@ -99,7 +99,8 @@ private:
     u32 enabled_stages_mask{};
 
     std::array<Shader::Info, 5> stage_infos{};
-    std::array<u32, 5> enabled_uniform_buffers{};
+    std::array<u32, 5> enabled_uniform_buffer_masks{};
+    VideoCommon::UniformBufferSizes uniform_buffer_sizes{};
     std::array<u32, 5> base_uniform_bindings{};
     std::array<u32, 5> base_storage_bindings{};
     std::array<u32, 5> num_texture_buffers{};
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;