2 files changed, 520 insertions, 0 deletions
diff --git a/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.cpp b/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.cpp
new file mode 100644
index 0000000000..27112846c6
--- /dev/null
+++ b/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.cpp
@@ -0,0 +1,376 @@
+// Copyright 2022 Dolphin Emulator Project
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+#include "VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.h"
+#include "VideoCommon/AbstractGfx.h"
+#include "VideoCommon/VideoConfig.h"
+
+CustomShaderCache::CustomShaderCache()
+{
+  m_api_type = g_ActiveConfig.backend_info.api_type;
+  m_host_config.bits = ShaderHostConfig::GetCurrent().bits;
+
+  m_async_shader_compiler = g_gfx->CreateAsyncShaderCompiler();
+  m_async_shader_compiler->StartWorkerThreads(1);  // TODO
+
+  m_async_uber_shader_compiler = g_gfx->CreateAsyncShaderCompiler();
+  m_async_uber_shader_compiler->StartWorkerThreads(1);  // TODO
+
+  m_frame_end_handler =
+      AfterFrameEvent::Register([this] { RetrieveAsyncShaders(); }, "RetreiveAsyncShaders");
+}
+
+CustomShaderCache::~CustomShaderCache()
+{
+  if (m_async_shader_compiler)
+    m_async_shader_compiler->StopWorkerThreads();
+
+  if (m_async_uber_shader_compiler)
+    m_async_uber_shader_compiler->StopWorkerThreads();
+}
+
+void CustomShaderCache::RetrieveAsyncShaders()
+{
+  m_async_shader_compiler->RetrieveWorkItems();
+  m_async_uber_shader_compiler->RetrieveWorkItems();
+}
+
+void CustomShaderCache::Reload()
+{
+  while (m_async_shader_compiler->HasPendingWork() || m_async_shader_compiler->HasCompletedWork())
+  {
+    m_async_shader_compiler->RetrieveWorkItems();
+  }
+
+  while (m_async_uber_shader_compiler->HasPendingWork() ||
+         m_async_uber_shader_compiler->HasCompletedWork())
+  {
+    m_async_uber_shader_compiler->RetrieveWorkItems();
+  }
+
+  m_ps_cache = {};
+  m_uber_ps_cache = {};
+  m_pipeline_cache = {};
+  m_uber_pipeline_cache = {};
+}
+
+std::optional<const AbstractPipeline*>
+CustomShaderCache::GetPipelineAsync(const VideoCommon::GXPipelineUid& uid,
+                                    const CustomShaderInstance& custom_shaders,
+                                    const AbstractPipelineConfig& pipeline_config)
+{
+  if (auto holder = m_pipeline_cache.GetHolder(uid, custom_shaders))
+  {
+    if (holder->pending)
+      return std::nullopt;
+    return holder->value.get();
+  }
+  AsyncCreatePipeline(uid, custom_shaders, pipeline_config);
+  return std::nullopt;
+}
+
+std::optional<const AbstractPipeline*>
+CustomShaderCache::GetPipelineAsync(const VideoCommon::GXUberPipelineUid& uid,
+                                    const CustomShaderInstance& custom_shaders,
+                                    const AbstractPipelineConfig& pipeline_config)
+{
+  if (auto holder = m_uber_pipeline_cache.GetHolder(uid, custom_shaders))
+  {
+    if (holder->pending)
+      return std::nullopt;
+    return holder->value.get();
+  }
+  AsyncCreatePipeline(uid, custom_shaders, pipeline_config);
+  return std::nullopt;
+}
+
+void CustomShaderCache::AsyncCreatePipeline(const VideoCommon::GXPipelineUid& uid,
+
+                                            const CustomShaderInstance& custom_shaders,
+                                            const AbstractPipelineConfig& pipeline_config)
+{
+  class PipelineWorkItem final : public VideoCommon::AsyncShaderCompiler::WorkItem
+  {
+  public:
+    PipelineWorkItem(CustomShaderCache* shader_cache, const VideoCommon::GXPipelineUid& uid,
+                     const CustomShaderInstance& custom_shaders, PipelineIterator iterator,
+                     const AbstractPipelineConfig& pipeline_config)
+        : m_shader_cache(shader_cache), m_uid(uid), m_iterator(iterator),
+          m_custom_shaders(custom_shaders), m_config(pipeline_config)
+    {
+      SetStagesReady();
+    }
+
+    void SetStagesReady()
+    {
+      m_stages_ready = true;
+
+      PixelShaderUid ps_uid = m_uid.ps_uid;
+      ClearUnusedPixelShaderUidBits(m_shader_cache->m_api_type, m_shader_cache->m_host_config,
+                                    &ps_uid);
+
+      if (auto holder = m_shader_cache->m_ps_cache.GetHolder(ps_uid, m_custom_shaders))
+      {
+        // If the pixel shader is no longer pending compilation
+        // and the shader compilation succeeded, set
+        // the pipeline to use the new pixel shader.
+        // Otherwise, use the existing shader.
+        if (!holder->pending && holder->value.get())
+        {
+          m_config.pixel_shader = holder->value.get();
+        }
+        m_stages_ready &= !holder->pending;
+      }
+      else
+      {
+        m_stages_ready &= false;
+        m_shader_cache->QueuePixelShaderCompile(ps_uid, m_custom_shaders);
+      }
+    }
+
+    bool Compile() override
+    {
+      if (m_stages_ready)
+      {
+        m_pipeline = g_gfx->CreatePipeline(m_config);
+      }
+      return true;
+    }
+
+    void Retrieve() override
+    {
+      if (m_stages_ready)
+      {
+        m_shader_cache->NotifyPipelineFinished(m_iterator, std::move(m_pipeline));
+      }
+      else
+      {
+        // Re-queue for next frame.
+        auto wi = m_shader_cache->m_async_shader_compiler->CreateWorkItem<PipelineWorkItem>(
+            m_shader_cache, m_uid, m_custom_shaders, m_iterator, m_config);
+        m_shader_cache->m_async_shader_compiler->QueueWorkItem(std::move(wi), 0);
+      }
+    }
+
+  private:
+    CustomShaderCache* m_shader_cache;
+    std::unique_ptr<AbstractPipeline> m_pipeline;
+    VideoCommon::GXPipelineUid m_uid;
+    PipelineIterator m_iterator;
+    AbstractPipelineConfig m_config;
+    CustomShaderInstance m_custom_shaders;
+    bool m_stages_ready;
+  };
+
+  auto list_iter = m_pipeline_cache.InsertElement(uid, custom_shaders);
+  auto work_item = m_async_shader_compiler->CreateWorkItem<PipelineWorkItem>(
+      this, uid, custom_shaders, list_iter, pipeline_config);
+  m_async_shader_compiler->QueueWorkItem(std::move(work_item), 0);
+}
+
+void CustomShaderCache::AsyncCreatePipeline(const VideoCommon::GXUberPipelineUid& uid,
+
+                                            const CustomShaderInstance& custom_shaders,
+                                            const AbstractPipelineConfig& pipeline_config)
+{
+  class PipelineWorkItem final : public VideoCommon::AsyncShaderCompiler::WorkItem
+  {
+  public:
+    PipelineWorkItem(CustomShaderCache* shader_cache, const VideoCommon::GXUberPipelineUid& uid,
+                     const CustomShaderInstance& custom_shaders, UberPipelineIterator iterator,
+                     const AbstractPipelineConfig& pipeline_config)
+        : m_shader_cache(shader_cache), m_uid(uid), m_iterator(iterator),
+          m_custom_shaders(custom_shaders), m_config(pipeline_config)
+    {
+      SetStagesReady();
+    }
+
+    void SetStagesReady()
+    {
+      m_stages_ready = true;
+
+      UberShader::PixelShaderUid ps_uid = m_uid.ps_uid;
+      ClearUnusedPixelShaderUidBits(m_shader_cache->m_api_type, m_shader_cache->m_host_config,
+                                    &ps_uid);
+
+      if (auto holder = m_shader_cache->m_uber_ps_cache.GetHolder(ps_uid, m_custom_shaders))
+      {
+        if (!holder->pending && holder->value.get())
+        {
+          m_config.pixel_shader = holder->value.get();
+        }
+        m_stages_ready &= !holder->pending;
+      }
+      else
+      {
+        m_stages_ready &= false;
+        m_shader_cache->QueuePixelShaderCompile(ps_uid, m_custom_shaders);
+      }
+    }
+
+    bool Compile() override
+    {
+      if (m_stages_ready)
+      {
+        if (m_config.pixel_shader == nullptr || m_config.vertex_shader == nullptr)
+          return false;
+
+        m_pipeline = g_gfx->CreatePipeline(m_config);
+      }
+      return true;
+    }
+
+    void Retrieve() override
+    {
+      if (m_stages_ready)
+      {
+        m_shader_cache->NotifyPipelineFinished(m_iterator, std::move(m_pipeline));
+      }
+      else
+      {
+        // Re-queue for next frame.
+        auto wi = m_shader_cache->m_async_uber_shader_compiler->CreateWorkItem<PipelineWorkItem>(
+            m_shader_cache, m_uid, m_custom_shaders, m_iterator, m_config);
+        m_shader_cache->m_async_uber_shader_compiler->QueueWorkItem(std::move(wi), 0);
+      }
+    }
+
+  private:
+    CustomShaderCache* m_shader_cache;
+    std::unique_ptr<AbstractPipeline> m_pipeline;
+    VideoCommon::GXUberPipelineUid m_uid;
+    UberPipelineIterator m_iterator;
+    AbstractPipelineConfig m_config;
+    CustomShaderInstance m_custom_shaders;
+    bool m_stages_ready;
+  };
+
+  auto list_iter = m_uber_pipeline_cache.InsertElement(uid, custom_shaders);
+  auto work_item = m_async_uber_shader_compiler->CreateWorkItem<PipelineWorkItem>(
+      this, uid, custom_shaders, list_iter, pipeline_config);
+  m_async_uber_shader_compiler->QueueWorkItem(std::move(work_item), 0);
+}
+
+void CustomShaderCache::NotifyPipelineFinished(PipelineIterator iterator,
+                                               std::unique_ptr<AbstractPipeline> pipeline)
+{
+  iterator->second.pending = false;
+  iterator->second.value = std::move(pipeline);
+}
+
+void CustomShaderCache::NotifyPipelineFinished(UberPipelineIterator iterator,
+                                               std::unique_ptr<AbstractPipeline> pipeline)
+{
+  iterator->second.pending = false;
+  iterator->second.value = std::move(pipeline);
+}
+
+void CustomShaderCache::QueuePixelShaderCompile(const PixelShaderUid& uid,
+
+                                                const CustomShaderInstance& custom_shaders)
+{
+  class PixelShaderWorkItem final : public VideoCommon::AsyncShaderCompiler::WorkItem
+  {
+  public:
+    PixelShaderWorkItem(CustomShaderCache* shader_cache, const PixelShaderUid& uid,
+                        const CustomShaderInstance& custom_shaders, PixelShaderIterator iter)
+        : m_shader_cache(shader_cache), m_uid(uid), m_custom_shaders(custom_shaders), m_iter(iter)
+    {
+    }
+
+    bool Compile() override
+    {
+      m_shader = m_shader_cache->CompilePixelShader(m_uid, m_custom_shaders);
+      return true;
+    }
+
+    void Retrieve() override
+    {
+      m_shader_cache->NotifyPixelShaderFinished(m_iter, std::move(m_shader));
+    }
+
+  private:
+    CustomShaderCache* m_shader_cache;
+    std::unique_ptr<AbstractShader> m_shader;
+    PixelShaderUid m_uid;
+    CustomShaderInstance m_custom_shaders;
+    PixelShaderIterator m_iter;
+  };
+
+  auto list_iter = m_ps_cache.InsertElement(uid, custom_shaders);
+  auto work_item = m_async_shader_compiler->CreateWorkItem<PixelShaderWorkItem>(
+      this, uid, custom_shaders, list_iter);
+  m_async_shader_compiler->QueueWorkItem(std::move(work_item), 0);
+}
+
+void CustomShaderCache::QueuePixelShaderCompile(const UberShader::PixelShaderUid& uid,
+
+                                                const CustomShaderInstance& custom_shaders)
+{
+  class PixelShaderWorkItem final : public VideoCommon::AsyncShaderCompiler::WorkItem
+  {
+  public:
+    PixelShaderWorkItem(CustomShaderCache* shader_cache, const UberShader::PixelShaderUid& uid,
+                        const CustomShaderInstance& custom_shaders, UberPixelShaderIterator iter)
+        : m_shader_cache(shader_cache), m_uid(uid), m_custom_shaders(custom_shaders), m_iter(iter)
+    {
+    }
+
+    bool Compile() override
+    {
+      m_shader = m_shader_cache->CompilePixelShader(m_uid, m_custom_shaders);
+      return true;
+    }
+
+    void Retrieve() override
+    {
+      m_shader_cache->NotifyPixelShaderFinished(m_iter, std::move(m_shader));
+    }
+
+  private:
+    CustomShaderCache* m_shader_cache;
+    std::unique_ptr<AbstractShader> m_shader;
+    UberShader::PixelShaderUid m_uid;
+    CustomShaderInstance m_custom_shaders;
+    UberPixelShaderIterator m_iter;
+  };
+
+  auto list_iter = m_uber_ps_cache.InsertElement(uid, custom_shaders);
+  auto work_item = m_async_uber_shader_compiler->CreateWorkItem<PixelShaderWorkItem>(
+      this, uid, custom_shaders, list_iter);
+  m_async_uber_shader_compiler->QueueWorkItem(std::move(work_item), 0);
+}
+
+std::unique_ptr<AbstractShader>
+CustomShaderCache::CompilePixelShader(const PixelShaderUid& uid,
+                                      const CustomShaderInstance& custom_shaders) const
+{
+  const ShaderCode source_code = GeneratePixelShaderCode(
+      m_api_type, m_host_config, uid.GetUidData(), custom_shaders.pixel_contents);
+  return g_gfx->CreateShaderFromSource(ShaderStage::Pixel, source_code.GetBuffer(),
+                                       "Custom Pixel Shader");
+}
+
+std::unique_ptr<AbstractShader>
+CustomShaderCache::CompilePixelShader(const UberShader::PixelShaderUid& uid,
+                                      const CustomShaderInstance& custom_shaders) const
+{
+  const ShaderCode source_code =
+      GenPixelShader(m_api_type, m_host_config, uid.GetUidData(), custom_shaders.pixel_contents);
+  return g_gfx->CreateShaderFromSource(ShaderStage::Pixel, source_code.GetBuffer(),
+                                       "Custom Uber Pixel Shader");
+}
+
+void CustomShaderCache::NotifyPixelShaderFinished(PixelShaderIterator iterator,
+                                                  std::unique_ptr<AbstractShader> shader)
+{
+  iterator->second.pending = false;
+  iterator->second.value = std::move(shader);
+}
+
+void CustomShaderCache::NotifyPixelShaderFinished(UberPixelShaderIterator iterator,
+                                                  std::unique_ptr<AbstractShader> shader)
+{
+  iterator->second.pending = false;
+  iterator->second.value = std::move(shader);
+}
diff --git a/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.h b/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.h
new file mode 100644
index 0000000000..ff2aba2823
--- /dev/null
+++ b/Source/Core/VideoCommon/GraphicsModSystem/Runtime/CustomShaderCache.h
@@ -0,0 +1,144 @@
+// Copyright 2022 Dolphin Emulator Project
+// SPDX-License-Identifier: GPL-2.0-or-later
+
+#pragma once
+
+#include <array>
+#include <list>
+#include <map>
+#include <memory>
+#include <optional>
+#include <string>
+#include <string_view>
+
+#include "VideoCommon/AbstractPipeline.h"
+#include "VideoCommon/AbstractShader.h"
+#include "VideoCommon/AsyncShaderCompiler.h"
+#include "VideoCommon/GXPipelineTypes.h"
+#include "VideoCommon/PixelShaderGen.h"
+#include "VideoCommon/ShaderGenCommon.h"
+#include "VideoCommon/UberShaderPixel.h"
+#include "VideoCommon/VideoEvents.h"
+
+struct CustomShaderInstance
+{
+  CustomPixelShaderContents pixel_contents;
+
+  bool operator==(const CustomShaderInstance& other) const = default;
+};
+
+class CustomShaderCache
+{
+public:
+  CustomShaderCache();
+  ~CustomShaderCache();
+  CustomShaderCache(const CustomShaderCache&) = delete;
+  CustomShaderCache(CustomShaderCache&&) = delete;
+  CustomShaderCache& operator=(const CustomShaderCache&) = delete;
+  CustomShaderCache& operator=(CustomShaderCache&&) = delete;
+
+  // Changes the shader host config. Shaders should be reloaded afterwards.
+  void SetHostConfig(const ShaderHostConfig& host_config) { m_host_config.bits = host_config.bits; }
+
+  // Retrieves all pending shaders/pipelines from the async compiler.
+  void RetrieveAsyncShaders();
+
+  // Reloads/recreates all shaders and pipelines.
+  void Reload();
+
+  // The optional will be empty if this pipeline is now background compiling.
+  std::optional<const AbstractPipeline*>
+  GetPipelineAsync(const VideoCommon::GXPipelineUid& uid,
+                   const CustomShaderInstance& custom_shaders,
+                   const AbstractPipelineConfig& pipeline_config);
+  std::optional<const AbstractPipeline*>
+  GetPipelineAsync(const VideoCommon::GXUberPipelineUid& uid,
+                   const CustomShaderInstance& custom_shaders,
+                   const AbstractPipelineConfig& pipeline_config);
+
+private:
+  // Configuration bits.
+  APIType m_api_type = APIType::Nothing;
+  ShaderHostConfig m_host_config = {};
+  std::unique_ptr<VideoCommon::AsyncShaderCompiler> m_async_shader_compiler;
+  std::unique_ptr<VideoCommon::AsyncShaderCompiler> m_async_uber_shader_compiler;
+
+  void AsyncCreatePipeline(const VideoCommon::GXPipelineUid& uid,
+                           const CustomShaderInstance& custom_shaders,
+                           const AbstractPipelineConfig& pipeline_config);
+  void AsyncCreatePipeline(const VideoCommon::GXUberPipelineUid& uid,
+                           const CustomShaderInstance& custom_shaders,
+                           const AbstractPipelineConfig& pipeline_config);
+
+  // Shader/Pipeline cache helper
+  template <typename Uid, typename ValueType>
+  struct Cache
+  {
+    struct CacheHolder
+    {
+      std::unique_ptr<ValueType> value = nullptr;
+      bool pending = true;
+    };
+    using CacheElement = std::pair<CustomShaderInstance, CacheHolder>;
+    using CacheList = std::list<CacheElement>;
+    std::map<Uid, CacheList> uid_to_cachelist;
+
+    const CacheHolder* GetHolder(const Uid& uid, const CustomShaderInstance& custom_shaders) const
+    {
+      if (auto uuid_it = uid_to_cachelist.find(uid); uuid_it != uid_to_cachelist.end())
+      {
+        for (const auto& [custom_shader_val, holder] : uuid_it->second)
+        {
+          if (custom_shaders == custom_shader_val)
+          {
+            return &holder;
+          }
+        }
+      }
+
+      return nullptr;
+    }
+
+    typename CacheList::iterator InsertElement(const Uid& uid,
+                                               const CustomShaderInstance& custom_shaders)
+    {
+      CacheList& cachelist = uid_to_cachelist[uid];
+      CacheElement e{custom_shaders, CacheHolder{}};
+      return cachelist.emplace(cachelist.begin(), std::move(e));
+    }
+  };
+
+  Cache<PixelShaderUid, AbstractShader> m_ps_cache;
+  Cache<UberShader::PixelShaderUid, AbstractShader> m_uber_ps_cache;
+  Cache<VideoCommon::GXPipelineUid, AbstractPipeline> m_pipeline_cache;
+  Cache<VideoCommon::GXUberPipelineUid, AbstractPipeline> m_uber_pipeline_cache;
+
+  using PipelineIterator = Cache<VideoCommon::GXPipelineUid, AbstractPipeline>::CacheList::iterator;
+  using UberPipelineIterator =
+      Cache<VideoCommon::GXUberPipelineUid, AbstractPipeline>::CacheList::iterator;
+  using PixelShaderIterator = Cache<PixelShaderUid, AbstractShader>::CacheList::iterator;
+  using UberPixelShaderIterator =
+      Cache<UberShader::PixelShaderUid, AbstractShader>::CacheList::iterator;
+
+  void NotifyPipelineFinished(PipelineIterator iterator,
+                              std::unique_ptr<AbstractPipeline> pipeline);
+  void NotifyPipelineFinished(UberPipelineIterator iterator,
+                              std::unique_ptr<AbstractPipeline> pipeline);
+
+  std::unique_ptr<AbstractShader>
+  CompilePixelShader(const PixelShaderUid& uid, const CustomShaderInstance& custom_shaders) const;
+  void NotifyPixelShaderFinished(PixelShaderIterator iterator,
+                                 std::unique_ptr<AbstractShader> shader);
+  std::unique_ptr<AbstractShader>
+  CompilePixelShader(const UberShader::PixelShaderUid& uid,
+                     const CustomShaderInstance& custom_shaders) const;
+  void NotifyPixelShaderFinished(UberPixelShaderIterator iterator,
+                                 std::unique_ptr<AbstractShader> shader);
+
+  void QueuePixelShaderCompile(const PixelShaderUid& uid,
+                               const CustomShaderInstance& custom_shaders);
+  void QueuePixelShaderCompile(const UberShader::PixelShaderUid& uid,
+                               const CustomShaderInstance& custom_shaders);
+
+  Common::EventHook m_frame_end_handler;
+};