vk_pipeline_pool.hh

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2024 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

 
#pragma once
 
#include <mutex>
 
#include "BLI_map.hh"
#include "BLI_utility_mixins.hh"
 
#include "gpu_state_private.hh"
 
#include "vk_common.hh"
 
namespace blender {
namespace gpu {

struct VKComputeInfo {
  VkShaderModule vk_shader_module;
  VkPipelineLayout vk_pipeline_layout;
  Vector<shader::SpecializationConstant::Value> specialization_constants;
 
  bool operator==(const VKComputeInfo &other) const
  {
    return vk_shader_module == other.vk_shader_module &&
           vk_pipeline_layout == other.vk_pipeline_layout &&
           specialization_constants == other.specialization_constants;
  };
};

struct VKGraphicsInfo {
  struct VertexIn {
    VkPrimitiveTopology vk_topology;
    Vector<VkVertexInputAttributeDescription> attributes;
    Vector<VkVertexInputBindingDescription> bindings;
 
    bool operator==(const VertexIn &other) const
    {
      /* TODO: use an exact implementation and remove the hash compare. */
#if 0
      return vk_topology == other.vk_topology && attributes.hash() == other.attributes.hash() &&
             bindings.hash() == other.bindings.hash();
#endif
      return hash() == other.hash();
    }
 
    uint64_t hash() const
    {
      uint64_t hash = 0;
      hash = hash * 33 ^ uint64_t(vk_topology);
      for (const VkVertexInputAttributeDescription &attribute : attributes) {
        hash = hash * 33 ^ uint64_t(attribute.location);
        hash = hash * 33 ^ uint64_t(attribute.binding);
        hash = hash * 33 ^ uint64_t(attribute.format);
        hash = hash * 33 ^ uint64_t(attribute.offset);
      }
      for (const VkVertexInputBindingDescription &binding : bindings) {
        hash = hash * 33 ^ uint64_t(binding.binding);
        hash = hash * 33 ^ uint64_t(binding.inputRate);
        hash = hash * 33 ^ uint64_t(binding.stride);
      }
      return hash;
    }
  };
  struct PreRasterization {
    VkShaderModule vk_vertex_module;
    VkShaderModule vk_geometry_module;
    bool operator==(const PreRasterization &other) const
    {
      return vk_vertex_module == other.vk_vertex_module &&
             vk_geometry_module == other.vk_geometry_module;
    }
    uint64_t hash() const
    {
      uint64_t hash = 0;
      hash = hash * 33 ^ uint64_t(vk_vertex_module);
      hash = hash * 33 ^ uint64_t(vk_geometry_module);
      return hash;
    }
  };
  struct FragmentShader {
    VkShaderModule vk_fragment_module;
    Vector<VkViewport> viewports;
    Vector<VkRect2D> scissors;
    std::optional<uint64_t> cached_hash;
 
    bool operator==(const FragmentShader &other) const
    {
      /* TODO: Do not use hash. */
      return vk_fragment_module == other.vk_fragment_module && hash() == other.hash();
    }
 
    uint64_t hash() const
    {
      if (cached_hash.has_value()) {
        return *cached_hash;
      }
      return calc_hash();
    }
 
    void update_hash()
    {
      cached_hash = calc_hash();
    }
 
   private:
    uint64_t calc_hash() const
    {
      uint64_t hash = 0;
      hash = hash * 33 ^ uint64_t(vk_fragment_module);
      for (const VkViewport &vk_viewport : viewports) {
        hash = hash * 33 ^ uint64_t(vk_viewport.x);
        hash = hash * 33 ^ uint64_t(vk_viewport.y);
        hash = hash * 33 ^ uint64_t(vk_viewport.width);
        hash = hash * 33 ^ uint64_t(vk_viewport.height);
        hash = hash * 33 ^ uint64_t(vk_viewport.minDepth);
        hash = hash * 33 ^ uint64_t(vk_viewport.maxDepth);
      }
      for (const VkRect2D &scissor : scissors) {
        hash = hash * 33 ^ uint64_t(scissor.offset.x);
        hash = hash * 33 ^ uint64_t(scissor.offset.y);
        hash = hash * 33 ^ uint64_t(scissor.extent.width);
        hash = hash * 33 ^ uint64_t(scissor.extent.height);
      }
      return hash;
    }
  };
  struct FragmentOut {
    VkFormat depth_attachment_format;
    VkFormat stencil_attachment_format;
    Vector<VkFormat> color_attachment_formats;
 
    bool operator==(const FragmentOut &other) const
    {
      return hash() == other.hash();
    }
 
    uint64_t hash() const
    {
      uint64_t hash = 0;
      hash = hash * 33 ^ uint64_t(depth_attachment_format);
      hash = hash * 33 ^ uint64_t(stencil_attachment_format);
      for (VkFormat color_attachment_format : color_attachment_formats) {
        hash = hash * 33 ^ uint64_t(color_attachment_format);
      }
 
      return hash;
    }
  };
 
