[dxvk] Allow creating pipeline libraries with more than one shader

src/dxvk/dxvk_pipemanager.cpp

@@ -248,8 +248,11 @@ namespace dxvk {
     if (pair != m_computePipelines.end())
       return &pair->second;
 
+    DxvkShaderPipelineLibraryKey key;
+    key.addShader(shaders.cs);
+
     auto layout = createPipelineLayout(shaders.cs->getBindings());
-    auto library = findPipelineLibraryLocked(shaders.cs);
+    auto library = findPipelineLibraryLocked(key);
 
     auto iter = m_computePipelines.emplace(
       std::piecewise_construct,
@@ -290,9 +293,28 @@ namespace dxvk {
     DxvkShaderPipelineLibrary* vsLibrary = nullptr;
     DxvkShaderPipelineLibrary* fsLibrary = nullptr;
 
-    if (shaders.tcs == nullptr && shaders.tes == nullptr && shaders.gs == nullptr) {
-      vsLibrary = findPipelineLibraryLocked(shaders.vs);
-      fsLibrary = findPipelineLibraryLocked(shaders.fs);
+    if (m_device->canUseGraphicsPipelineLibrary()) {
+      DxvkShaderPipelineLibraryKey vsKey;
+      vsKey.addShader(shaders.vs);
+
+      if (shaders.tcs != nullptr)
+        vsKey.addShader(shaders.tcs);
+      if (shaders.tes != nullptr)
+        vsKey.addShader(shaders.tes);
+      if (shaders.gs != nullptr)
+        vsKey.addShader(shaders.gs);
+
+      if (vsKey.canUsePipelineLibrary())
+        vsLibrary = findPipelineLibraryLocked(vsKey);
+
+      if (vsLibrary) {
+        DxvkShaderPipelineLibraryKey fsKey;
+
+        if (shaders.fs != nullptr)
+          fsKey.addShader(shaders.fs);
+
+        fsLibrary = findPipelineLibraryLocked(fsKey);
+      }
     }
 
     auto iter = m_graphicsPipelines.emplace(
@@ -339,7 +361,10 @@ namespace dxvk {
   void DxvkPipelineManager::registerShader(
     const Rc<DxvkShader>&         shader) {
     if (canPrecompileShader(shader)) {
-      auto library = createPipelineLibrary(shader);
+      DxvkShaderPipelineLibraryKey key;
+      key.addShader(shader);
+
+      auto library = createPipelineLibrary(key);
       m_workers.compilePipelineLibrary(library, DxvkPipelinePriority::Normal);
     }
 
@@ -353,7 +378,10 @@ namespace dxvk {
       return;
 
     // Dispatch high-priority compile job
-    auto library = findPipelineLibrary(shader);
+    DxvkShaderPipelineLibraryKey key;
+    key.addShader(shader);
+
+    auto library = findPipelineLibrary(key);
 
     if (library)
       m_workers.compilePipelineLibrary(library, DxvkPipelinePriority::High);
@@ -416,46 +444,49 @@ namespace dxvk {
 
 
   DxvkShaderPipelineLibrary* DxvkPipelineManager::createPipelineLibrary(
-    const Rc<DxvkShader>&     shader) {
+    const DxvkShaderPipelineLibraryKey& key) {
     std::lock_guard<dxvk::mutex> lock(m_mutex);
-    auto layout = createPipelineLayout(shader->getBindings());
+    return createPipelineLibraryLocked(key);
+  }
 
-    DxvkShaderPipelineLibraryKey key;
-    key.shader = shader;
+
+  DxvkShaderPipelineLibrary* DxvkPipelineManager::createPipelineLibraryLocked(
+    const DxvkShaderPipelineLibraryKey& key) {
+    auto bindings = key.getBindings();
+    auto layout = createPipelineLayout(bindings);
 
     auto iter = m_shaderLibraries.emplace(
       std::piecewise_construct,
       std::tuple(key),
-      std::tuple(m_device, this, shader.ptr(), layout));
+      std::tuple(m_device, this, key, layout));
     return &iter.first->second;
   }
 
 
   DxvkShaderPipelineLibrary* DxvkPipelineManager::createNullFsPipelineLibrary() {
     std::lock_guard<dxvk::mutex> lock(m_mutex);
+    DxvkShaderPipelineLibraryKey key;
+
     DxvkBindingLayout bindings(VK_SHADER_STAGE_FRAGMENT_BIT);
     auto layout = createPipelineLayout(bindings);
 
     auto iter = m_shaderLibraries.emplace(
       std::piecewise_construct,
       std::tuple(),
-      std::tuple(m_device, this, nullptr, layout));
+      std::tuple(m_device, this, key, layout));
     return &iter.first->second;
   }
 
 
   DxvkShaderPipelineLibrary* DxvkPipelineManager::findPipelineLibrary(
-    const Rc<DxvkShader>&     shader) {
+    const DxvkShaderPipelineLibraryKey& key) {
     std::lock_guard<dxvk::mutex> lock(m_mutex);
-    return findPipelineLibraryLocked(shader);
+    return findPipelineLibraryLocked(key);
   }
 
 
   DxvkShaderPipelineLibrary* DxvkPipelineManager::findPipelineLibraryLocked(
-    const Rc<DxvkShader>&     shader) {
-    DxvkShaderPipelineLibraryKey key;
-    key.shader = shader;
-
+    const DxvkShaderPipelineLibraryKey& key) {
     auto pair = m_shaderLibraries.find(key);
     if (pair == m_shaderLibraries.end())
       return nullptr;
@@ -465,7 +496,7 @@ namespace dxvk {
 
 
   bool DxvkPipelineManager::canPrecompileShader(
-    const Rc<DxvkShader>&     shader) const {
+    const Rc<DxvkShader>& shader) const {
     if (!shader->canUsePipelineLibrary(true))
       return false;
 

src/dxvk/dxvk_pipemanager.h

@@ -295,18 +295,21 @@ namespace dxvk {
       const DxvkBindingLayout& layout);
 
     DxvkShaderPipelineLibrary* createPipelineLibrary(
-      const Rc<DxvkShader>&     shader);
+      const DxvkShaderPipelineLibraryKey& key);
+
+    DxvkShaderPipelineLibrary* createPipelineLibraryLocked(
+      const DxvkShaderPipelineLibraryKey& key);
 
     DxvkShaderPipelineLibrary* createNullFsPipelineLibrary();
 
     DxvkShaderPipelineLibrary* findPipelineLibrary(
-      const Rc<DxvkShader>&     shader);
+      const DxvkShaderPipelineLibraryKey& key);
 
     DxvkShaderPipelineLibrary* findPipelineLibraryLocked(
-      const Rc<DxvkShader>&     shader);
+      const DxvkShaderPipelineLibraryKey& key);
 
     bool canPrecompileShader(
-      const Rc<DxvkShader>&     shader) const;
+      const Rc<DxvkShader>& shader) const;
 
   };
   

src/dxvk/dxvk_shader.cpp

@@ -897,28 +897,108 @@ namespace dxvk {
   }
 
 
-  DxvkShaderPipelineLibrary::DxvkShaderPipelineLibrary(
-    const DxvkDevice*               device,
-          DxvkPipelineManager*      manager,
-          DxvkShader*               shader,
-    const DxvkBindingLayoutObjects* layout)
-  : m_device      (device),
-    m_stats       (&manager->m_stats),
-    m_layout      (layout) {
-    if (shader) {
+  DxvkShaderPipelineLibraryKey::DxvkShaderPipelineLibraryKey() {
+
+  }
+
+
+  DxvkShaderPipelineLibraryKey::~DxvkShaderPipelineLibraryKey() {
+
+  }
+
+
+  DxvkShaderSet DxvkShaderPipelineLibraryKey::getShaderSet() const {
+    DxvkShaderSet result;
+
+    for (uint32_t i = 0; i < m_shaderCount; i++) {
+      auto shader = m_shaders[i].ptr();
+
       switch (shader->info().stage) {
-        case VK_SHADER_STAGE_VERTEX_BIT:
-          m_shaders.vs = shader;
-          break;
-        case VK_SHADER_STAGE_FRAGMENT_BIT:
-          m_shaders.fs = shader;
-          break;
-        case VK_SHADER_STAGE_COMPUTE_BIT:
-          m_shaders.cs = shader;
-          break;
+        case VK_SHADER_STAGE_VERTEX_BIT:                  result.vs = shader; break;
+        case VK_SHADER_STAGE_TESSELLATION_CONTROL_BIT:    result.tcs = shader; break;
+        case VK_SHADER_STAGE_TESSELLATION_EVALUATION_BIT: result.tes = shader; break;
+        case VK_SHADER_STAGE_GEOMETRY_BIT:                result.gs = shader; break;
+        case VK_SHADER_STAGE_FRAGMENT_BIT:                result.fs = shader; break;
+        case VK_SHADER_STAGE_COMPUTE_BIT:                 result.cs = shader; break;
         default: ;
       }
     }
+
+    return result;
+  }
+
+
+  DxvkBindingLayout DxvkShaderPipelineLibraryKey::getBindings() const {
+    DxvkBindingLayout mergedLayout(m_shaderStages);
+
+    for (uint32_t i = 0; i < m_shaderCount; i++)
+      mergedLayout.merge(m_shaders[i]->getBindings());
+
+    return mergedLayout;
+  }
+
+
+  void DxvkShaderPipelineLibraryKey::addShader(
+    const Rc<DxvkShader>&               shader) {
+    m_shaderStages |= shader->info().stage;
+    m_shaders[m_shaderCount++] = shader;
+  }
+
+
+  bool DxvkShaderPipelineLibraryKey::canUsePipelineLibrary() const {
+    // Ensure that each individual shader can be used in a library
+    bool standalone = m_shaderCount <= 1;
+
+    for (uint32_t i = 0; i < m_shaderCount; i++) {
+      if (!m_shaders[i]->canUsePipelineLibrary(standalone))
+        return false;
+    }
+
+    // Ensure that stage I/O is compatible between stages
+    for (uint32_t i = 0; i + 1 < m_shaderCount; i++) {
+      uint32_t currStageIoMask = m_shaders[i]->info().outputMask;
+      uint32_t nextStageIoMask = m_shaders[i + 1]->info().inputMask;
+
+      if ((currStageIoMask & nextStageIoMask) != nextStageIoMask)
+        return false;
+    }
+
+    return true;
+  }
+
+
+  bool DxvkShaderPipelineLibraryKey::eq(
+    const DxvkShaderPipelineLibraryKey& other) const {
+    bool eq = m_shaderStages == other.m_shaderStages;
+
+    for (uint32_t i = 0; i < m_shaderCount && eq; i++)
+      eq = m_shaders[i] == other.m_shaders[i];
+
+    return eq;
+  }
+
+
+  size_t DxvkShaderPipelineLibraryKey::hash() const {
+    DxvkHashState hash;
+    hash.add(uint32_t(m_shaderStages));
+
+    for (uint32_t i = 0; i < m_shaderCount; i++)
+      hash.add(m_shaders[i]->getHash());
+
+    return hash;
+  }
+
+
+  DxvkShaderPipelineLibrary::DxvkShaderPipelineLibrary(
+    const DxvkDevice*               device,
+          DxvkPipelineManager*      manager,
+    const DxvkShaderPipelineLibraryKey& key,
+    const DxvkBindingLayoutObjects* layout)
+  : m_device      (device),
+    m_stats       (&manager->m_stats),
+    m_shaders     (key.getShaderSet()),
+    m_layout      (layout) {
+
   }
 
 

src/dxvk/dxvk_shader.h

@@ -385,22 +385,6 @@ namespace dxvk {
   };
 
 
-  /**
-   * \brief Shader pipeline library key
-   */
-  struct DxvkShaderPipelineLibraryKey {
-    Rc<DxvkShader> shader;
-
-    bool eq(const DxvkShaderPipelineLibraryKey& other) const {
-      return shader == other.shader;
-    }
-
-    size_t hash() const {
-      return DxvkShader::getHash(shader);
-    }
-  };
-
-
   /**
    * \brief Shader set
    *
@@ -433,6 +417,68 @@ namespace dxvk {
   };
 
 
+  /**
+   * \brief Shader pipeline library key
+   */
+  class DxvkShaderPipelineLibraryKey {
+
+  public:
+
+    DxvkShaderPipelineLibraryKey();
+
+    ~DxvkShaderPipelineLibraryKey();
+
+    /**
+     * \brief Creates shader set from key
+     * \returns Shader set
+     */
+    DxvkShaderSet getShaderSet() const;
+
+    /**
+     * \brief Generates merged binding layout
+     * \returns Binding layout
+     */
+    DxvkBindingLayout getBindings() const;
+
+    /**
+     * \brief Adds a shader to the key
+     *
+     * Shaders must be added in stage order.
+     * \param [in] shader Shader to add
+     */
+    void addShader(
+      const Rc<DxvkShader>&               shader);
+
+    /**
+     * \brief Checks wether a pipeline library can be created
+     * \returns \c true if all added shaders are compatible
+     */
+    bool canUsePipelineLibrary() const;
+
+    /**
+     * \brief Checks for equality
+     *
+     * \param [in] other Key to compare to
+     * \returns \c true if the keys are equal
+     */
+    bool eq(
+      const DxvkShaderPipelineLibraryKey& other) const;
+
+    /**
+     * \brief Computes key hash
+     * \returns Key hash
+     */
+    size_t hash() const;
+
+  private:
+
+    uint32_t                      m_shaderCount   = 0;
+    VkShaderStageFlags            m_shaderStages  = 0;
+    std::array<Rc<DxvkShader>, 4> m_shaders;
+
+  };
+
+
   /**
    * \brief Shader pipeline library
    *
@@ -449,7 +495,7 @@ namespace dxvk {
     DxvkShaderPipelineLibrary(
       const DxvkDevice*               device,
             DxvkPipelineManager*      manager,
-            DxvkShader*               shader,
+      const DxvkShaderPipelineLibraryKey& key,
       const DxvkBindingLayoutObjects* layout);
 
     ~DxvkShaderPipelineLibrary();