From 0c7081290ac3c38d29b43caf740b41e2a9e94566 Mon Sep 17 00:00:00 2001
From: Vinicius Rangel <me@viniciusrangel.dev>
Date: Sun, 8 Dec 2024 14:05:31 -0300
Subject: [PATCH] devtools: patch shader at runtime

---
 src/core/debug_state.cpp                      |  7 +-
 src/core/debug_state.h                        | 19 ++++--
 src/core/devtools/widget/shader_list.cpp      | 65 +++++++++++++------
 src/core/devtools/widget/shader_list.h        | 14 ++--
 .../renderer_vulkan/vk_compute_pipeline.cpp   |  2 +-
 .../renderer_vulkan/vk_compute_pipeline.h     | 24 ++++++-
 .../renderer_vulkan/vk_graphics_pipeline.h    |  2 +
 .../renderer_vulkan/vk_pipeline_cache.cpp     | 62 ++++++++++++++----
 .../renderer_vulkan/vk_pipeline_cache.h       | 27 ++++++--
 src/video_core/renderer_vulkan/vk_presenter.h |  4 ++
 .../renderer_vulkan/vk_rasterizer.h           |  4 ++
 11 files changed, 170 insertions(+), 60 deletions(-)

diff --git a/src/core/debug_state.cpp b/src/core/debug_state.cpp
index 4b7cd181b98..8123c0b38f5 100644
--- a/src/core/debug_state.cpp
+++ b/src/core/debug_state.cpp
@@ -177,9 +177,10 @@ void DebugStateImpl::PushRegsDump(uintptr_t base_addr, uintptr_t header_addr,
     }
 }
 