  VertexIn vertex_in;
  PreRasterization pre_rasterization;
  FragmentShader fragment_shader;
  FragmentOut fragment_out;
 
  GPUState state;
  GPUStateMutable mutable_state;
  VkPipelineLayout vk_pipeline_layout;
  Vector<shader::SpecializationConstant::Value> specialization_constants;
 
  bool operator==(const VKGraphicsInfo &other) const
  {
    return vertex_in == other.vertex_in && pre_rasterization == other.pre_rasterization &&
           fragment_shader == other.fragment_shader && fragment_out == other.fragment_out &&
           vk_pipeline_layout == other.vk_pipeline_layout &&
           specialization_constants == other.specialization_constants && state == other.state &&
           mutable_state == other.mutable_state;
  };
  uint64_t hash() const
  {
    uint64_t hash = 0;
    hash = hash * 33 ^ vertex_in.hash();
    hash = hash * 33 ^ pre_rasterization.hash();
    hash = hash * 33 ^ fragment_shader.hash();
    hash = hash * 33 ^ fragment_out.hash();
    hash = hash * 33 ^ uint64_t(vk_pipeline_layout);
    hash = hash * 33 ^ get_default_hash(specialization_constants);
    hash = hash * 33 ^ state.data;
    hash = hash * 33 ^ mutable_state.data[0];
    hash = hash * 33 ^ mutable_state.data[1];
    hash = hash * 33 ^ mutable_state.data[2];
    return hash;
  }
};
 
}  // namespace gpu
 
template<> struct DefaultHash<gpu::VKComputeInfo> {
  uint64_t operator()(const gpu::VKComputeInfo &key) const
  {
    uint64_t hash = uint64_t(key.vk_shader_module);
    hash = hash * 33 ^ uint64_t(key.vk_pipeline_layout);
    hash = hash * 33 ^ get_default_hash(key.specialization_constants);
    return hash;
  }
};
 
namespace gpu {
class VKDevice;

class VKPipelinePool : public NonCopyable {
  friend class VKDevice;
 
 public:
 private:
  Map<VKComputeInfo, VkPipeline> compute_pipelines_;
  Map<VKGraphicsInfo, VkPipeline> graphic_pipelines_;
  /* Partially initialized structures to reuse. */
  VkComputePipelineCreateInfo vk_compute_pipeline_create_info_;
 
  VkGraphicsPipelineCreateInfo vk_graphics_pipeline_create_info_;
  VkPipelineRenderingCreateInfo vk_pipeline_rendering_create_info_;
  VkPipelineShaderStageCreateInfo vk_pipeline_shader_stage_create_info_[3];
  VkPipelineInputAssemblyStateCreateInfo vk_pipeline_input_assembly_state_create_info_;
  VkPipelineVertexInputStateCreateInfo vk_pipeline_vertex_input_state_create_info_;
 
  VkPipelineRasterizationStateCreateInfo vk_pipeline_rasterization_state_create_info_;
  VkPipelineViewportStateCreateInfo vk_pipeline_viewport_state_create_info_;
  VkPipelineDepthStencilStateCreateInfo vk_pipeline_depth_stencil_state_create_info_;
 
  VkPipelineMultisampleStateCreateInfo vk_pipeline_multisample_state_create_info_;
 
  Vector<VkPipelineColorBlendAttachmentState> vk_pipeline_color_blend_attachment_states_;
  VkPipelineColorBlendStateCreateInfo vk_pipeline_color_blend_state_create_info_;
  VkPipelineColorBlendAttachmentState vk_pipeline_color_blend_attachment_state_template_;
 
  VkSpecializationInfo vk_specialization_info_;
  Vector<VkSpecializationMapEntry> vk_specialization_map_entries_;
  VkPushConstantRange vk_push_constant_range_;
 
  VkPipelineCache vk_pipeline_cache_static_;
  VkPipelineCache vk_pipeline_cache_non_static_;
 
  std::mutex mutex_;
 
 public:
  VKPipelinePool();
 
  void init();

  VkPipeline get_or_create_compute_pipeline(VKComputeInfo &compute_info,
                                            bool is_static_shader,
                                            VkPipeline vk_pipeline_base);

  VkPipeline get_or_create_graphics_pipeline(VKGraphicsInfo &graphics_info,
                                             bool is_static_shader,
                                             VkPipeline vk_pipeline_base);

  void remove(Span<VkShaderModule> vk_shader_modules);

  void free_data();

  void read_from_disk();

  void write_to_disk();
 
 private:
  VkSpecializationInfo *specialization_info_update(
      Span<shader::SpecializationConstant::Value> specialization_constants);
  void specialization_info_reset();
};
 
}  // namespace gpu
 
}  // namespace blender