From ddc9622b905655fb512b58d19ea6fb10b46a04c9 Mon Sep 17 00:00:00 2001
From: PixelyIon <pixelyion@protonmail.com>
Date: Mon, 18 Apr 2022 12:16:54 +0530
Subject: [PATCH] Fix Shader Module Cache

As bindings weren't correctly handled due to the fact that `EmitSPIRV` would change the bindings, the shader module cache would not correctly function and have no cache hits in `find` and rather have them in `try_emplace` which negated any performance benefit of it. This has now been fixed by retaining the initial cache key for insertion into the cache while also storing the post-emit bindings and restoring them during a cache hit.
---
 .../main/cpp/skyline/gpu/shader_manager.cpp   | 21 ++++++++++++-------
 app/src/main/cpp/skyline/gpu/shader_manager.h | 11 ++++++++--
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/app/src/main/cpp/skyline/gpu/shader_manager.cpp b/app/src/main/cpp/skyline/gpu/shader_manager.cpp
index 6967099e..d8c59132 100644
--- a/app/src/main/cpp/skyline/gpu/shader_manager.cpp
+++ b/app/src/main/cpp/skyline/gpu/shader_manager.cpp
@@ -235,7 +235,7 @@ namespace skyline::gpu {
 
         #define RIH(member) boost::hash_combine(hash, state.runtimeInfo.member)
 
-        hash = XXH64(state.runtimeInfo.generic_input_types.data(), state.runtimeInfo.generic_input_types.size() * sizeof(u32), hash);
+        hash = XXH64(state.runtimeInfo.generic_input_types.data(), state.runtimeInfo.generic_input_types.size() * sizeof(Shader::AttributeType), hash);
         hash = XXH64(&state.runtimeInfo.previous_stage_stores.mask, sizeof(state.runtimeInfo.previous_stage_stores.mask), hash);
         RIH(convert_depth_mode);
         RIH(force_early_z);
@@ -255,11 +255,18 @@ namespace skyline::gpu {
         return hash;
     }
 
-    vk::ShaderModule ShaderManager::CompileShader(Shader::RuntimeInfo &runtimeInfo, const std::shared_ptr<ShaderProgram> &program, Shader::Backend::Bindings &bindings) {
-        auto it{shaderModuleCache.find(ShaderModuleState{program, bindings, runtimeInfo})};
-        if (it != shaderModuleCache.end())
-            return *it->second;
+    ShaderManager::ShaderModule::ShaderModule(const vk::raii::Device &device, const vk::ShaderModuleCreateInfo &createInfo, Shader::Backend::Bindings bindings) : vkModule(device, createInfo), bindings(bindings) {}
 
+    vk::ShaderModule ShaderManager::CompileShader(Shader::RuntimeInfo &runtimeInfo, const std::shared_ptr<ShaderProgram> &program, Shader::Backend::Bindings &bindings) {
+        ShaderModuleState shaderModuleState{program, bindings, runtimeInfo};
+        auto it{shaderModuleCache.find(shaderModuleState)};
+        if (it != shaderModuleCache.end()) {
+            const auto &entry{it->second};
+            bindings = entry.bindings;
+            return *entry.vkModule;
+        }
+
+        // Note: EmitSPIRV will change bindings so we explicitly have pre/post emit bindings
         auto spirv{Shader::Backend::SPIRV::EmitSPIRV(profile, runtimeInfo, program->program, bindings)};
 
         vk::ShaderModuleCreateInfo createInfo{
@@ -267,7 +274,7 @@ namespace skyline::gpu {
             .codeSize = spirv.size() * sizeof(u32),
         };
 
-        auto shaderModule{shaderModuleCache.try_emplace(ShaderModuleState{program, bindings, runtimeInfo}, gpu.vkDevice, createInfo)};
-        return *shaderModule.first->second;
+        auto shaderModule{shaderModuleCache.try_emplace(shaderModuleState, gpu.vkDevice, createInfo, bindings)};
+        return *(shaderModule.first->second.vkModule);
     }
 }
diff --git a/app/src/main/cpp/skyline/gpu/shader_manager.h b/app/src/main/cpp/skyline/gpu/shader_manager.h
index 49e90492..1da7e593 100644
--- a/app/src/main/cpp/skyline/gpu/shader_manager.h
+++ b/app/src/main/cpp/skyline/gpu/shader_manager.h
@@ -69,7 +69,7 @@ namespace skyline::gpu {
          */
         struct ShaderModuleState {
             std::shared_ptr<ShaderProgram> program;
-            Shader::Backend::Bindings bindings;
+            Shader::Backend::Bindings bindings; //!< The bindings prior to the shader being compiled
             Shader::RuntimeInfo runtimeInfo;
 
             bool operator==(const ShaderModuleState &) const;
@@ -79,7 +79,14 @@ namespace skyline::gpu {
             constexpr size_t operator()(const ShaderModuleState &state) const;
         };
 
-        std::unordered_map<ShaderModuleState, vk::raii::ShaderModule, ShaderModuleStateHash> shaderModuleCache; //!< A map from shader module state to the corresponding Vulkan shader module
+        struct ShaderModule {
+            vk::raii::ShaderModule vkModule;
+            Shader::Backend::Bindings bindings; //!< The bindings after the shader has been compiled
+
+            ShaderModule(const vk::raii::Device& device, const vk::ShaderModuleCreateInfo& createInfo, Shader::Backend::Bindings bindings);
+        };
+
+        std::unordered_map<ShaderModuleState, ShaderModule, ShaderModuleStateHash> shaderModuleCache; //!< A map from shader module state to the corresponding Vulkan shader module
 
       public:
         ShaderManager(const DeviceState &state, GPU &gpu);