source/blender/gpu/CMakeLists.txt

@@ -531,6 +531,7 @@ set(GLSL_SRC_TEST
   tests/shaders/gpu_compute_ssbo_test.glsl
   tests/shaders/gpu_compute_vbo_test.glsl
   tests/shaders/gpu_compute_dummy_test.glsl
+  tests/shaders/gpu_push_constants_test.glsl
 )
 
 set(MTL_BACKEND_GLSL_SRC

source/blender/gpu/shaders/infos/gpu_shader_test_info.hh

@@ -55,6 +55,19 @@ GPU_SHADER_CREATE_INFO(gpu_compute_ssbo_binding_test)
     .compute_source("gpu_compute_dummy_test.glsl")
     .do_static_compilation(true);
 
+GPU_SHADER_CREATE_INFO(gpu_push_constants_base_test)
+    .local_group_size(1)
+    .storage_buf(0, Qualifier::WRITE, "float", "data_out[]")
+    .compute_source("gpu_push_constants_test.glsl");
+
+GPU_SHADER_CREATE_INFO(gpu_push_constants_packed_test)
+    .additional_info("gpu_push_constants_base_test")
+    .push_constant(Type::FLOAT, "float_in")
+    .push_constant(Type::VEC2, "vec2_in")
+    .push_constant(Type::VEC3, "vec3_in")
+    .push_constant(Type::VEC4, "vec4_in")
+    .do_static_compilation(true);
+
 GPU_SHADER_CREATE_INFO(eevee_shadow_test)
     .fragment_source("eevee_shadow_test.glsl")
     .additional_info("gpu_shader_test")

source/blender/gpu/tests/gpu_push_constants_test.cc

@@ -2,18 +2,79 @@
 
 #include "testing/testing.h"
 
+#include "GPU_capabilities.h"
+#include "GPU_compute.h"
+#include "GPU_shader.h"
 #include "GPU_storage_buffer.h"
 
+#include "BLI_math_vector.hh"
 #include "BLI_vector.hh"
 
 #include "gpu_testing.hh"
 
 namespace blender::gpu::tests {
 
-static void test_push_constants()
+static void push_constants(const char *info_name)
 {
+  if (!GPU_compute_shader_support() && !GPU_shader_storage_buffer_objects_support()) {
+    /* We can't test as a the platform does not support compute shaders. */
+    std::cout << "Skipping test: platform not supported";
+    return;
+  }
+
+  static constexpr uint SIZE = 16;
+
+  /* Build compute shader. */
+  GPUShader *shader = GPU_shader_create_from_info_name(info_name);
+  EXPECT_NE(shader, nullptr);
+  GPU_shader_bind(shader);
+
+  /* Construct IBO. */
+  GPUStorageBuf *ssbo = GPU_storagebuf_create_ex(
+      SIZE * sizeof(float), nullptr, GPU_USAGE_DEVICE_ONLY, __func__);
+  GPU_storagebuf_bind(ssbo, GPU_shader_get_ssbo_binding(shader, "data_out"));
+
+  const float float_in = 10.0f;
+  const float2 vec2_in(20.0f, 21.0f);
+  const float3 vec3_in(30.0f, 31.0f, 32.0f);
+  const float4 vec4_in(40.0f, 41.0f, 42.0f, 43.0f);
+  GPU_shader_uniform_1f(shader, "float_in", float_in);
+  GPU_shader_uniform_2fv(shader, "vec2_in", vec2_in);
+  GPU_shader_uniform_3fv(shader, "vec3_in", vec3_in);
+  GPU_shader_uniform_4fv(shader, "vec4_in", vec4_in);
+
+  /* Dispatch compute task. */
+  GPU_compute_dispatch(shader, 1, 1, 1);
+
+  /* Check if compute has been done. */
+  GPU_memory_barrier(GPU_BARRIER_SHADER_STORAGE);
+
+  /* Download the index buffer. */
+  float data[SIZE];
+  GPU_storagebuf_read(ssbo, data);
+
+  /* Check the results. */
+  EXPECT_EQ(data[0], float_in);
+  EXPECT_EQ(data[1], vec2_in.x);
+  EXPECT_EQ(data[2], vec2_in.y);
+  EXPECT_EQ(data[3], vec3_in.x);
+  EXPECT_EQ(data[4], vec3_in.y);
+  EXPECT_EQ(data[5], vec3_in.z);
+  EXPECT_EQ(data[6], vec4_in.x);
+  EXPECT_EQ(data[7], vec4_in.y);
+  EXPECT_EQ(data[8], vec4_in.z);
+  EXPECT_EQ(data[9], vec4_in.w);
+
+  /* Cleanup. */
+  GPU_shader_unbind();
+  GPU_storagebuf_free(ssbo);
+  GPU_shader_free(shader);
 }
 
-GPU_TEST(push_constants);
+static void test_push_constants_packed()
+{
+  push_constants("gpu_push_constants_packed_test");
+}
+GPU_TEST(push_constants_packed)
 
 }  // namespace blender::gpu::tests

source/blender/gpu/tests/shaders/gpu_push_constants_test.glsl

@@ -0,0 +1,16 @@
+void main()
+{
+  data_out[0] = float_in;
+
+  data_out[1] = vec2_in.x;
+  data_out[2] = vec2_in.y;
+
+  data_out[3] = vec3_in.x;
+  data_out[4] = vec3_in.y;
+  data_out[5] = vec3_in.z;
+
+  data_out[6] = vec4_in.x;
+  data_out[7] = vec4_in.y;
+  data_out[8] = vec4_in.z;
+  data_out[9] = vec4_in.w;
+}

source/blender/gpu/vulkan/vk_backend.cc

@@ -70,6 +70,8 @@ void VKBackend::compute_dispatch(int groups_x_len, int groups_y_len, int groups_
   descriptor_set.update(context.device_get());
   command_buffer.bind(
       descriptor_set, shader->vk_pipeline_layout_get(), VK_PIPELINE_BIND_POINT_COMPUTE);
+  command_buffer.push_constants(
+      pipeline.push_constants_get(), shader->vk_pipeline_layout_get(), VK_SHADER_STAGE_ALL);
   command_buffer.dispatch(groups_x_len, groups_y_len, groups_z_len);
 }
 

source/blender/gpu/vulkan/vk_command_buffer.cc

@@ -61,6 +61,7 @@ void VKCommandBuffer::bind(const VKPipeline &pipeline, VkPipelineBindPoint bind_
 {
   vkCmdBindPipeline(vk_command_buffer_, bind_point, pipeline.vk_handle());
 }
+
 void VKCommandBuffer::bind(const VKDescriptorSet &descriptor_set,
                            const VkPipelineLayout vk_pipeline_layout,
                            VkPipelineBindPoint bind_point)
@@ -70,6 +71,21 @@ void VKCommandBuffer::bind(const VKDescriptorSet &descriptor_set,
       vk_command_buffer_, bind_point, vk_pipeline_layout, 0, 1, &vk_descriptor_set, 0, 0);
 }
 
+void VKCommandBuffer::push_constants(const VKPushConstants &push_constants,
+                                     const VkPipelineLayout vk_pipeline_layout,
+                                     const VkShaderStageFlags vk_shader_stages)
+{
+  if (push_constants.size_in_bytes() == 0) {
+    return;
+  }
+  vkCmdPushConstants(vk_command_buffer_,
+                     vk_pipeline_layout,
+                     vk_shader_stages,
+                     push_constants.offset(),
+                     push_constants.size_in_bytes(),
+                     push_constants.data());
+}
+
 void VKCommandBuffer::copy(VKBuffer &dst_buffer,
                            VKTexture &src_texture,
                            Span<VkBufferImageCopy> regions)

source/blender/gpu/vulkan/vk_command_buffer.hh

@@ -13,6 +13,7 @@
 namespace blender::gpu {
 class VKBuffer;
 class VKTexture;
+class VKPushConstants;
 
 /** Command buffer to keep track of the life-time of a command buffer.*/
 class VKCommandBuffer : NonCopyable, NonMovable {
@@ -33,6 +34,9 @@ class VKCommandBuffer : NonCopyable, NonMovable {
   void bind(const VKDescriptorSet &descriptor_set,
             const VkPipelineLayout vk_pipeline_layout,
             VkPipelineBindPoint bind_point);
+  void push_constants(const VKPushConstants &push_constants,
+                      const VkPipelineLayout vk_pipeline_layout,
+                      const VkShaderStageFlags vk_shader_stages);
   void dispatch(int groups_x_len, int groups_y_len, int groups_z_len);
   /* Copy the contents of a texture mip level to the dst buffer.*/
   void copy(VKBuffer &dst_buffer, VKTexture &src_texture, Span<VkBufferImageCopy> regions);

source/blender/gpu/vulkan/vk_context.cc

@@ -30,6 +30,7 @@ VKContext::VKContext(void *ghost_window, void *ghost_context)
                          &vk_device_,
                          &vk_queue_family_,
                          &vk_queue_);
+  init_physical_device_limits();
 
   /* Initialize the memory allocator. */
   VmaAllocatorCreateInfo info = {};
@@ -54,6 +55,13 @@ VKContext::~VKContext()
   vmaDestroyAllocator(mem_allocator_);
 }
 
+void VKContext::init_physical_device_limits()
+{
+  VkPhysicalDeviceProperties properties = {};
+  vkGetPhysicalDeviceProperties(vk_physical_device_, &properties);
+  vk_physical_device_limits_ = properties.limits;
+}
+
 void VKContext::activate()
 {
 }

source/blender/gpu/vulkan/vk_context.hh

@@ -30,6 +30,9 @@ class VKContext : public Context {
   VmaAllocator mem_allocator_ = VK_NULL_HANDLE;
   VKDescriptorPools descriptor_pools_;
 
+  /** Limits of the device linked to this context. */
+  VkPhysicalDeviceLimits vk_physical_device_limits_;
+
   void *ghost_context_;
 
  public:
@@ -59,6 +62,11 @@ class VKContext : public Context {
     return vk_physical_device_;
   }
 
+  const VkPhysicalDeviceLimits &physical_device_limits_get() const
+  {
+    return vk_physical_device_limits_;
+  }
+
   VkDevice device_get() const
   {
     return vk_device_;
@@ -88,6 +96,9 @@ class VKContext : public Context {
   {
     return mem_allocator_;
   }
+
+ private:
+  void init_physical_device_limits();
 };
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_descriptor_set.cc

@@ -14,6 +14,15 @@
 #include "BLI_assert.h"
 
 namespace blender::gpu {
+
+VKDescriptorSet::VKDescriptorSet(VKDescriptorSet &&other)
+    : vk_descriptor_pool_(other.vk_descriptor_pool_),
+      vk_descriptor_set_(other.vk_descriptor_set_),
+      bindings_(std::move(other.bindings_))
+{
+  other.mark_freed();
+}
+
 VKDescriptorSet::~VKDescriptorSet()
 {
   if (vk_descriptor_set_ != VK_NULL_HANDLE) {

source/blender/gpu/vulkan/vk_descriptor_set.hh

@@ -108,6 +108,7 @@ class VKDescriptorSet : NonCopyable {
       : vk_descriptor_pool_(vk_descriptor_pool), vk_descriptor_set_(vk_descriptor_set)
   {
   }
+  VKDescriptorSet(VKDescriptorSet &&other);
   virtual ~VKDescriptorSet();
 
   VKDescriptorSet &operator=(VKDescriptorSet &&other)

source/blender/gpu/vulkan/vk_pipeline.cc

@@ -11,10 +11,13 @@
 
 namespace blender::gpu {
 
-VKPipeline::VKPipeline(VkPipeline vk_pipeline, VKDescriptorSet &&vk_descriptor_set)
-    : vk_pipeline_(vk_pipeline)
+VKPipeline::VKPipeline(VkPipeline vk_pipeline,
+                       VKDescriptorSet &&descriptor_set,
+                       VKPushConstants &&push_constants)
+    : vk_pipeline_(vk_pipeline),
+      descriptor_set_(std::move(descriptor_set)),
+      push_constants_(std::move(push_constants))
 {
-  descriptor_set_ = std::move(vk_descriptor_set);
 }
 
 VKPipeline::~VKPipeline()
@@ -29,7 +32,8 @@ VKPipeline::~VKPipeline()
 VKPipeline VKPipeline::create_compute_pipeline(VKContext &context,
                                                VkShaderModule compute_module,
                                                VkDescriptorSetLayout &descriptor_set_layout,
-                                               VkPipelineLayout &pipeline_layout)
+                                               VkPipelineLayout &pipeline_layout,
+                                               VKPushConstantsLayout &push_constants_layout)
 {
   VK_ALLOCATION_CALLBACKS
   VkDevice vk_device = context.device_get();
@@ -44,15 +48,17 @@ VKPipeline VKPipeline::create_compute_pipeline(VKContext &context,
   pipeline_info.layout = pipeline_layout;
   pipeline_info.stage.pName = "main";
 
-  VkPipeline pipeline;
+  VkPipeline vk_pipeline;
   if (vkCreateComputePipelines(
-          vk_device, nullptr, 1, &pipeline_info, vk_allocation_callbacks, &pipeline) !=
+          vk_device, nullptr, 1, &pipeline_info, vk_allocation_callbacks, &vk_pipeline) !=
       VK_SUCCESS) {
     return VKPipeline();
   }
 
   VKDescriptorSet descriptor_set = context.descriptor_pools_get().allocate(descriptor_set_layout);
-  return VKPipeline(pipeline, std::move(descriptor_set));
+  VKPushConstants push_constants(push_constants_layout);
+  return VKPipeline(
+      vk_pipeline, std::move(descriptor_set), std::move(push_constants));
 }
 
 VkPipeline VKPipeline::vk_handle() const

source/blender/gpu/vulkan/vk_pipeline.hh

@@ -7,42 +7,55 @@
 
 #pragma once
 
+#include <optional>
+
 #include "BLI_utility_mixins.hh"
 #include "BLI_vector.hh"
 
 #include "vk_common.hh"
 #include "vk_descriptor_set.hh"
+#include "vk_push_constants.hh"
 
 namespace blender::gpu {
 class VKContext;
 
 class VKPipeline : NonCopyable {
-  VKDescriptorSet descriptor_set_;
   VkPipeline vk_pipeline_ = VK_NULL_HANDLE;
+  VKDescriptorSet descriptor_set_;
+  VKPushConstants push_constants_;
 
  public:
   VKPipeline() = default;
 
   virtual ~VKPipeline();
-  VKPipeline(VkPipeline vk_pipeline, VKDescriptorSet &&vk_descriptor_set);
+  VKPipeline(VkPipeline vk_pipeline,
+             VKDescriptorSet &&vk_descriptor_set,
+             VKPushConstants &&push_constants);
   VKPipeline &operator=(VKPipeline &&other)
   {
     vk_pipeline_ = other.vk_pipeline_;
     other.vk_pipeline_ = VK_NULL_HANDLE;
     descriptor_set_ = std::move(other.descriptor_set_);
+    push_constants_ = std::move(other.push_constants_);
     return *this;
   }
 
   static VKPipeline create_compute_pipeline(VKContext &context,
                                             VkShaderModule compute_module,
                                             VkDescriptorSetLayout &descriptor_set_layout,
-                                            VkPipelineLayout &pipeline_layouts);
+                                            VkPipelineLayout &pipeline_layouts,
+                                            VKPushConstantsLayout &push_constants_layout);
 
   VKDescriptorSet &descriptor_set_get()
   {
     return descriptor_set_;
   }
 
+  VKPushConstants &push_constants_get()
+  {
+    return push_constants_;
+  }
+
   VkPipeline vk_handle() const;
   bool is_valid() const;
 };

source/blender/gpu/vulkan/vk_push_constants.cc

@@ -71,9 +71,12 @@ static uint32_t to_alignment(const shader::Type type)
 }
 
 static VKPushConstantsLayout::PushConstantLayout init_constant(
-    const shader::ShaderCreateInfo::PushConst &push_constant, uint32_t *r_offset)
+    const shader::ShaderCreateInfo::PushConst &push_constant,
+    const ShaderInput &shader_input,
+    uint32_t *r_offset)
 {
   VKPushConstantsLayout::PushConstantLayout layout;
+  layout.location = shader_input.location;
   layout.type = push_constant.type;
   layout.array_size = push_constant.array_size;
   layout.offset = *r_offset;
@@ -101,9 +104,50 @@ void VKPushConstantsLayout::init(const shader::ShaderCreateInfo &info,
   BLI_assert(push_constants.is_empty());
   uint32_t offset = 0;
   for (const shader::ShaderCreateInfo::PushConst &push_constant : info.push_constants_) {
-    push_constants.append(init_constant(push_constant, &offset));
+    const ShaderInput *shader_input = interface.uniform_get(push_constant.name.c_str());
+    BLI_assert(shader_input);
+    push_constants.append(init_constant(push_constant, *shader_input, &offset));
   }
   size_in_bytes_ = offset;
 }
 
+const VKPushConstantsLayout::PushConstantLayout *VKPushConstantsLayout::find(
+    int32_t location) const
+{
+  for (const PushConstantLayout &push_constant : push_constants) {
+    if (push_constant.location == location) {
+      return &push_constant;
+    }
+  }
+  return nullptr;
+}
+
+VKPushConstants::VKPushConstants(VKPushConstantsLayout &layout) : layout_(layout)
+{
+  data_ = MEM_mallocN(layout.size_in_bytes(), __func__);
+}
+
+VKPushConstants::VKPushConstants(VKPushConstants &&other) : layout_(other.layout_)
+{
+  data_ = other.data_;
+  other.data_ = nullptr;
+}
+
+VKPushConstants::~VKPushConstants()
+{
+  if (data_ != nullptr) {
+    MEM_freeN(data_);
+    data_ = nullptr;
+  }
+}
+
+VKPushConstants &VKPushConstants::operator=(VKPushConstants &&other)
+{
+  layout_ = other.layout_;
+  data_ = other.data_;
+  other.data_ = nullptr;
+
+  return *this;
+}
+
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_push_constants.hh

@@ -7,6 +7,7 @@
 
 #pragma once
 
+#include "BLI_utility_mixins.hh"
 #include "BLI_vector.hh"
 
 #include "vk_common.hh"
@@ -16,6 +17,11 @@ namespace blender::gpu {
 
 struct VKPushConstantsLayout {
   struct PushConstantLayout {
+    /* TODO: location requires sequential lookups, we should make the location index based for
+     * quicker access. */
+    int32_t location;
+
+    /** Offset in the push constant data (in bytes). */
     uint32_t offset;
     shader::Type type;
     int array_size;
@@ -32,13 +38,59 @@ struct VKPushConstantsLayout {
   {
     return size_in_bytes_;
   }
+
+  const PushConstantLayout *find(int32_t location) const;
 };
 
-struct VKPushConstants {
-  const VKPushConstantsLayout &layout_;
-  uint32_t *data = nullptr;
+static VKPushConstantsLayout dummy_layout;
+class VKPushConstants : NonCopyable {
 
-  VKPushConstants(const VKPushConstantsLayout &layout);
+ private:
+  VKPushConstantsLayout &layout_ = dummy_layout;
+  void *data_ = nullptr;
+
+ public:
+  VKPushConstants() = default;
+  VKPushConstants(VKPushConstantsLayout &layout);
+  VKPushConstants(VKPushConstants &&other);
+  virtual ~VKPushConstants();
+
+  VKPushConstants &operator=(VKPushConstants &&other);
+
+  size_t offset() const
+  {
+    return 0;
+  }
+
+  size_t size_in_bytes() const
+  {
+    return layout_.size_in_bytes();
+  }
+
+  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 VKPushConstantsLayout::PushConstantLayout *push_constant_layout = layout_.find(location);
+    if (push_constant_layout == nullptr) {
+      /* Currently the builtin uniforms are set using a predefined location each time a shader is
+       * bound.*/
+      return;
+    }
+    BLI_assert_msg(push_constant_layout->offset + comp_len * array_size * sizeof(T) <=
+                       layout_.size_in_bytes(),
+                   "Tried to write outside the push constant allocated memory.");
+    uint8_t *bytes = static_cast<uint8_t *>(data_);
+    T *dst = static_cast<T *>(static_cast<void *>(&bytes[push_constant_layout->offset]));
+    memcpy(dst, input_data, comp_len * array_size * sizeof(T));
+  }
 };
 
 }  // namespace blender::gpu

source/blender/gpu/vulkan/vk_shader.cc

@@ -700,7 +700,7 @@ bool VKShader::finalize(const shader::ShaderCreateInfo *info)
     BLI_assert(fragment_module_ == VK_NULL_HANDLE);
     BLI_assert(compute_module_ != VK_NULL_HANDLE);
     compute_pipeline_ = VKPipeline::create_compute_pipeline(
-        *context_, compute_module_, layout_, pipeline_layout_);
+        *context_, compute_module_, layout_, pipeline_layout_, push_constants_layout_);
     result = compute_pipeline_.is_valid();
   }
 
@@ -976,12 +976,11 @@ void VKShader::unbind()
   }
 }
 
-void VKShader::uniform_float(int /*location*/,
-                             int /*comp_len*/,
-                             int /*array_size*/,
-                             const float * /*data*/)
+void VKShader::uniform_float(int location, int comp_len, int array_size, const float *data)
 {
+  pipeline_get().push_constants_get().push_constant_set(location, comp_len, array_size, data);
 }
+
 void VKShader::uniform_int(int /*location*/,
                            int /*comp_len*/,
                            int /*array_size*/,
@@ -991,6 +990,9 @@ void VKShader::uniform_int(int /*location*/,
 
 std::string VKShader::resources_declare(const shader::ShaderCreateInfo &info) const
 {
+  if (info.name_ == "workbench_next_prepass_ptcloud_opaque_flat_texture_clip") {
+    printf("%s\n", info.name_.c_str());
+  }
   VKShaderInterface interface;
   interface.init(info);
   std::stringstream ss;

source/blender/gpu/vulkan/vk_shader_interface.cc

@@ -14,7 +14,7 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
   using namespace blender::gpu::shader;
 
   attr_len_ = 0;
-  uniform_len_ = 0;
+  uniform_len_ = info.push_constants_.size();
   ssbo_len_ = 0;
   ubo_len_ = 0;
   image_offset_ = -1;
@@ -51,7 +51,7 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
   name_buffer_ = (char *)MEM_mallocN(info.interface_names_size_, "name_buffer");
   uint32_t name_buffer_offset = 0;
 
-  int location = 0;
+  int32_t location = 0;
 
   /* Uniform blocks */
   for (const ShaderCreateInfo::Resource &res : all_resources) {
@@ -79,6 +79,18 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
     }
   }
 
+  /* Push constants. */
+  /* NOTE: Push constants must be added after other uniform resources as resources have strict
+   * rules for their 'location' due to descriptor sets. Push constants only need an unique location
+   * as it is only used by the GPU module internally.*/
+  int32_t push_constant_location = location + 1024;
+  for (const ShaderCreateInfo::PushConst &push_constant : info.push_constants_) {
+    copy_input_name(input, push_constant.name, name_buffer_, name_buffer_offset);
+    input->location = push_constant_location++;
+    input->binding = -1;
+    input++;
+  }
+
   /* Storage buffers */
   for (const ShaderCreateInfo::Resource &res : all_resources) {
     if (res.bind_type == ShaderCreateInfo::Resource::BindType::STORAGE_BUFFER) {