From c931b4ba87283f9c83e25c1ca5f8a87ef15d7574 Mon Sep 17 00:00:00 2001
From: Philip Rebohle <philip.rebohle@tu-dortmund.de>
Date: Fri, 16 Mar 2018 01:24:03 +0100
Subject: [PATCH] [dxvk] Implement proper compute pipeline lookup

Fixes correctness issues and potential GPU lockups in case
a compute shader resource is not bound at dispatch time.
---
 src/dxvk/dxvk_compute.cpp | 61 +++++++++++++++++++++++++++++++--------
 src/dxvk/dxvk_compute.h   | 18 ++++++++++--
 2 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/src/dxvk/dxvk_compute.cpp b/src/dxvk/dxvk_compute.cpp
index e97c316a2..511a5aae6 100644
--- a/src/dxvk/dxvk_compute.cpp
+++ b/src/dxvk/dxvk_compute.cpp
@@ -29,25 +29,35 @@ namespace dxvk {
       slotMapping.bindingInfos());
     
     m_cs = cs->createShaderModule(m_vkd, slotMapping);
-    
-    this->compilePipeline();
   }
   
   
   DxvkComputePipeline::~DxvkComputePipeline() {
-    if (m_pipeline != VK_NULL_HANDLE)
-      m_vkd->vkDestroyPipeline(m_vkd->device(), m_pipeline, nullptr);
+    this->destroyPipelines();
   }
   
   
   VkPipeline DxvkComputePipeline::getPipelineHandle(
-    const DxvkComputePipelineStateInfo& state) const {
-    // TODO take pipeine state into account
-    return m_pipeline;
+    const DxvkComputePipelineStateInfo& state) {
+    std::lock_guard<std::mutex> lock(m_mutex);
+    
+    for (const PipelineStruct& pair : m_pipelines) {
+      if (pair.stateVector == state)
+        return pair.pipeline;
+    }
+    
+    VkPipeline pipeline = this->compilePipeline(state, m_basePipeline);
+    m_pipelines.push_back({ state, pipeline });
+    
+    if (m_basePipeline == VK_NULL_HANDLE)
+      m_basePipeline = pipeline;
+    return pipeline;
   }
   
   
-  void DxvkComputePipeline::compilePipeline() {
+  VkPipeline DxvkComputePipeline::compilePipeline(
+    const DxvkComputePipelineStateInfo& state,
+          VkPipeline                    baseHandle) const {
     std::vector<VkDescriptorSetLayoutBinding> bindings;
 
     if (Logger::logLevel() <= LogLevel::Debug) {
@@ -55,18 +65,45 @@ namespace dxvk {
       Logger::debug(str::format("  cs  : ", m_cs ->debugName()));
     }
     
+    std::array<VkBool32,                 MaxNumActiveBindings> specData;
+    std::array<VkSpecializationMapEntry, MaxNumActiveBindings> specMap;
+    
+    for (uint32_t i = 0; i < MaxNumActiveBindings; i++) {
+      specData[i] = state.bsBindingState.isBound(i) ? VK_TRUE : VK_FALSE;
+      specMap [i] = { i, static_cast<uint32_t>(sizeof(VkBool32)) * i, sizeof(VkBool32) };
+    }
+    
+    VkSpecializationInfo specInfo;
+    specInfo.mapEntryCount    = specMap.size();
+    specInfo.pMapEntries      = specMap.data();
+    specInfo.dataSize         = specData.size() * sizeof(VkBool32);
+    specInfo.pData            = specData.data();
+    
     VkComputePipelineCreateInfo info;
     info.sType                = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
     info.pNext                = nullptr;
-    info.flags                = 0;
-    info.stage                = m_cs->stageInfo(nullptr);
+    info.flags                = baseHandle == VK_NULL_HANDLE
+      ? VK_PIPELINE_CREATE_ALLOW_DERIVATIVES_BIT
+      : VK_PIPELINE_CREATE_DERIVATIVE_BIT;
+    info.stage                = m_cs->stageInfo(&specInfo);
     info.layout               = m_layout->pipelineLayout();
-    info.basePipelineHandle   = VK_NULL_HANDLE;
+    info.basePipelineHandle   = baseHandle;
     info.basePipelineIndex    = -1;
     
+    VkPipeline pipeline = VK_NULL_HANDLE;
     if (m_vkd->vkCreateComputePipelines(m_vkd->device(),
-          m_cache->handle(), 1, &info, nullptr, &m_pipeline) != VK_SUCCESS)
+          m_cache->handle(), 1, &info, nullptr, &pipeline) != VK_SUCCESS) {
       Logger::err("DxvkComputePipeline: Failed to compile pipeline");
+      return VK_NULL_HANDLE;
+    }
+    
+    return pipeline;
+  }
+  
+  
+  void DxvkComputePipeline::destroyPipelines() {
+    for (const PipelineStruct& pair : m_pipelines)
+      m_vkd->vkDestroyPipeline(m_vkd->device(), pair.pipeline, nullptr);
   }
   
 }
\ No newline at end of file
diff --git a/src/dxvk/dxvk_compute.h b/src/dxvk/dxvk_compute.h
index f63dbc82f..85ff02bb0 100644
--- a/src/dxvk/dxvk_compute.h
+++ b/src/dxvk/dxvk_compute.h
@@ -58,10 +58,15 @@ namespace dxvk {
      * \returns Pipeline handle
      */
     VkPipeline getPipelineHandle(
-      const DxvkComputePipelineStateInfo& state) const;
+      const DxvkComputePipelineStateInfo& state);
     
   private:
     
+    struct PipelineStruct {
+      DxvkComputePipelineStateInfo stateVector;
+      VkPipeline                   pipeline;
+    };
+    
     const DxvkDevice* const m_device;
     const Rc<vk::DeviceFn>  m_vkd;
     
@@ -69,9 +74,16 @@ namespace dxvk {
     Rc<DxvkPipelineLayout>  m_layout;
     Rc<DxvkShaderModule>    m_cs;
     
-    VkPipeline m_pipeline = VK_NULL_HANDLE;
+    std::mutex                  m_mutex;
+    std::vector<PipelineStruct> m_pipelines;
     
-    void compilePipeline();
+    VkPipeline m_basePipeline = VK_NULL_HANDLE;
+    
+    VkPipeline compilePipeline(
+      const DxvkComputePipelineStateInfo& state,
+            VkPipeline                    baseHandle) const;
+    
+    void destroyPipelines();
     
   };