From cb6e5ef0488ef61e5909172e49f802e14e6848e5 Mon Sep 17 00:00:00 2001
From: Ameer J <52414509+ameerj@users.noreply.github.com>
Date: Sun, 26 Nov 2023 22:32:32 -0500
Subject: [PATCH] Vulkan: Max UBO checks

---
 src/video_core/renderer_opengl/gl_shader_cache.cpp  |  8 ++++----
 .../renderer_vulkan/vk_pipeline_cache.cpp           | 13 ++++++++++---
 src/video_core/vulkan_common/vulkan_device.h        |  5 +++++
 3 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/video_core/renderer_opengl/gl_shader_cache.cpp b/src/video_core/renderer_opengl/gl_shader_cache.cpp
index 226faccdfc..d972976b37 100644
--- a/src/video_core/renderer_opengl/gl_shader_cache.cpp
+++ b/src/video_core/renderer_opengl/gl_shader_cache.cpp
@@ -581,10 +581,10 @@ std::unique_ptr<ComputePipeline> ShaderCache::CreateComputePipeline(
 
     auto program{TranslateProgram(pools.inst, pools.block, env, cfg, host_info)};
     const u32 num_storage_buffers{Shader::NumDescriptors(program.info.storage_buffers_descriptors)};
-    Shader::RuntimeInfo info{};
-    info.glasm_use_storage_buffers = num_storage_buffers <= device.GetMaxGLASMStorageBufferBlocks();
-    info.max_num_cbufs = device.GetMaxUniformBuffers(program.stage);
-
+    const Shader::RuntimeInfo info{
+        .glasm_use_storage_buffers = num_storage_buffers <= device.GetMaxGLASMStorageBufferBlocks(),
+        .max_num_cbufs = device.GetMaxUniformBuffers(program.stage),
+    };
     std::string code{};
     std::vector<u32> code_spirv;
     switch (device.GetShaderBackend()) {
diff --git a/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp b/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
index 89b455bff8..b7bda9e7ef 100644
--- a/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
+++ b/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
@@ -145,7 +145,8 @@ Shader::AttributeType AttributeType(const FixedPipelineState& state, size_t inde
 Shader::RuntimeInfo MakeRuntimeInfo(std::span<const Shader::IR::Program> programs,
                                     const GraphicsPipelineCacheKey& key,
                                     const Shader::IR::Program& program,
-                                    const Shader::IR::Program* previous_program) {
+                                    const Shader::IR::Program* previous_program,
+                                    u32 max_num_cbufs) {
     Shader::RuntimeInfo info;
     if (previous_program) {
         info.previous_stage_stores = previous_program->info.stores;
@@ -261,6 +262,7 @@ Shader::RuntimeInfo MakeRuntimeInfo(std::span<const Shader::IR::Program> program
     }
     info.force_early_z = key.state.early_z != 0;
     info.y_negate = key.state.y_negate != 0;
+    info.max_num_cbufs = max_num_cbufs;
     return info;
 }
 
@@ -653,6 +655,7 @@ std::unique_ptr<GraphicsPipeline> PipelineCache::CreateGraphicsPipeline(
 
     const Shader::IR::Program* previous_stage{};
     Shader::Backend::Bindings binding;
+    const u32 max_num_cbufs{static_cast<u32>(device.GetMaxPerStageUniformBuffers())};
     for (size_t index = uses_vertex_a && uses_vertex_b ? 1 : 0; index < Maxwell::MaxShaderProgram;
          ++index) {
         const bool is_emulated_stage = layer_source_program != nullptr &&
@@ -666,7 +669,8 @@ std::unique_ptr<GraphicsPipeline> PipelineCache::CreateGraphicsPipeline(
         const size_t stage_index{index - 1};
         infos[stage_index] = &program.info;
 
-        const auto runtime_info{MakeRuntimeInfo(programs, key, program, previous_stage)};
+        const auto runtime_info{
+            MakeRuntimeInfo(programs, key, program, previous_stage, max_num_cbufs)};
         ConvertLegacyToGeneric(program, runtime_info);
         const std::vector<u32> code{EmitSPIRV(profile, runtime_info, program, binding)};
         device.SaveShader(code);
@@ -762,7 +766,10 @@ std::unique_ptr<ComputePipeline> PipelineCache::CreateComputePipeline(
     }
 
     auto program{TranslateProgram(pools.inst, pools.block, env, cfg, host_info)};
-    const std::vector<u32> code{EmitSPIRV(profile, program)};
+    const Shader::RuntimeInfo info{
+        .max_num_cbufs = static_cast<u32>(device.GetMaxPerStageUniformBuffers()),
+    };
+    const std::vector<u32> code{EmitSPIRV(profile, info, program)};
     device.SaveShader(code);
     vk::ShaderModule spv_module{BuildShader(device, code)};
     if (device.HasDebuggingToolAttached()) {
diff --git a/src/video_core/vulkan_common/vulkan_device.h b/src/video_core/vulkan_common/vulkan_device.h
index 4f3846345c..75882334ab 100644
--- a/src/video_core/vulkan_common/vulkan_device.h
+++ b/src/video_core/vulkan_common/vulkan_device.h
@@ -304,6 +304,11 @@ public:
         return properties.properties.limits.maxComputeSharedMemorySize;
     }
 
+    /// Returns the maximum number of uniform buffers allowed per stage.
+    VkDeviceSize GetMaxPerStageUniformBuffers() const {
+        return properties.properties.limits.maxPerStageDescriptorUniformBuffers;
+    }
+
     /// Returns float control properties of the device.
     const VkPhysicalDeviceFloatControlsPropertiesKHR& FloatControlProperties() const {
         return properties.float_controls;