From 764de6ff82d187362d9b4964111662f76d4eadb5 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Mon, 8 Aug 2022 16:12:26 +0200
Subject: [PATCH] [dxvk] Add use counter to pipeline libraries

And destroy Vulkan pipeline objects once the counter reaches zero.
---
 src/dxvk/dxvk_compute.cpp  |  5 ++++-
 src/dxvk/dxvk_graphics.cpp | 10 +++++++---
 src/dxvk/dxvk_shader.cpp   | 31 ++++++++++++++++++++++++++-----
 src/dxvk/dxvk_shader.h     | 17 +++++++++++++++--
 4 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/src/dxvk/dxvk_compute.cpp b/src/dxvk/dxvk_compute.cpp
index 0006ce145..fdc7394d2 100644
--- a/src/dxvk/dxvk_compute.cpp
+++ b/src/dxvk/dxvk_compute.cpp
@@ -30,6 +30,9 @@ namespace dxvk {
   
   
   DxvkComputePipeline::~DxvkComputePipeline() {
+    if (m_libraryHandle)
+      m_library->releasePipelineHandle();
+
     for (const auto& instance : m_pipelines)
       this->destroyPipeline(instance.handle);
   }
@@ -45,7 +48,7 @@ namespace dxvk {
       // Retrieve actual pipeline handle on first use. This
       // may wait for an ongoing compile job to finish, or
       // compile the pipeline immediately on the calling thread.
-      m_libraryHandle = m_library->getPipelineHandle(
+      m_libraryHandle = m_library->acquirePipelineHandle(
         DxvkShaderPipelineLibraryCompileArgs());
 
       return m_libraryHandle;
diff --git a/src/dxvk/dxvk_graphics.cpp b/src/dxvk/dxvk_graphics.cpp
index 1070859d8..dbd8c7d02 100644
--- a/src/dxvk/dxvk_graphics.cpp
+++ b/src/dxvk/dxvk_graphics.cpp
@@ -880,9 +880,13 @@ namespace dxvk {
   
   
   DxvkGraphicsPipeline::~DxvkGraphicsPipeline() {
-    for (const auto& instance : m_fastPipelines)
+    for (const auto& instance : m_fastPipelines) {
       this->destroyPipeline(instance.second);
 
+      m_vsLibrary->releasePipelineHandle();
+      m_fsLibrary->releasePipelineHandle();
+    }
+
     for (const auto& instance : m_basePipelines)
       this->destroyPipeline(instance.second);
   }
@@ -1102,8 +1106,8 @@ namespace dxvk {
 
     std::array<VkPipeline, 4> libraries = {{
       key.viLibrary->getHandle(),
-      m_vsLibrary->getPipelineHandle(key.args),
-      m_fsLibrary->getPipelineHandle(key.args),
+      m_vsLibrary->acquirePipelineHandle(key.args),
+      m_fsLibrary->acquirePipelineHandle(key.args),
       key.foLibrary->getHandle(),
     }};
 
diff --git a/src/dxvk/dxvk_shader.cpp b/src/dxvk/dxvk_shader.cpp
index 4d7c468db..d5eba8499 100644
--- a/src/dxvk/dxvk_shader.cpp
+++ b/src/dxvk/dxvk_shader.cpp
@@ -890,10 +890,7 @@ namespace dxvk {
 
 
   DxvkShaderPipelineLibrary::~DxvkShaderPipelineLibrary() {
-    auto vk = m_device->vkd();
-
-    vk->vkDestroyPipeline(vk->device(), m_pipeline, nullptr);
-    vk->vkDestroyPipeline(vk->device(), m_pipelineNoDepthClip, nullptr);
+    this->destroyShaderPipelinesLocked();
   }
 
 
@@ -911,10 +908,13 @@ namespace dxvk {
   }
 
 
-  VkPipeline DxvkShaderPipelineLibrary::getPipelineHandle(
+  VkPipeline DxvkShaderPipelineLibrary::acquirePipelineHandle(
     const DxvkShaderPipelineLibraryCompileArgs& args) {
     std::lock_guard lock(m_mutex);
 
+    if (m_device->mustTrackPipelineLifetime())
+      m_useCount += 1;
+
     VkShaderStageFlagBits stage = getShaderStage();
 
     VkPipeline& pipeline = (stage == VK_SHADER_STAGE_VERTEX_BIT && !args.depthClipEnable)
@@ -929,6 +929,16 @@ namespace dxvk {
   }
 
 
+  void DxvkShaderPipelineLibrary::releasePipelineHandle() {
+    if (m_device->mustTrackPipelineLifetime()) {
+      std::lock_guard lock(m_mutex);
+
+      if (!(--m_useCount))
+        this->destroyShaderPipelinesLocked();
+    }
+  }
+
+
   void DxvkShaderPipelineLibrary::compilePipeline() {
     std::lock_guard lock(m_mutex);
 
@@ -955,6 +965,17 @@ namespace dxvk {
   }
 
 
+  void DxvkShaderPipelineLibrary::destroyShaderPipelinesLocked() {
+    auto vk = m_device->vkd();
+
+    vk->vkDestroyPipeline(vk->device(), m_pipeline, nullptr);
+    vk->vkDestroyPipeline(vk->device(), m_pipelineNoDepthClip, nullptr);
+
+    m_pipeline = VK_NULL_HANDLE;
+    m_pipelineNoDepthClip = VK_NULL_HANDLE;
+  }
+
+
   VkPipeline DxvkShaderPipelineLibrary::compileShaderPipelineLocked(
     const DxvkShaderPipelineLibraryCompileArgs& args) {
     VkShaderStageFlagBits stage = getShaderStage();
diff --git a/src/dxvk/dxvk_shader.h b/src/dxvk/dxvk_shader.h
index a794e0361..3c273a28e 100644
--- a/src/dxvk/dxvk_shader.h
+++ b/src/dxvk/dxvk_shader.h
@@ -396,16 +396,26 @@ namespace dxvk {
     VkShaderModuleIdentifierEXT getModuleIdentifier();
 
     /**
-     * \brief Queries pipeline handle for the given set of arguments
+     * \brief Acquires pipeline handle for the given set of arguments
      *
      * Either returns an already compiled pipeline library object, or
      * performs the compilation step if that has not happened yet.
+     * Increments the use count by one.
      * \param [in] args Compile arguments
      * \returns Vulkan pipeline handle
      */
-    VkPipeline getPipelineHandle(
+    VkPipeline acquirePipelineHandle(
       const DxvkShaderPipelineLibraryCompileArgs& args);
 
+    /**
+     * \brief Releases pipeline
+     *
+     * Decrements the use count by 1. If the use count reaches 0,
+     * any previously compiled pipeline library object may be
+     * destroyed in order to save memory.
+     */
+    void releasePipelineHandle();
+
     /**
      * \brief Compiles the pipeline with default arguments
      *
@@ -425,11 +435,14 @@ namespace dxvk {
     dxvk::mutex     m_mutex;
     VkPipeline      m_pipeline             = VK_NULL_HANDLE;
     VkPipeline      m_pipelineNoDepthClip  = VK_NULL_HANDLE;
+    uint32_t        m_useCount             = 0u;
     bool            m_compiledOnce         = false;
 
     dxvk::mutex                 m_identifierMutex;
     VkShaderModuleIdentifierEXT m_identifier = { VK_STRUCTURE_TYPE_SHADER_MODULE_IDENTIFIER_EXT };
 
+    void destroyShaderPipelinesLocked();
+
     VkPipeline compileShaderPipelineLocked(
       const DxvkShaderPipelineLibraryCompileArgs& args);