vk_push_constants.hh

Go to the documentation of this file.

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

 
#pragma once
 
#include "BLI_utility_mixins.hh"
#include "BLI_vector.hh"
 
#include "gpu_shader_create_info.hh"
 
#include "vk_common.hh"
#include "vk_descriptor_set.hh"
 
namespace blender::gpu {
class VKShaderInterface;
class VKUniformBuffer;
class VKContext;
class VKDevice;

class VKPushConstants {
  friend class VKContext;
 
 public:
  enum class StorageType {
    NONE,

    PUSH_CONSTANTS,

    UNIFORM_BUFFER,
  };

  struct Layout {
    static constexpr StorageType STORAGE_TYPE_DEFAULT = StorageType::PUSH_CONSTANTS;
    static constexpr StorageType STORAGE_TYPE_FALLBACK = StorageType::UNIFORM_BUFFER;
 
    struct PushConstant {
      /* Used as lookup based on ShaderInput. */
      int32_t location;

      uint32_t offset;
      shader::Type type;
      int array_size;
      uint inner_row_padding;
    };
 
   private:
    Vector<PushConstant> push_constants;
    uint32_t size_in_bytes_ = 0;
    StorageType storage_type_ = StorageType::NONE;
    VKDescriptorSet::Location descriptor_set_location_;
 
   public:
    static StorageType determine_storage_type(const shader::ShaderCreateInfo &info,
                                              const VKDevice &device);

    void init(const shader::ShaderCreateInfo &info,
              const VKShaderInterface &interface,
              StorageType storage_type,
              VKDescriptorSet::Location location);

    StorageType storage_type_get() const
    {
      return storage_type_;
    }

    VKDescriptorSet::Location descriptor_set_location_get() const
    {
      return descriptor_set_location_;
    }

    uint32_t size_in_bytes() const
    {
      return size_in_bytes_;
    }

    const PushConstant *find(int32_t location) const;
 
    void debug_print() const;
  };
 
 private:
  const Layout *layout_ = nullptr;
  void *data_ = nullptr;

  std::unique_ptr<VKUniformBuffer> uniform_buffer_;
 
 public:
  VKPushConstants();
  VKPushConstants(const Layout *layout);
  VKPushConstants(VKPushConstants &&other);
  virtual ~VKPushConstants();
 
  VKPushConstants &operator=(VKPushConstants &&other);
 
  size_t offset() const
  {
    return 0;
  }
 
  const Layout &layout_get() const
  {
    return *layout_;
  }

  const void *data() const
  {
    return data_;
  }

  template<typename T>
  void push_constant_set(int32_t location,
                         int32_t comp_len,
                         int32_t array_size,
                         const T *input_data)
  {
    const Layout::PushConstant *push_constant_layout = layout_->find(location);
    if (push_constant_layout == nullptr) {
      /* Legacy code can still try to update push constants when they don't exist. For example
       * `immDrawPixelsTexSetup` will bind an image slot manually. This works in OpenGL, but in
       * vulkan images aren't stored as push constants. */
      return;
    }
 
    uint8_t *bytes = static_cast<uint8_t *>(data_);
    T *dst = static_cast<T *>(static_cast<void *>(&bytes[push_constant_layout->offset]));
    const int inner_row_padding = push_constant_layout->inner_row_padding;
    const bool is_tightly_std140_packed = (comp_len % 4) == 0;
    /* Vec3[] are not tightly packed in std430. */
    const bool is_tightly_std430_packed = comp_len != 3 || array_size == 0;
    if (inner_row_padding == 0 &&
        ((layout_->storage_type_get() == StorageType::PUSH_CONSTANTS &&
          is_tightly_std430_packed) ||
         array_size == 0 || push_constant_layout->array_size == 0 || is_tightly_std140_packed))
    {
      const size_t copy_size_in_bytes = comp_len * max_ii(array_size, 1) * sizeof(T);
      BLI_assert_msg(push_constant_layout->offset + copy_size_in_bytes <= layout_->size_in_bytes(),
                     "Tried to write outside the push constant allocated memory.");
      memcpy(dst, input_data, copy_size_in_bytes);
      return;
    }
 
    /* Store elements in uniform buffer as array. In Std140 arrays have an element stride of 16
     * bytes. */
    BLI_assert(sizeof(T) == 4);
    const T *src = input_data;
    if (inner_row_padding == 0) {
      for (const int i : IndexRange(array_size)) {
        UNUSED_VARS(i);
        memcpy(dst, src, comp_len * sizeof(T));
        src += comp_len;
        dst += 4;
      }
    }
    else {
      BLI_assert_msg(array_size == 1, "No support for MAT3 arrays, but can be added when needed");
      for (const int component_index : IndexRange(comp_len)) {
        *dst = *src;
        dst += 1;
        src += 1;
        if ((component_index % inner_row_padding) == (inner_row_padding - 1)) {
          dst += 1;
        }
      }
    }
  }

  void update_uniform_buffer();

  std::unique_ptr<VKUniformBuffer> &uniform_buffer_get();
};
 
}  // namespace blender::gpu