devtools: patch shader at runtime

viniciuslrangel · Dec 8, 2024 · 0c70812 · 0c70812
1 parent ee18fc8
commit 0c70812
Show file tree

Hide file tree

Showing 11 changed files with 170 additions and 60 deletions.
diff --git 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
@@ -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,26 +76,30 @@ struct FrameDump {
 
 struct ShaderDump {
     std::string name;
+    vk::ShaderModule module;
 
     std::vector<u32> spv;
     std::vector<u32> isa;
 
     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
@@ -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,31 +84,36 @@ 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;
     }
 
     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;
             }
             isa_editor.SetText(value.cache_patch_disasm);
             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,14 +188,15 @@ bool ShaderList::Selection::DrawShader(DebugStateType::ShaderDump& value) {
                                                     res);
                     } else {
                         isa_editor.SetText(value.cache_patch_disasm);
+                        ReloadShader(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
@@ -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
@@ -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
@@ -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
@@ -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"