From cb56e16a4b6743e43e1c2ed783e4a58830b5caf6 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Wed, 6 Jul 2022 01:14:02 +0200
Subject: [PATCH] [dxvk] Use precompiled compute pipelines whenever possible

---
 src/dxvk/dxvk_compute.cpp     | 48 ++++++++++++++++++++++++-----------
 src/dxvk/dxvk_compute.h       |  6 ++++-
 src/dxvk/dxvk_pipemanager.cpp | 16 +++++++++++-
 src/dxvk/dxvk_pipemanager.h   |  3 +++
 4 files changed, 56 insertions(+), 17 deletions(-)

diff --git a/src/dxvk/dxvk_compute.cpp b/src/dxvk/dxvk_compute.cpp
index ee908d20..ffe8c983 100644
--- a/src/dxvk/dxvk_compute.cpp
+++ b/src/dxvk/dxvk_compute.cpp
@@ -14,11 +14,14 @@ namespace dxvk {
           DxvkDevice*                 device,
           DxvkPipelineManager*        pipeMgr,
           DxvkComputePipelineShaders  shaders,
-          DxvkBindingLayoutObjects*   layout)
+          DxvkBindingLayoutObjects*   layout,
+          DxvkShaderPipelineLibrary*  library)
   : m_device        (device),
     m_cache         (&pipeMgr->m_cache),
     m_stateCache    (&pipeMgr->m_stateCache),
     m_stats         (&pipeMgr->m_stats),
+    m_library       (library),
+    m_libraryHandle (VK_NULL_HANDLE),
     m_shaders       (std::move(shaders)),
     m_bindings      (layout) {
 
@@ -33,28 +36,43 @@ namespace dxvk {
   
   VkPipeline DxvkComputePipeline::getPipelineHandle(
     const DxvkComputePipelineStateInfo& state) {
-    DxvkComputePipelineInstance* instance = this->findInstance(state);
-
-    if (unlikely(!instance)) {
-      std::lock_guard<dxvk::mutex> lock(m_mutex);
-      instance = this->findInstance(state);
-
-      if (!instance) {
-        instance = this->createInstance(state);
-        this->writePipelineStateToCache(state);
+    if (m_library) {
+      // For compute pipelines that can be precompiled, we can use that
+      // pipeline variant unconditionally since there is no state for us
+      // to worry about other than specialization constants
+      if (unlikely(!m_libraryHandle)) {
+        m_libraryHandle = m_library->getPipelineHandle(m_cache->handle(),
+          DxvkShaderPipelineLibraryCompileArgs());
+        m_stats->numComputePipelines += 1;
       }
-    }
 
-    return instance->pipeline();
+      return m_libraryHandle;
+    } else {
+      DxvkComputePipelineInstance* instance = this->findInstance(state);
+
+      if (unlikely(!instance)) {
+        std::lock_guard<dxvk::mutex> lock(m_mutex);
+        instance = this->findInstance(state);
+
+        if (!instance) {
+          instance = this->createInstance(state);
+          this->writePipelineStateToCache(state);
+        }
+      }
+
+      return instance->pipeline();
+    }
   }
 
 
   void DxvkComputePipeline::compilePipeline(
     const DxvkComputePipelineStateInfo& state) {
-    std::lock_guard<dxvk::mutex> lock(m_mutex);
+    if (!m_library) {
+      std::lock_guard<dxvk::mutex> lock(m_mutex);
 
-    if (!this->findInstance(state))
-      this->createInstance(state);
+      if (!this->findInstance(state))
+        this->createInstance(state);
+    }
   }
   
   
diff --git a/src/dxvk/dxvk_compute.h b/src/dxvk/dxvk_compute.h
index 4da2c14a..d492f4d6 100644
--- a/src/dxvk/dxvk_compute.h
+++ b/src/dxvk/dxvk_compute.h
@@ -96,7 +96,8 @@ namespace dxvk {
             DxvkDevice*                 device,
             DxvkPipelineManager*        pipeMgr,
             DxvkComputePipelineShaders  shaders,
-            DxvkBindingLayoutObjects*   layout);
+            DxvkBindingLayoutObjects*   layout,
+            DxvkShaderPipelineLibrary*  library);
 
     ~DxvkComputePipeline();
     
@@ -146,6 +147,9 @@ namespace dxvk {
     DxvkStateCache*             m_stateCache;
     DxvkPipelineStats*          m_stats;
 
+    DxvkShaderPipelineLibrary*  m_library;
+    VkPipeline                  m_libraryHandle;
+
     DxvkComputePipelineShaders  m_shaders;
     DxvkBindingLayoutObjects*   m_bindings;
     
diff --git a/src/dxvk/dxvk_pipemanager.cpp b/src/dxvk/dxvk_pipemanager.cpp
index a2980ebc..578dda39 100644
--- a/src/dxvk/dxvk_pipemanager.cpp
+++ b/src/dxvk/dxvk_pipemanager.cpp
@@ -189,11 +189,12 @@ namespace dxvk {
       return &pair->second;
 
     auto layout = createPipelineLayout(shaders.cs->getBindings());
+    auto library = findPipelineLibrary(shaders.cs);
 
     auto iter = m_computePipelines.emplace(
       std::piecewise_construct,
       std::tuple(shaders),
-      std::tuple(m_device, this, shaders, layout));
+      std::tuple(m_device, this, shaders, layout, library));
     return &iter.first->second;
   }
   
@@ -341,5 +342,18 @@ namespace dxvk {
       std::tuple(m_device, shader.ptr(), layout));
     return &iter.first->second;
   }
+
+
+  DxvkShaderPipelineLibrary* DxvkPipelineManager::findPipelineLibrary(
+    const Rc<DxvkShader>&     shader) {
+    DxvkShaderPipelineLibraryKey key;
+    key.shader = shader;
+
+    auto pair = m_shaderLibraries.find(key);
+    if (pair == m_shaderLibraries.end())
+      return nullptr;
+
+    return &pair->second;
+  }
   
 }
diff --git a/src/dxvk/dxvk_pipemanager.h b/src/dxvk/dxvk_pipemanager.h
index b10389a7..a3739acd 100644
--- a/src/dxvk/dxvk_pipemanager.h
+++ b/src/dxvk/dxvk_pipemanager.h
@@ -271,6 +271,9 @@ namespace dxvk {
     DxvkShaderPipelineLibrary* createPipelineLibrary(
       const Rc<DxvkShader>&     shader);
 
+    DxvkShaderPipelineLibrary* findPipelineLibrary(
+      const Rc<DxvkShader>&     shader);
+
   };
   
 }
\ No newline at end of file