-void DebugStateImpl::CollectShader(const std::string& name, std::span<const u32> spv,
-                                   std::span<const u32> raw_code, std::span<const u32> patch_spv) {
-    shader_dump_list.emplace_back(name, std::vector<u32>{spv.begin(), spv.end()},
+void DebugStateImpl::CollectShader(const std::string& name, vk::ShaderModule module,
+                                   std::span<const u32> spv, std::span<const u32> raw_code,
+                                   std::span<const u32> patch_spv) {
+    shader_dump_list.emplace_back(name, module, std::vector<u32>{spv.begin(), spv.end()},
                                   std::vector<u32>{raw_code.begin(), raw_code.end()},
                                   std::vector<u32>{patch_spv.begin(), patch_spv.end()});
 }
diff --git a/src/core/debug_state.h b/src/core/debug_state.h
index 5e905bb5a52..4f47eaf0853 100644
--- a/src/core/debug_state.h
+++ b/src/core/debug_state.h
@@ -12,7 +12,7 @@
 
 #include "common/types.h"
 #include "video_core/amdgpu/liverpool.h"
-#include "video_core/renderer_vulkan/vk_pipeline_cache.h"
+#include "video_core/renderer_vulkan/vk_graphics_pipeline.h"
 
 #ifdef _WIN32
 #ifndef WIN32_LEAN_AND_MEAN
@@ -76,6 +76,7 @@ struct FrameDump {
 
 struct ShaderDump {
     std::string name;
+    vk::ShaderModule module;
 
     std::vector<u32> spv;
     std::vector<u32> isa;
@@ -83,19 +84,22 @@ struct ShaderDump {
     std::vector<u32> patch_spv;
     std::string patch_source{};
 
+    bool loaded_data = false;
+    bool is_patched = false;
     std::string cache_spv_disasm{};
     std::string cache_isa_disasm{};
     std::string cache_patch_disasm{};
 
-    ShaderDump(std::string name, std::vector<u32> spv, std::vector<u32> isa,
-               std::vector<u32> patch_spv)
-        : name(std::move(name)), spv(std::move(spv)), isa(std::move(isa)),
+    ShaderDump(std::string name, vk::ShaderModule module, std::vector<u32> spv,
+               std::vector<u32> isa, std::vector<u32> patch_spv)
+        : name(std::move(name)), module(module), spv(std::move(spv)), isa(std::move(isa)),
           patch_spv(std::move(patch_spv)) {}
 
     ShaderDump(const ShaderDump& other) = delete;
     ShaderDump(ShaderDump&& other) noexcept
-        : name{std::move(other.name)}, spv{std::move(other.spv)}, isa{std::move(other.isa)},
-          patch_spv{std::move(other.patch_spv)}, patch_source{std::move(other.patch_source)},
+        : name{std::move(other.name)}, module{std::move(other.module)}, spv{std::move(other.spv)},
+          isa{std::move(other.isa)}, patch_spv{std::move(other.patch_spv)},
+          patch_source{std::move(other.patch_source)},
           cache_spv_disasm{std::move(other.cache_spv_disasm)},
           cache_isa_disasm{std::move(other.cache_isa_disasm)},
           cache_patch_disasm{std::move(other.cache_patch_disasm)} {}
@@ -104,6 +108,7 @@ struct ShaderDump {
         if (this == &other)
             return *this;
         name = std::move(other.name);
+        module = std::move(other.module);
         spv = std::move(other.spv);
         isa = std::move(other.isa);
         patch_spv = std::move(other.patch_spv);
@@ -198,7 +203,7 @@ class DebugStateImpl {
     void PushRegsDump(uintptr_t base_addr, uintptr_t header_addr,
                       const AmdGpu::Liverpool::Regs& regs, bool is_compute = false);
 
-    void CollectShader(const std::string& name, std::span<const u32> spv,
+    void CollectShader(const std::string& name, vk::ShaderModule module, std::span<const u32> spv,
                        std::span<const u32> raw_code, std::span<const u32> patch_spv);
 };
 } // namespace DebugStateType
diff --git a/src/core/devtools/widget/shader_list.cpp b/src/core/devtools/widget/shader_list.cpp
index 1eea21934a8..be743fa8e84 100644
--- a/src/core/devtools/widget/shader_list.cpp
+++ b/src/core/devtools/widget/shader_list.cpp
@@ -16,6 +16,7 @@
 #include "imgui/imgui_std.h"
 #include "sdl_window.h"
 #include "video_core/renderer_vulkan/vk_presenter.h"
+#include "video_core/renderer_vulkan/vk_rasterizer.h"
 
 extern std::unique_ptr<Vulkan::Presenter> presenter;
 
@@ -35,9 +36,20 @@ ShaderList::Selection::~Selection() {
     presenter->GetWindow().ReleaseKeyboard();
 }
 
+void ShaderList::Selection::ReloadShader(DebugStateType::ShaderDump& value) {
+    auto& spv = value.is_patched ? value.patch_spv : value.spv;
+    if (spv.empty()) {
+        return;
+    }
+    auto& cache = presenter->GetRasterizer().GetPipelineCache();
+    if (const auto m = cache.ReplaceShader(value.module, spv); m) {
+        value.module = *m;
+    }
+}
+
 bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
-    if (!loaded_data) {
-        loaded_data = true;
+    if (!value.loaded_data) {
+        value.loaded_data = true;
         if (value.cache_isa_disasm.empty()) {
             value.cache_isa_disasm = RunDisassembler(Options.disassembler_cli_isa, value.isa);
         }
@@ -58,12 +70,11 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
                 std::string{std::istreambuf_iterator{file}, std::istreambuf_iterator<char>{}};
         }
 
-        if (value.patch_spv.empty()) { // No patch
-            showing_patch = false;
+        value.is_patched = !value.patch_spv.empty();
+        if (!value.is_patched) { // No patch
             isa_editor.SetText(value.cache_isa_disasm);
             glsl_editor.SetText(value.cache_spv_disasm);
         } else {
-            showing_patch = true;
             isa_editor.SetText(value.cache_patch_disasm);
             isa_editor.SetLanguageDefinition(TextEditor::LanguageDefinition::SPIRV());
             glsl_editor.SetText(value.patch_source);
@@ -73,6 +84,7 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
 
     char name[64];
     snprintf(name, sizeof(name), "Shader %s", value.name.c_str());
+    SetNextWindowSize({450.0f, 600.0f}, ImGuiCond_FirstUseEver);
     if (!Begin(name, &open, ImGuiWindowFlags_NoNav)) {
         End();
         return open;
@@ -80,8 +92,8 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
 
     Text("%s", value.name.c_str());
     SameLine(0.0f, 7.0f);
-    if (Checkbox("Enable patch", &showing_patch)) {
-        if (showing_patch) {
+    if (Checkbox("Enable patch", &value.is_patched)) {
+        if (value.is_patched) {
             if (value.patch_source.empty()) {
                 value.patch_source = value.cache_spv_disasm;
             }
@@ -89,15 +101,19 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
             isa_editor.SetLanguageDefinition(TextEditor::LanguageDefinition::SPIRV());
             glsl_editor.SetText(value.patch_source);
             glsl_editor.SetReadOnly(false);
+            if (!value.patch_spv.empty()) {
+                ReloadShader(value);
+            }
         } else {
             isa_editor.SetText(value.cache_isa_disasm);
             isa_editor.SetLanguageDefinition(TextEditor::LanguageDefinition());
             glsl_editor.SetText(value.cache_spv_disasm);
             glsl_editor.SetReadOnly(true);
+            ReloadShader(value);
         }
     }
 
-    if (showing_patch) {
+    if (value.is_patched) {
         if (BeginCombo("Shader type", showing_bin ? "SPIRV" : "GLSL",
                        ImGuiComboFlags_WidthFitPreview)) {
             if (Selectable("GLSL")) {
@@ -121,7 +137,7 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
         }
     }
 
-    if (showing_patch) {
+    if (value.is_patched) {
         bool save = false;
         bool compile = false;
         SameLine(0.0f, 3.0f);
@@ -172,6 +188,7 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
                                                     res);
                     } else {
                         isa_editor.SetText(value.cache_patch_disasm);
+                        ReloadShader(value);
                     }
                 }
             }
@@ -179,7 +196,7 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
     }
 
     if (showing_bin) {
-        isa_editor.Render(showing_patch ? "SPIRV" : "ISA", GetContentRegionAvail());
+        isa_editor.Render(value.is_patched ? "SPIRV" : "ISA", GetContentRegionAvail());
     } else {
         glsl_editor.Render("GLSL", GetContentRegionAvail());
     }
@@ -189,6 +206,16 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
 }
 
 void ShaderList::Draw() {
+    for (auto it = open_shaders.begin(); it != open_shaders.end();) {
+        auto& selection = *it;
+        auto& shader = DebugState.shader_dump_list[selection.index];
+        if (!selection.DrawShader(shader)) {
+            it = open_shaders.erase(it);
+        } else {
+            ++it;
+        }
+    }
+
     SetNextWindowSize({500.0f, 600.0f}, ImGuiCond_FirstUseEver);
     if (!Begin("Shader list", &open)) {
         End();
@@ -201,20 +228,16 @@ void ShaderList::Draw() {
         return;
     }
 
-    for (auto it = open_shaders.begin(); it != open_shaders.end();) {
-        auto& selection = *it;
-        auto& shader = DebugState.shader_dump_list[selection.index];
-        if (!selection.DrawShader(shader)) {
-            it = open_shaders.erase(it);
-        } else {
-            ++it;
-        }
-    }
-
     auto width = GetContentRegionAvail().x;
     int i = 0;
     for (const auto& shader : DebugState.shader_dump_list) {
-        if (ButtonEx(shader.name.c_str(), {width, 20.0f}, ImGuiButtonFlags_NoHoveredOnFocus)) {
+        char name[128];
+        if (shader.patch_spv.empty()) {
+            snprintf(name, sizeof(name), "%s", shader.name.c_str());
+        } else {
+            snprintf(name, sizeof(name), "%s (PATCH)", shader.name.c_str());
+        }
+        if (ButtonEx(name, {width, 20.0f}, ImGuiButtonFlags_NoHoveredOnFocus)) {
             open_shaders.emplace_back(i);
         }
         i++;
diff --git a/src/core/devtools/widget/shader_list.h b/src/core/devtools/widget/shader_list.h
index 0d4dd5ce8ca..2534ded3517 100644
--- a/src/core/devtools/widget/shader_list.h
+++ b/src/core/devtools/widget/shader_list.h
@@ -12,21 +12,21 @@ namespace Core::Devtools::Widget {
 
 class ShaderList {
     struct Selection {
+        explicit Selection(int index);
+        ~Selection();
+
+        void ReloadShader(DebugStateType::ShaderDump& value);
+
+        bool DrawShader(DebugStateType::ShaderDump& value);
+
         int index;
         TextEditor isa_editor{};
         TextEditor glsl_editor{};
         bool open = true;
-        bool loaded_data = false;
         bool showing_bin = false;
-        bool showing_patch = false;
 
         std::filesystem::path patch_path;
         std::filesystem::path patch_bin_path;
-
-        explicit Selection(int index);
-        ~Selection();
-
-        bool DrawShader(DebugStateType::ShaderDump& value);
     };
 
     std::vector<Selection> open_shaders{};
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
index 09d4e4195f4..8d495ab0608 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.cpp
@@ -13,7 +13,7 @@ namespace Vulkan {
 
 ComputePipeline::ComputePipeline(const Instance& instance_, Scheduler& scheduler_,
                                  DescriptorHeap& desc_heap_, vk::PipelineCache pipeline_cache,
-                                 u64 compute_key_, const Shader::Info& info_,
+                                 ComputePipelineKey compute_key_, const Shader::Info& info_,
                                  vk::ShaderModule module)
     : Pipeline{instance_, scheduler_, desc_heap_, pipeline_cache, true}, compute_key{compute_key_} {
     auto& info = stages[int(Shader::Stage::Compute)];
diff --git a/src/video_core/renderer_vulkan/vk_compute_pipeline.h b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
index ca429b58d5e..1c28e461c0a 100644
--- a/src/video_core/renderer_vulkan/vk_compute_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_compute_pipeline.h
@@ -17,15 +17,33 @@ class Instance;
 class Scheduler;
 class DescriptorHeap;
 
+struct ComputePipelineKey {
+    size_t value;
+
+    friend bool operator==(const ComputePipelineKey& lhs, const ComputePipelineKey& rhs) {
+        return lhs.value == rhs.value;
+    }
+    friend bool operator!=(const ComputePipelineKey& lhs, const ComputePipelineKey& rhs) {
+        return !(lhs == rhs);
+    }
+};
+
 class ComputePipeline : public Pipeline {
 public:
     ComputePipeline(const Instance& instance, Scheduler& scheduler, DescriptorHeap& desc_heap,
-                    vk::PipelineCache pipeline_cache, u64 compute_key, const Shader::Info& info,
-                    vk::ShaderModule module);
+                    vk::PipelineCache pipeline_cache, ComputePipelineKey compute_key,
+                    const Shader::Info& info, vk::ShaderModule module);
     ~ComputePipeline();
 
 private:
-    u64 compute_key;
+    ComputePipelineKey compute_key;
 };
 
 } // namespace Vulkan
+
+template <>
+struct std::hash<Vulkan::ComputePipelineKey> {
+    std::size_t operator()(const Vulkan::ComputePipelineKey& key) const noexcept {
+        return std::hash<size_t>{}(key.value);
+    }
+};
diff --git a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
index 4f4abfd16de..8f0c372457d 100644
--- a/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
+++ b/src/video_core/renderer_vulkan/vk_graphics_pipeline.h
@@ -1,6 +1,8 @@
 // SPDX-FileCopyrightText: Copyright 2024 shadPS4 Emulator Project
 // SPDX-License-Identifier: GPL-2.0-or-later
 
+#pragma once
+
 #include <xxhash.h>
 
 #include "common/types.h"
diff --git a/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp b/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
index 2cd3a3a815e..7f573600e8f 100644
--- a/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
+++ b/src/video_core/renderer_vulkan/vk_pipeline_cache.cpp
@@ -184,10 +184,18 @@ const GraphicsPipeline* PipelineCache::GetGraphicsPipeline() {
     }
     const auto [it, is_new] = graphics_pipelines.try_emplace(graphics_key);
     if (is_new) {
-        it.value() = graphics_pipeline_pool.Create(instance, scheduler, desc_heap, graphics_key,
-                                                   *pipeline_cache, infos, modules);
+        it.value() = std::make_unique<GraphicsPipeline>(
+            instance, scheduler, desc_heap, graphics_key, *pipeline_cache, infos, modules);
+        if (Config::collectShadersForDebug()) {
+            for (auto stage = 0; stage < MaxShaderStages; ++stage) {
+                if (infos[stage]) {
+                    auto& m = modules[stage];
+                    module_related_pipelines[m].emplace_back(graphics_key);
+                }
+            }
+        }
     }
-    return it->second;
+    return it->second.get();
 }
 
 const ComputePipeline* PipelineCache::GetComputePipeline() {
@@ -196,10 +204,14 @@ const ComputePipeline* PipelineCache::GetComputePipeline() {
     }
     const auto [it, is_new] = compute_pipelines.try_emplace(compute_key);
     if (is_new) {
-        it.value() = compute_pipeline_pool.Create(instance, scheduler, desc_heap, *pipeline_cache,
-                                                  compute_key, *infos[0], modules[0]);
+        it.value() = std::make_unique<ComputePipeline>(
+            instance, scheduler, desc_heap, *pipeline_cache, compute_key, *infos[0], modules[0]);
+        if (Config::collectShadersForDebug()) {
+            auto& m = modules[0];
+            module_related_pipelines[m].emplace_back(compute_key);
+        }
     }
-    return it->second;
+    return it->second.get();
 }
 
 bool PipelineCache::RefreshGraphicsKey() {
@@ -392,7 +404,7 @@ bool PipelineCache::RefreshComputeKey() {
     Shader::Backend::Bindings binding{};
     const auto* cs_pgm = &liverpool->regs.cs_program;
     const auto cs_params = Liverpool::GetParams(*cs_pgm);
-    std::tie(infos[0], modules[0], compute_key) =
+    std::tie(infos[0], modules[0], compute_key.value) =
         GetProgram(Shader::Stage::Compute, cs_params, binding);
     return true;
 }
@@ -422,7 +434,7 @@ vk::ShaderModule PipelineCache::CompileModule(Shader::Info& info,
     const auto name = fmt::format("{}_{:#018x}_{}", info.stage, info.pgm_hash, perm_idx);
     Vulkan::SetObjectName(instance.GetDevice(), module, name);
     if (Config::collectShadersForDebug()) {
-        DebugState.CollectShader(name, spv, code, patch ? *patch : std::span<const u32>{});
+        DebugState.CollectShader(name, module, spv, code, patch ? *patch : std::span<const u32>{});
     }
     return module;
 }
@@ -432,16 +444,16 @@ std::tuple<const Shader::Info*, vk::ShaderModule, u64> PipelineCache::GetProgram
     const auto runtime_info = BuildRuntimeInfo(stage);
     auto [it_pgm, new_program] = program_cache.try_emplace(params.hash);
     if (new_program) {
-        Program* program = program_pool.Create(stage, params);
+        it_pgm.value() = std::make_unique<Program>(stage, params);
+        auto& program = it_pgm.value();
         auto start = binding;
         const auto module = CompileModule(program->info, runtime_info, params.code, 0, binding);
         const auto spec = Shader::StageSpecialization(program->info, runtime_info, start);
         program->AddPermut(module, std::move(spec));
-        it_pgm.value() = program;
         return std::make_tuple(&program->info, module, HashCombine(params.hash, 0));
     }
 
-    Program* program = it_pgm->second;
+    auto& program = it_pgm.value();
     auto& info = program->info;
     info.RefreshFlatBuf();
     const auto spec = Shader::StageSpecialization(info, runtime_info, binding);
@@ -461,6 +473,34 @@ std::tuple<const Shader::Info*, vk::ShaderModule, u64> PipelineCache::GetProgram
     return std::make_tuple(&info, module, HashCombine(params.hash, perm_idx));
 }
 
+std::optional<vk::ShaderModule> PipelineCache::ReplaceShader(vk::ShaderModule module,
+                                                             std::span<const u32> spv_code) {
+    std::optional<vk::ShaderModule> new_module{};
+    for (const auto& [_, program] : program_cache) {
+        for (auto& m : program->modules) {
+            if (m.module == module) {
+                const auto& d = instance.GetDevice();
+                d.destroyShaderModule(m.module);
+                m.module = CompileSPV(spv_code, d);
+                new_module = m.module;
+            }
+        }
+    }
+    if (module_related_pipelines.contains(module)) {
+        auto& pipeline_keys = module_related_pipelines[module];
+        for (auto& key : pipeline_keys) {
+            if (std::holds_alternative<GraphicsPipelineKey>(key)) {
+                auto& graphics_key = std::get<GraphicsPipelineKey>(key);
+                graphics_pipelines.erase(graphics_key);
+            } else if (std::holds_alternative<ComputePipelineKey>(key)) {
+                auto& compute_key = std::get<ComputePipelineKey>(key);
+                compute_pipelines.erase(compute_key);
+            }
+        }
+    }
+    return new_module;
+}
+
 void PipelineCache::DumpShader(std::span<const u32> code, u64 hash, Shader::Stage stage,
                                size_t perm_idx, std::string_view ext) {
     if (!Config::dumpShaders()) {
diff --git a/src/video_core/renderer_vulkan/vk_pipeline_cache.h b/src/video_core/renderer_vulkan/vk_pipeline_cache.h
index 662bcbd80f1..829481d07d6 100644
--- a/src/video_core/renderer_vulkan/vk_pipeline_cache.h
+++ b/src/video_core/renderer_vulkan/vk_pipeline_cache.h
@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <variant>
 #include <tsl/robin_map.h>
 #include "shader_recompiler/profile.h"
 #include "shader_recompiler/recompiler.h"
@@ -11,6 +12,13 @@
 #include "video_core/renderer_vulkan/vk_graphics_pipeline.h"
 #include "video_core/renderer_vulkan/vk_resource_pool.h"
 
+template <>
+struct std::hash<vk::ShaderModule> {
+    std::size_t operator()(const vk::ShaderModule& module) const noexcept {
+        return std::hash<size_t>{}(reinterpret_cast<size_t>((VkShaderModule)module));
+    }
+};
+
 namespace Shader {
 struct Info;
 }
@@ -50,6 +58,9 @@ class PipelineCache {
     std::tuple<const Shader::Info*, vk::ShaderModule, u64> GetProgram(
         Shader::Stage stage, Shader::ShaderParams params, Shader::Backend::Bindings& binding);
 
+    std::optional<vk::ShaderModule> ReplaceShader(vk::ShaderModule module,
+                                                  std::span<const u32> spv_code);
+
 private:
     bool RefreshGraphicsKey();
     bool RefreshComputeKey();
@@ -72,16 +83,18 @@ class PipelineCache {
     vk::UniquePipelineLayout pipeline_layout;
     Shader::Profile profile{};
     Shader::Pools pools;
-    tsl::robin_map<size_t, Program*> program_cache;
-    Common::ObjectPool<Program> program_pool;
-    Common::ObjectPool<GraphicsPipeline> graphics_pipeline_pool;
-    Common::ObjectPool<ComputePipeline> compute_pipeline_pool;
-    tsl::robin_map<size_t, ComputePipeline*> compute_pipelines;
-    tsl::robin_map<GraphicsPipelineKey, GraphicsPipeline*> graphics_pipelines;
+    tsl::robin_map<size_t, std::unique_ptr<Program>> program_cache;
+    tsl::robin_map<ComputePipelineKey, std::unique_ptr<ComputePipeline>> compute_pipelines;
+    tsl::robin_map<GraphicsPipelineKey, std::unique_ptr<GraphicsPipeline>> graphics_pipelines;
     std::array<const Shader::Info*, MaxShaderStages> infos{};
     std::array<vk::ShaderModule, MaxShaderStages> modules{};
     GraphicsPipelineKey graphics_key{};
-    u64 compute_key{};
+    ComputePipelineKey compute_key{};
+
+    // Only if Config::collectShadersForDebug()
+    tsl::robin_map<vk::ShaderModule,
+                   std::vector<std::variant<GraphicsPipelineKey, ComputePipelineKey>>>
+        module_related_pipelines;
 };
 
 } // namespace Vulkan
diff --git a/src/video_core/renderer_vulkan/vk_presenter.h b/src/video_core/renderer_vulkan/vk_presenter.h
index 98b3b5dae73..4c29af0f0d6 100644
--- a/src/video_core/renderer_vulkan/vk_presenter.h
+++ b/src/video_core/renderer_vulkan/vk_presenter.h
@@ -94,6 +94,10 @@ class Presenter {
         draw_scheduler.Flush(info);
     }
 
+    Rasterizer& GetRasterizer() const {
+        return *rasterizer.get();
+    }
+
 private:
     void CreatePostProcessPipeline();
     Frame* PrepareFrameInternal(VideoCore::ImageId image_id, bool is_eop = true);
diff --git a/src/video_core/renderer_vulkan/vk_rasterizer.h b/src/video_core/renderer_vulkan/vk_rasterizer.h
index fe8aceba723..1936276a21a 100644
--- a/src/video_core/renderer_vulkan/vk_rasterizer.h
+++ b/src/video_core/renderer_vulkan/vk_rasterizer.h
@@ -54,6 +54,10 @@ class Rasterizer {
     u64 Flush();
     void Finish();
 
+    PipelineCache& GetPipelineCache() {
+        return pipeline_cache;
+    }
+
 private:
     RenderState PrepareRenderState(u32 mrt_mask);
     void BeginRendering(const GraphicsPipeline& pipeline, RenderState& state);