From 4bc2d713fbbba2950af81105e7e07602f0bf67a4 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Mon, 8 Aug 2022 14:25:34 +0200
Subject: [PATCH] [dxvk] Destroy shader pipeline libraries after initial
 compile on 32-bit

This alone saves ~700 MiB of address space in the Resident Evil 6 main
menu on Nvidia.
---
 src/dxvk/dxvk_shader.cpp | 80 +++++++++++++++++++++++++++++-----------
 src/dxvk/dxvk_shader.h   |  6 +++
 2 files changed, 64 insertions(+), 22 deletions(-)

diff --git a/src/dxvk/dxvk_shader.cpp b/src/dxvk/dxvk_shader.cpp
index 71f8abc9c..c044b465f 100644
--- a/src/dxvk/dxvk_shader.cpp
+++ b/src/dxvk/dxvk_shader.cpp
@@ -915,10 +915,7 @@ namespace dxvk {
     const DxvkShaderPipelineLibraryCompileArgs& args) {
     std::lock_guard lock(m_mutex);
 
-    VkShaderStageFlagBits stage = VK_SHADER_STAGE_FRAGMENT_BIT;
-
-    if (m_shader)
-      stage = m_shader->info().stage;
+    VkShaderStageFlagBits stage = getShaderStage();
 
     VkPipeline& pipeline = (stage == VK_SHADER_STAGE_VERTEX_BIT && !args.depthClipEnable)
       ? m_pipelineNoDepthClip
@@ -927,28 +924,54 @@ namespace dxvk {
     if (pipeline)
       return pipeline;
 
-    bool usesDefaultArgs = (args == DxvkShaderPipelineLibraryCompileArgs());
+    pipeline = compileShaderPipelineLocked(args);
+    return pipeline;
+  }
 
+
+  void DxvkShaderPipelineLibrary::compilePipeline() {
+    std::lock_guard lock(m_mutex);
+
+    // Skip if a pipeline has already been compiled
+    if (m_pipeline)
+      return;
+
+    // Compile the pipeline with default args
+    VkPipeline pipeline = compileShaderPipelineLocked(
+      DxvkShaderPipelineLibraryCompileArgs());
+
+    // On 32-bit, destroy the pipeline immediately in order to
+    // save memory. We should hit the driver's disk cache once
+    // we need to recreate the pipeline.
+    if (m_device->mustTrackPipelineLifetime()) {
+      auto vk = m_device->vkd();
+      vk->vkDestroyPipeline(vk->device(), pipeline, nullptr);
+
+      pipeline = VK_NULL_HANDLE;
+    }
+
+    // Write back pipeline handle for future use
+    m_pipeline = pipeline;
+  }
+
+
+  VkPipeline DxvkShaderPipelineLibrary::compileShaderPipelineLocked(
+    const DxvkShaderPipelineLibraryCompileArgs& args) {
+    VkShaderStageFlagBits stage = getShaderStage();
+    VkPipeline pipeline = VK_NULL_HANDLE;
+
+    // Compile pipeline of the appropriate type
     switch (stage) {
       case VK_SHADER_STAGE_VERTEX_BIT:
         pipeline = compileVertexShaderPipeline(args);
-
-        if (usesDefaultArgs)
-          m_stats->numGraphicsLibraries += 1;
         break;
 
       case VK_SHADER_STAGE_FRAGMENT_BIT:
         pipeline = compileFragmentShaderPipeline();
-
-        if (usesDefaultArgs)
-          m_stats->numGraphicsLibraries += 1;
         break;
 
       case VK_SHADER_STAGE_COMPUTE_BIT:
         pipeline = compileComputeShaderPipeline();
-
-        if (usesDefaultArgs)
-          m_stats->numComputePipelines += 1;
         break;
 
       default:
@@ -956,18 +979,21 @@ namespace dxvk {
         return VK_NULL_HANDLE;
     }
 
+    // Increment stat counter the first time this
+    // shader pipeline gets compiled successfully
+    if (!m_compiledOnce && pipeline) {
+      if (stage == VK_SHADER_STAGE_COMPUTE_BIT)
+        m_stats->numComputePipelines += 1;
+      else
+        m_stats->numGraphicsLibraries += 1;
+
+      m_compiledOnce = true;
+    }
+
     return pipeline;
   }
 
 
-  void DxvkShaderPipelineLibrary::compilePipeline() {
-    // Just compile the pipeline with default args. Implicitly skips
-    // this step if another thread has compiled the pipeline in the
-    // meantime, in order to avoid duplicate work.
-    getPipelineHandle(DxvkShaderPipelineLibraryCompileArgs());
-  }
-
-
   VkPipeline DxvkShaderPipelineLibrary::compileVertexShaderPipeline(
     const DxvkShaderPipelineLibraryCompileArgs& args) {
     auto vk = m_device->vkd();
@@ -1175,4 +1201,14 @@ namespace dxvk {
       vk->device(), &info, &m_identifier);
   }
 
+
+  VkShaderStageFlagBits DxvkShaderPipelineLibrary::getShaderStage() const {
+    VkShaderStageFlagBits stage = VK_SHADER_STAGE_FRAGMENT_BIT;
+
+    if (m_shader != nullptr)
+      stage = m_shader->info().stage;
+
+    return stage;
+  }
+
 }
\ No newline at end of file
diff --git a/src/dxvk/dxvk_shader.h b/src/dxvk/dxvk_shader.h
index 6363e3a41..6086b1c87 100644
--- a/src/dxvk/dxvk_shader.h
+++ b/src/dxvk/dxvk_shader.h
@@ -425,10 +425,14 @@ namespace dxvk {
     dxvk::mutex     m_mutex;
     VkPipeline      m_pipeline             = VK_NULL_HANDLE;
     VkPipeline      m_pipelineNoDepthClip  = VK_NULL_HANDLE;
+    bool            m_compiledOnce         = false;
 
     dxvk::mutex                 m_identifierMutex;
     VkShaderModuleIdentifierEXT m_identifier = { VK_STRUCTURE_TYPE_SHADER_MODULE_IDENTIFIER_EXT };
 
+    VkPipeline compileShaderPipelineLocked(
+      const DxvkShaderPipelineLibraryCompileArgs& args);
+
     VkPipeline compileVertexShaderPipeline(
       const DxvkShaderPipelineLibraryCompileArgs& args);
 
@@ -444,6 +448,8 @@ namespace dxvk {
     void generateModuleIdentifierLocked(
       const SpirvCodeBuffer& spirvCode);
 
+    VkShaderStageFlagBits getShaderStage() const;
+
   };
   
 }
\ No newline at end of file