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 "xxhash.h"
 
#include "BLI_map.hh"
#include "BLI_mutex.hh"
#include "BLI_utility_mixins.hh"
 
#include "gpu_state_private.hh"
 
#include "vk_common.hh"
 
namespace blender::gpu {
class VKDevice;
class VKDiscardPool;

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;
  };
 
  uint64_t hash() const
  {
    uint64_t hash = uint64_t(vk_shader_module);
    hash = hash * 33 ^ uint64_t(vk_pipeline_layout);
    hash = hash * 33 ^ specialization_constants.hash();
    return hash;
  }
};

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 = uint64_t(vk_topology);
      hash = hash * 33 ^
             XXH3_64bits(attributes.data(),
                         attributes.size() * sizeof(VkVertexInputAttributeDescription));
      hash = hash * 33 ^ XXH3_64bits(bindings.data(),
                                     bindings.size() * sizeof(VkVertexInputBindingDescription));
      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
    {
      if (vk_fragment_module != other.vk_fragment_module ||
          viewports.size() != other.viewports.size() || scissors.size() != other.scissors.size() ||
          hash() != other.hash())
      {
        return false;
      }
 
      return true;
    }
 
    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 = uint64_t(vk_fragment_module);
      hash = hash * 33 ^ uint64_t(viewports.size());
      hash = hash * 33 ^ uint64_t(scissors.size());
 
      return hash;
    }
  };
  struct FragmentOut {
    uint32_t color_attachment_size;
 
    /* Dynamic rendering */
    VkFormat depth_attachment_format;
    VkFormat stencil_attachment_format;
    Vector<VkFormat> color_attachment_formats;
 
    bool operator==(const FragmentOut &other) const
    {
#if 1
      return hash() == other.hash();
#else
      if (depth_attachment_format != other.depth_attachment_format ||
          stencil_attachment_format != other.stencil_attachment_format ||
          color_attachment_formats.size() != other.color_attachment_formats.size())
      {
        return false;
      }
 
      if (memcmp(color_attachment_formats.data(),
                 other.color_attachment_formats.data(),
                 color_attachment_formats.size() * sizeof(VkFormat)) == 0)
      {
        return false;
      }
      return true;
#endif
    }
 
    uint64_t hash() const
    {
      uint64_t hash = uint64_t(depth_attachment_format);
      hash = hash * 33 ^ uint64_t(stencil_attachment_format);
      hash = hash * 33 ^ XXH3_64bits(color_attachment_formats.data(),
                                     color_attachment_formats.size() * sizeof(VkFormat));
      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 ^ specialization_constants.hash();
    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;
  }
};

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_;
  VkPipelineRasterizationProvokingVertexStateCreateInfoEXT
      vk_pipeline_rasterization_provoking_vertex_state_info_;
 
  Vector<VkDynamicState> vk_dynamic_states_;
  VkPipelineDynamicStateCreateInfo vk_pipeline_dynamic_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_;
 
  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 discard(VKDiscardPool &discard_pool, VkPipelineLayout vk_pipeline_layout);

  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 blender::gpu