source/blender/gpu/CMakeLists.txt

@@ -201,9 +201,11 @@ set(VULKAN_SRC
   vulkan/vk_fence.cc
   vulkan/vk_framebuffer.cc
   vulkan/vk_index_buffer.cc
+  vulkan/vk_memory_layout.cc
   vulkan/vk_memory.cc
   vulkan/vk_pipeline.cc
   vulkan/vk_pixel_buffer.cc
+  vulkan/vk_push_constants.cc
   vulkan/vk_query.cc
   vulkan/vk_shader.cc
   vulkan/vk_shader_interface.cc
@@ -226,9 +228,11 @@ set(VULKAN_SRC
   vulkan/vk_fence.hh
   vulkan/vk_framebuffer.hh
   vulkan/vk_index_buffer.hh
+  vulkan/vk_memory_layout.hh
   vulkan/vk_memory.hh
   vulkan/vk_pipeline.hh
   vulkan/vk_pixel_buffer.hh
+  vulkan/vk_push_constants.hh
   vulkan/vk_query.hh
   vulkan/vk_shader.hh
   vulkan/vk_shader_interface.hh
@@ -532,6 +536,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
@@ -804,12 +809,19 @@ if(WITH_GTESTS)
       tests/gpu_testing.cc
 
       tests/gpu_index_buffer_test.cc
+      tests/gpu_push_constants_test.cc
       tests/gpu_shader_builtin_test.cc
       tests/gpu_shader_test.cc
       tests/gpu_storage_buffer_test.cc
 
       tests/gpu_testing.hh
     )
+    if(WITH_VULKAN_BACKEND)
+      list(APPEND TEST_SRC
+        tests/memory_layout_test.cc
+      )
+    endif()
+
     set(TEST_INC
     )
     set(TEST_LIB

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

@@ -55,6 +55,36 @@ GPU_SHADER_CREATE_INFO(gpu_compute_ssbo_binding_test)
     .compute_source("gpu_compute_dummy_test.glsl")
     .do_static_compilation(true);
 
+/* Push constants*/
+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_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);
+
+/* Push constants size test. */
+GPU_SHADER_CREATE_INFO(gpu_push_constants_128bytes_test)
+    .additional_info("gpu_push_constants_test")
+    .push_constant(Type::FLOAT, "filler", 20)
+    .do_static_compilation(true);
+
+GPU_SHADER_CREATE_INFO(gpu_push_constants_256bytes_test)
+    .additional_info("gpu_push_constants_128bytes_test")
+    .push_constant(Type::FLOAT, "filler2", 32)
+    .do_static_compilation(true);
+
+GPU_SHADER_CREATE_INFO(gpu_push_constants_512bytes_test)
+    .additional_info("gpu_push_constants_256bytes_test")
+    .push_constant(Type::FLOAT, "filler3", 64)
+    .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

@@ -0,0 +1,210 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+#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_utility_mixins.hh"
+#include "BLI_vector.hh"
+
+#include "gpu_testing.hh"
+
+namespace blender::gpu::tests {
+struct CallData {
+  GPUStorageBuf *ssbo = nullptr;
+  Vector<float> data;
+
+  float float_in;
+  float2 vec2_in;
+  float3 vec3_in;
+  float4 vec4_in;
+
+  void init_ssbo(size_t num_floats)
+  {
+    if (ssbo == nullptr) {
+      ssbo = GPU_storagebuf_create_ex(
+          num_floats * sizeof(float), nullptr, GPU_USAGE_DEVICE_ONLY, __func__);
+      data.resize(num_floats);
+    }
+  }
+
+  ~CallData()
+  {
+    if (ssbo != nullptr) {
+      GPU_storagebuf_free(ssbo);
+      ssbo = nullptr;
+    }
+  }
+
+  void generate_test_data(const float vector_mul, const float scalar_mul)
+  {
+    float_in = vector_mul;
+    vec2_in = float2(vector_mul * 2.0, vector_mul * 2.0 + scalar_mul);
+    vec3_in = float3(
+        vector_mul * 3.0, vector_mul * 3.0 + scalar_mul, vector_mul * 3.0 + scalar_mul * 2.0);
+    vec4_in = float4(vector_mul * 4.0,
+                     vector_mul * 4.0 + scalar_mul,
+                     vector_mul * 4.0 + scalar_mul * 2.0,
+                     vector_mul * 4.0 + scalar_mul * 3.0);
+  }
+
+  void read_back()
+  {
+    GPU_memory_barrier(GPU_BARRIER_SHADER_STORAGE);
+    GPU_storagebuf_read(ssbo, data.data());
+  }
+
+  void validate()
+  {
+    /* 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);
+  }
+};
+
+struct Shader {
+  GPUShader *shader = nullptr;
+  Vector<CallData> call_datas;
+
+  ~Shader()
+  {
+    if (shader != nullptr) {
+      GPU_shader_unbind();
+      GPU_shader_free(shader);
+    }
+  }
+
+  void init_shader(const char *info_name)
+  {
+    if (shader == nullptr) {
+      shader = GPU_shader_create_from_info_name(info_name);
+      EXPECT_NE(shader, nullptr);
+      GPU_shader_bind(shader);
+    }
+  }
+
+  CallData &new_call()
+  {
+    CallData call_data;
+    call_datas.append(call_data);
+    return call_datas.last();
+  }
+
+  void bind(CallData &call_data)
+  {
+    GPU_storagebuf_bind(call_data.ssbo, GPU_shader_get_ssbo_binding(shader, "data_out"));
+  }
+
+  void update_push_constants(const CallData &call_data)
+  {
+    GPU_shader_uniform_1f(shader, "float_in", call_data.float_in);
+    GPU_shader_uniform_2fv(shader, "vec2_in", call_data.vec2_in);
+    GPU_shader_uniform_3fv(shader, "vec3_in", call_data.vec3_in);
+    GPU_shader_uniform_4fv(shader, "vec4_in", call_data.vec4_in);
+  }
+
+  void dispatch()
+  {
+    GPU_compute_dispatch(shader, 1, 1, 1);
+  }
+};
+
+/** Test the given info when doing a single call. */
+static void do_push_constants_test(const char *info_name, const int num_calls_simultaneously = 1)
+{
+  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;
+
+  Shader shader;
+  shader.init_shader(info_name);
+
+  for (const int call_index : IndexRange(num_calls_simultaneously)) {
+    CallData &call_data = shader.new_call();
+    call_data.generate_test_data(call_index * 10.0, call_index * 1.0);
+    call_data.init_ssbo(SIZE);
+    shader.bind(call_data);
+    shader.update_push_constants(call_data);
+    shader.dispatch();
+  }
+  /* All calls will be "simultaneously" in flight. First readback will wait until the dispatches
+   * have finished execution.*/
+  for (const int call_index : IndexRange(num_calls_simultaneously)) {
+    CallData &call_data = shader.call_datas[call_index];
+    call_data.read_back();
+    call_data.validate();
+  }
+}
+
+/* Test case with single call as sanity check, before we make it more interesting.*/
+static void test_push_constants()
+{
+  do_push_constants_test("gpu_push_constants_test");
+}
+GPU_TEST(push_constants)
+
+static void test_push_constants_128bytes()
+{
+  do_push_constants_test("gpu_push_constants_128bytes_test");
+}
+GPU_TEST(push_constants_128bytes)
+
+static void test_push_constants_256bytes()
+{
+  do_push_constants_test("gpu_push_constants_256bytes_test");
+}
+GPU_TEST(push_constants_256bytes)
+
+static void test_push_constants_512bytes()
+{
+  do_push_constants_test("gpu_push_constants_512bytes_test");
+}
+GPU_TEST(push_constants_512bytes)
+
+#if 0
+/* Schedule multiple simultaneously. */
+/* These test have been disabled for now as this will to be solved in a separate PR.
+ * - DescriptorSets may not be altered, when they are in the command queue or being executed.
+ */
+static void test_push_constants_multiple()
+{
+  do_push_constants_test("gpu_push_constants_test", 10);
+}
+GPU_TEST(push_constants_multiple)
+
+static void test_push_constants_multiple_128bytes()
+{
+  do_push_constants_test("gpu_push_constants_128bytes_test", 10);
+}
+GPU_TEST(push_constants_multiple_128bytes)
+
+static void test_push_constants_multiple_256bytes()
+{
+  do_push_constants_test("gpu_push_constants_256bytes_test", 10);
+}
+GPU_TEST(push_constants_multiple_256bytes)
+
+static void test_push_constants_multiple_512bytes()
+{
+  do_push_constants_test("gpu_push_constants_512bytes_test", 10);
+}
+GPU_TEST(push_constants_multiple_512bytes)
+#endif
+
+}  // namespace blender::gpu::tests

source/blender/gpu/tests/memory_layout_test.cc

@@ -0,0 +1,100 @@
+/* SPDX-License-Identifier: Apache-2.0 */
+
+#include "testing/testing.h"
+
+#include "../vulkan/vk_memory_layout.hh"
+
+namespace blender::gpu {
+
+template<typename Layout>
+static void def_attr(const shader::Type type,
+                     const int array_size,
+                     const uint32_t expected_alignment,
+                     const uint32_t expected_reserve,
+                     uint32_t *r_offset)
+{
+  align<Layout>(type, array_size, r_offset);
+  EXPECT_EQ(*r_offset, expected_alignment);
+  reserve<Layout>(type, array_size, r_offset);
+  EXPECT_EQ(*r_offset, expected_reserve);
+}
+
+TEST(std140, fl)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 16);
+}
+
+TEST(std140, _2fl)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 4, 8, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 16);
+}
+
+TEST(std140, _3fl)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 4, 8, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 8, 12, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 16);
+}
+
+TEST(std140, _4fl)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 4, 8, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 8, 12, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 0, 12, 16, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 16);
+}
+
+TEST(std140, fl2)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 2, 0, 32, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 32);
+}
+
+TEST(std140, fl_fl2)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+  def_attr<Std140>(shader::Type::FLOAT, 2, 16, 48, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 48);
+}
+
+TEST(std140, fl_vec2)
+{
+  uint32_t offset = 0;
+
+  def_attr<Std140>(shader::Type::FLOAT, 0, 0, 4, &offset);
+  def_attr<Std140>(shader::Type::VEC2, 0, 8, 16, &offset);
+
+  align_end_of_struct<Std140>(&offset);
+  EXPECT_EQ(offset, 16);
+}
+
+}  // namespace blender::gpu

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

@@ -67,6 +67,9 @@ void VKBackend::compute_dispatch(int groups_x_len, int groups_y_len, int groups_
   VKCommandBuffer &command_buffer = context.command_buffer_get();
   VKPipeline &pipeline = shader->pipeline_get();
   VKDescriptorSet &descriptor_set = pipeline.descriptor_set_get();
+  VKPushConstants &push_constants = pipeline.push_constants_get();
+
+  push_constants.update(context);
   descriptor_set.update(context.device_get());
   command_buffer.bind(
       descriptor_set, shader->vk_pipeline_layout_get(), VK_PIPELINE_BIND_POINT_COMPUTE);

source/blender/gpu/vulkan/vk_command_buffer.cc

@@ -9,6 +9,7 @@
 #include "vk_buffer.hh"
 #include "vk_context.hh"
 #include "vk_memory.hh"
+#include "vk_pipeline.hh"
 #include "vk_texture.hh"
 
 #include "BLI_assert.h"
@@ -71,6 +72,20 @@ 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)
+{
+  BLI_assert(push_constants.layout_get().storage_type_get() ==
+             VKPushConstants::StorageType::PUSH_CONSTANTS);
+  vkCmdPushConstants(vk_command_buffer_,
+                     vk_pipeline_layout,
+                     vk_shader_stages,
+                     push_constants.offset(),
+                     push_constants.layout_get().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

@@ -8,11 +8,15 @@
 #pragma once
 
 #include "vk_common.hh"
-#include "vk_pipeline.hh"
+
+#include "BLI_utility_mixins.hh"
 
 namespace blender::gpu {
 class VKBuffer;
 class VKTexture;
+class VKPushConstants;
+class VKPipeline;
+class VKDescriptorSet;
 
 /** Command buffer to keep track of the life-time of a command buffer. */
 class VKCommandBuffer : NonCopyable, NonMovable {
@@ -33,6 +37,14 @@ class VKCommandBuffer : NonCopyable, NonMovable {
   void bind(const VKDescriptorSet &descriptor_set,
             const VkPipelineLayout vk_pipeline_layout,
             VkPipelineBindPoint bind_point);
+  /**
+   * Add a push constant command to the command buffer.
+   *
+   * Only valid when the storage type of push_constants is StorageType::PUSH_CONSTANTS.
+   */
+  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_descriptor_set.cc

@@ -15,6 +15,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

@@ -111,6 +111,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_memory_layout.cc

@@ -0,0 +1,172 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#include "vk_memory_layout.hh"
+
+namespace blender::gpu {
+
+uint32_t Std430::component_mem_size(const shader::Type /*type*/)
+{
+  return 4;
+}
+
+uint32_t Std430::element_alignment(const shader::Type type, const bool is_array)
+{
+  if (is_array) {
+    return 16;
+  }
+  switch (type) {
+    case shader::Type::FLOAT:
+    case shader::Type::UINT:
+    case shader::Type::INT:
+    case shader::Type::BOOL:
+      return 4;
+    case shader::Type::VEC2:
+    case shader::Type::UVEC2:
+    case shader::Type::IVEC2:
+      return 8;
+    case shader::Type::VEC3:
+    case shader::Type::UVEC3:
+    case shader::Type::IVEC3:
+    case shader::Type::VEC4:
+    case shader::Type::UVEC4:
+    case shader::Type::IVEC4:
+    case shader::Type::MAT3:
+    case shader::Type::MAT4:
+      return 16;
+    default:
+      BLI_assert_msg(false, "Type not supported in dynamic structs.");
+  }
+  return 0;
+}
+
+uint32_t Std430::element_components_len(const shader::Type type)
+{
+  switch (type) {
+    case shader::Type::FLOAT:
+    case shader::Type::UINT:
+    case shader::Type::INT:
+    case shader::Type::BOOL:
+      return 1;
+    case shader::Type::VEC2:
+    case shader::Type::UVEC2:
+    case shader::Type::IVEC2:
+      return 2;
+    case shader::Type::VEC3:
+    case shader::Type::UVEC3:
+    case shader::Type::IVEC3:
+    case shader::Type::VEC4:
+    case shader::Type::UVEC4:
+    case shader::Type::IVEC4:
+      return 4;
+    case shader::Type::MAT3:
+      return 12;
+    case shader::Type::MAT4:
+      return 16;
+    default:
+      BLI_assert_msg(false, "Type not supported in dynamic structs.");
+  }
+  return 0;
+}
+
+uint32_t Std430::array_components_len(const shader::Type type)
+{
+  return Std430::element_components_len(type);
+}
+
+uint32_t Std140::component_mem_size(const shader::Type /*type*/)
+{
+  return 4;
+}
+
+uint32_t Std140::element_alignment(const shader::Type type, const bool is_array)
+{
+  if (is_array) {
+    return 16;
+  }
+  switch (type) {
+    case shader::Type::FLOAT:
+    case shader::Type::UINT:
+    case shader::Type::INT:
+    case shader::Type::BOOL:
+      return 4;
+    case shader::Type::VEC2:
+    case shader::Type::UVEC2:
+    case shader::Type::IVEC2:
+      return 8;
+    case shader::Type::VEC3:
+    case shader::Type::UVEC3:
+    case shader::Type::IVEC3:
+    case shader::Type::VEC4:
+    case shader::Type::UVEC4:
+    case shader::Type::IVEC4:
+    case shader::Type::MAT3:
+    case shader::Type::MAT4:
+      return 16;
+    default:
+      BLI_assert_msg(false, "Type not supported in dynamic structs.");
+  }
+  return 0;
+}
+
+uint32_t Std140::element_components_len(const shader::Type type)
+{
+  switch (type) {
+    case shader::Type::FLOAT:
+    case shader::Type::UINT:
+    case shader::Type::INT:
+    case shader::Type::BOOL:
+      return 1;
+    case shader::Type::VEC2:
+    case shader::Type::UVEC2:
+    case shader::Type::IVEC2:
+      return 2;
+    case shader::Type::VEC3:
+    case shader::Type::UVEC3:
+    case shader::Type::IVEC3:
+    case shader::Type::VEC4:
+    case shader::Type::UVEC4:
+    case shader::Type::IVEC4:
+      return 4;
+    case shader::Type::MAT3:
+      return 12;
+    case shader::Type::MAT4:
+      return 16;
+    default:
+      BLI_assert_msg(false, "Type not supported in dynamic structs.");
+  }
+  return 0;
+}
+
+uint32_t Std140::array_components_len(const shader::Type type)
+{
+  switch (type) {
+    case shader::Type::FLOAT:
+    case shader::Type::UINT:
+    case shader::Type::INT:
+    case shader::Type::BOOL:
+    case shader::Type::VEC2:
+    case shader::Type::UVEC2:
+    case shader::Type::IVEC2:
+    case shader::Type::VEC3:
+    case shader::Type::UVEC3:
+    case shader::Type::IVEC3:
+    case shader::Type::VEC4:
+    case shader::Type::UVEC4:
+    case shader::Type::IVEC4:
+      return 4;
+    case shader::Type::MAT3:
+      return 12;
+    case shader::Type::MAT4:
+      return 16;
+    default:
+      BLI_assert_msg(false, "Type not supported in dynamic structs.");
+  }
+  return 0;
+}
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_memory_layout.hh

@@ -0,0 +1,105 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#pragma once
+
+#include "gpu_shader_create_info.hh"
+
+namespace blender::gpu {
+
+/**
+ * Information about alignment/components and memory size for types when using std140 layout.
+ */
+struct Std140 {
+  /** Get the memory size in bytes of a single component using by the given type.*/
+  static uint32_t component_mem_size(const shader::Type type);
+  /** Get to alignment of the given type in bytes.*/
+  static uint32_t element_alignment(const shader::Type type, bool is_array);
+  /** Get the number of components that should be allocated for the given type.*/
+  static uint32_t element_components_len(const shader::Type type);
+  /** Get the number of components of the given type when used in an array.*/
+  static uint32_t array_components_len(const shader::Type type);
+};
+
+/**
+ * Information about alignment/components and memory size for types when using std430 layout.
+ */
+struct Std430 {
+  /** Get the memory size in bytes of a single component using by the given type.*/
+  static uint32_t component_mem_size(const shader::Type type);
+  /** Get to alignment of the given type in bytes.*/
+  static uint32_t element_alignment(const shader::Type type, bool is_array);
+  /** Get the number of components that should be allocated for the given type.*/
+  static uint32_t element_components_len(const shader::Type type);
+  /** Get the number of components of the given type when used in an array.*/
+  static uint32_t array_components_len(const shader::Type type);
+};
+
+template<typename LayoutT> static uint32_t element_stride(const shader::Type type)
+{
+  return LayoutT::element_components_len(type) * LayoutT::component_mem_size(type);
+}
+
+template<typename LayoutT> static uint32_t array_stride(const shader::Type type)
+{
+  return LayoutT::array_components_len(type) * LayoutT::component_mem_size(type);
+}
+
+/**
+ * Move the r_offset to the next alignment where the given type+array_size can be
+ * reserved.
+ *
+ * 'type': The type that needs to be aligned.
+ * 'array_size': The array_size that needs to be aligned. (0=no array).
+ * 'r_offset': After the call it will point to the byte where the reservation
+ *     can happen.
+ */
+template<typename LayoutT>
+static void align(const shader::Type &type, const int32_t array_size, uint32_t *r_offset)
+{
+  uint32_t alignment = LayoutT::element_alignment(type, array_size != 0);
+  uint32_t alignment_mask = alignment - 1;
+  uint32_t offset = *r_offset;
+  if ((offset & alignment_mask) != 0) {
+    offset &= ~alignment_mask;
+    offset += alignment;
+    *r_offset = offset;
+  }
+}
+
+/**
+ * Reserve space for the given type and array size.
+ *
+ * This function doesn't handle alignment this needs to be done up front by calling
+ * 'align<Layout>' function. Caller is responsible for this.
+ *
+ * 'type': The type that needs to be reserved.
+ * 'array_size': The array_size that needs to be reserved. (0=no array).
+ * 'r_offset': When calling needs to be pointing to the aligned location where to
+ *     reserve space. After the call it will point to the byte just after reserved
+ *     space.
+ */
+template<typename LayoutT>
+static void reserve(const shader::Type type, int32_t array_size, uint32_t *r_offset)
+{
+  uint32_t size = array_size == 0 ? element_stride<LayoutT>(type) :
+                                    array_stride<LayoutT>(type) * array_size;
+  *r_offset += size;
+}
+
+/**
+ * Update 'r_offset' to be aligned to the end of the struct.
+ *
+ * Call this function when all attributes have been added to make sure that the struct size is
+ * correct.
+ */
+template<typename LayoutT> static void align_end_of_struct(uint32_t *r_offset)
+{
+  align<LayoutT>(shader::Type::VEC4, 0, r_offset);
+}
+
+}  // namespace blender::gpu

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()
@@ -26,10 +29,12 @@ VKPipeline::~VKPipeline()
   }
 }
 
-VKPipeline VKPipeline::create_compute_pipeline(VKContext &context,
-                                               VkShaderModule compute_module,
-                                               VkDescriptorSetLayout &descriptor_set_layout,
-                                               VkPipelineLayout &pipeline_layout)
+VKPipeline VKPipeline::create_compute_pipeline(
+    VKContext &context,
+    VkShaderModule compute_module,
+    VkDescriptorSetLayout &descriptor_set_layout,
+    VkPipelineLayout &pipeline_layout,
+    const VKPushConstants::Layout &push_constants_layout)
 {
   VK_ALLOCATION_CALLBACKS
   VkDevice vk_device = context.device_get();
@@ -44,15 +49,16 @@ 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,
+                                            const VKPushConstants::Layout &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

@@ -0,0 +1,190 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#include "vk_push_constants.hh"
+#include "vk_backend.hh"
+#include "vk_memory_layout.hh"
+#include "vk_shader.hh"
+#include "vk_shader_interface.hh"
+#include "vk_storage_buffer.hh"
+#include "vk_uniform_buffer.hh"
+
+namespace blender::gpu {
+
+template<typename LayoutT>
+static VKPushConstants::Layout::PushConstant init_constant(
+    const shader::ShaderCreateInfo::PushConst &push_constant,
+    const ShaderInput &shader_input,
+    uint32_t *r_offset)
+{
+  align<LayoutT>(push_constant.type, push_constant.array_size, r_offset);
+
+  VKPushConstants::Layout::PushConstant layout;
+  layout.location = shader_input.location;
+  layout.type = push_constant.type;
+  layout.array_size = push_constant.array_size;
+  layout.offset = *r_offset;
+
+  reserve<LayoutT>(push_constant.type, push_constant.array_size, r_offset);
+  return layout;
+}
+
+template<typename LayoutT>
+uint32_t struct_size(Span<shader::ShaderCreateInfo::PushConst> push_constants)
+{
+  uint32_t offset = 0;
+  for (const shader::ShaderCreateInfo::PushConst &push_constant : push_constants) {
+    align<LayoutT>(push_constant.type, push_constant.array_size, &offset);
+    reserve<LayoutT>(push_constant.type, push_constant.array_size, &offset);
+  }
+
+  align_end_of_struct<LayoutT>(&offset);
+  return offset;
+}
+
+VKPushConstants::StorageType VKPushConstants::Layout::determine_storage_type(
+    const shader::ShaderCreateInfo &info, const VkPhysicalDeviceLimits &vk_physical_device_limits)
+{
+  if (info.push_constants_.is_empty()) {
+    return StorageType::NONE;
+  }
+
+  uint32_t size = struct_size<Std430>(info.push_constants_);
+  return size <= vk_physical_device_limits.maxPushConstantsSize ? STORAGE_TYPE_DEFAULT :
+                                                                  STORAGE_TYPE_FALLBACK;
+}
+
+template<typename LayoutT>
+void init_struct(const shader::ShaderCreateInfo &info,
+                 const VKShaderInterface &interface,
+                 Vector<VKPushConstants::Layout::PushConstant> &r_struct,
+                 uint32_t *r_offset)
+{
+  for (const shader::ShaderCreateInfo::PushConst &push_constant : info.push_constants_) {
+    const ShaderInput *shader_input = interface.uniform_get(push_constant.name.c_str());
+    r_struct.append(init_constant<LayoutT>(push_constant, *shader_input, r_offset));
+  }
+  align_end_of_struct<Std140>(r_offset);
+}
+
+void VKPushConstants::Layout::init(const shader::ShaderCreateInfo &info,
+                                   const VKShaderInterface &interface,
+                                   const StorageType storage_type,
+                                   const VKDescriptorSet::Location location)
+{
+  BLI_assert(push_constants.is_empty());
+  storage_type_ = storage_type;
+
+  size_in_bytes_ = 0;
+  if (storage_type == StorageType::UNIFORM_BUFFER) {
+    descriptor_set_location_ = location;
+    init_struct<Std140>(info, interface, push_constants, &size_in_bytes_);
+  }
+  else {
+    init_struct<Std430>(info, interface, push_constants, &size_in_bytes_);
+  }
+}
+
+const VKPushConstants::Layout::PushConstant *VKPushConstants::Layout::find(int32_t location) const
+{
+  for (const PushConstant &push_constant : push_constants) {
+    if (push_constant.location == location) {
+      return &push_constant;
+    }
+  }
+  return nullptr;
+}
+
+VKPushConstants::VKPushConstants() = default;
+VKPushConstants::VKPushConstants(const Layout *layout) : layout_(layout)
+{
+  data_ = MEM_mallocN(layout->size_in_bytes(), __func__);
+  switch (layout_->storage_type_get()) {
+    case StorageType::UNIFORM_BUFFER:
+      uniform_buffer_ = new VKUniformBuffer(layout_->size_in_bytes(), __func__);
+      break;
+
+    case StorageType::PUSH_CONSTANTS:
+    case StorageType::NONE:
+      break;
+  }
+}
+
+VKPushConstants::VKPushConstants(VKPushConstants &&other) : layout_(other.layout_)
+{
+  data_ = other.data_;
+  other.data_ = nullptr;
+
+  uniform_buffer_ = other.uniform_buffer_;
+  other.uniform_buffer_ = nullptr;
+}
+
+VKPushConstants::~VKPushConstants()
+{
+  if (data_ != nullptr) {
+    MEM_freeN(data_);
+    data_ = nullptr;
+  }
+
+  delete uniform_buffer_;
+  uniform_buffer_ = nullptr;
+}
+
+VKPushConstants &VKPushConstants::operator=(VKPushConstants &&other)
+{
+  layout_ = other.layout_;
+
+  data_ = other.data_;
+  other.data_ = nullptr;
+
+  uniform_buffer_ = other.uniform_buffer_;
+  other.uniform_buffer_ = nullptr;
+
+  return *this;
+}
+
+void VKPushConstants::update(VKContext &context)
+{
+  VKShader *shader = static_cast<VKShader *>(context.shader);
+  VKCommandBuffer &command_buffer = context.command_buffer_get();
+  VKPipeline &pipeline = shader->pipeline_get();
+  BLI_assert_msg(&pipeline.push_constants_get() == this,
+                 "Invalid state detected. Push constants doesn't belong to the active shader of "
+                 "the given context.");
+  VKDescriptorSet &descriptor_set = pipeline.descriptor_set_get();
+
+  switch (layout_get().storage_type_get()) {
+    case VKPushConstants::StorageType::NONE:
+      break;
+
+    case VKPushConstants::StorageType::PUSH_CONSTANTS:
+      command_buffer.push_constants(*this, shader->vk_pipeline_layout_get(), VK_SHADER_STAGE_ALL);
+      break;
+
+    case VKPushConstants::StorageType::UNIFORM_BUFFER:
+      update_uniform_buffer();
+      descriptor_set.bind(uniform_buffer_get(), layout_get().descriptor_set_location_get());
+      break;
+  }
+}
+
+void VKPushConstants::update_uniform_buffer()
+{
+  BLI_assert(layout_->storage_type_get() == StorageType::UNIFORM_BUFFER);
+  BLI_assert(uniform_buffer_ != nullptr);
+  BLI_assert(data_ != nullptr);
+  uniform_buffer_->update(data_);
+}
+
+VKUniformBuffer &VKPushConstants::uniform_buffer_get()
+{
+  BLI_assert(layout_->storage_type_get() == StorageType::UNIFORM_BUFFER);
+  BLI_assert(uniform_buffer_ != nullptr);
+  return *uniform_buffer_;
+}
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_push_constants.hh

@@ -0,0 +1,249 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ *
+ * Push constants is a way to quickly provide a small amount of uniform data to shaders. It should
+ * be much quicker than UBOs but a huge limitation is the size of data - spec requires 128 bytes to
+ * be available for a push constant range. Hardware vendors may support more, but compared to other
+ * means it is still very little (for example 256 bytes).
+ *
+ * Due to this size requirements we try to use push constants when it fits on the device. If it
+ * doesn't fit we fallback to use an uniform buffer.
+ *
+ * Shader developers are responsible to fine-tune the performance of the shader. One way to do this
+ * is to tailor what will be sent as a push constant to keep the push constants within the limits.
+ */
+
+#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;
+
+/**
+ * Container to store push constants in a buffer.
+ *
+ * Can handle buffers with different memory layouts (std140/std430)
+ * Which memory layout is used is based on the storage type.
+ *
+ * VKPushConstantsLayout only describes the buffer, an instance of this
+ * class can handle setting/modifying/duplicating push constants.
+ *
+ * It should also keep track of the submissions in order to reuse the allocated
+ * data.
+ */
+class VKPushConstants : NonCopyable {
+ public:
+  /** Different methods to store push constants.*/
+  enum class StorageType {
+    /** Push constants aren't in use.*/
+    NONE,
+
+    /** Store push constants as regular vulkan push constants.*/
+    PUSH_CONSTANTS,
+
+    /**
+     * Fallback when push constants doesn't meet the device requirements.
+     */
+    UNIFORM_BUFFER,
+  };
+
+  /**
+   * Describe the layout of the push constants and the storage type that should be used.
+   */
+  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;
+
+      /** Offset in the push constant data (in bytes). */
+      uint32_t offset;
+      shader::Type type;
+      int array_size;
+    };
+
+   private:
+    Vector<PushConstant> push_constants;
+    uint32_t size_in_bytes_ = 0;
+    StorageType storage_type_ = StorageType::NONE;
+    /**
+     * Binding index in the descriptor set when the push constants use an uniform buffer.
+     */
+    VKDescriptorSet::Location descriptor_set_location_;
+
+   public:
+    /**
+     * Return the desired storage type that can fit the push constants of the given shader create
+     * info, matching the device limits.
+     *
+     * Returns:
+     * - StorageType::NONE: No push constants are needed.
+     * - StorageType::PUSH_CONSTANTS: Regular vulkan push constants can be used.
+     * - StorageType::UNIFORM_BUFFER: The push constants don't fit in the limits of the given
+     * device. A uniform buffer should be used as a fallback method.
+     */
+    static StorageType determine_storage_type(
+        const shader::ShaderCreateInfo &info,
+        const VkPhysicalDeviceLimits &vk_physical_device_limits);
+
+    /**
+     * Initialize the push constants of the given shader create info with the
+     * binding location.
+     *
+     * interface: Uniform locations of the interface are used as lookup key.
+     * storage_type: The type of storage for push constants to use.
+     * location: When storage_type=StorageType::UNIFORM_BUFFER this contains
+     *    the location in the descriptor set where the uniform buffer can be
+     *    bound.
+     */
+    void init(const shader::ShaderCreateInfo &info,
+              const VKShaderInterface &interface,
+              StorageType storage_type,
+              VKDescriptorSet::Location location);
+
+    /**
+     * Return the storage type that is used.
+     */
+    StorageType storage_type_get() const
+    {
+      return storage_type_;
+    }
+
+    /**
+     * Get the binding location for the uniform buffer.
+     *
+     * Only valid when storage_type=StorageType::UNIFORM_BUFFER.
+     */
+    VKDescriptorSet::Location descriptor_set_location_get() const
+    {
+      return descriptor_set_location_;
+    }
+
+    /**
+     * Get the size needed to store the push constants.
+     */
+    uint32_t size_in_bytes() const
+    {
+      return size_in_bytes_;
+    }
+
+    /**
+     * Find the push constant layout for the given location.
+     * Location = ShaderInput.location.
+     */
+    const PushConstant *find(int32_t location) const;
+  };
+
+ private:
+  const Layout *layout_ = nullptr;
+  void *data_ = nullptr;
+  VKUniformBuffer *uniform_buffer_ = nullptr;
+
+ 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_;
+  }
+
+  /**
+   * Get the reference to the active data.
+   *
+   * Data can get inactive when push constants are modified, after being added to the command
+   * queue. We still keep track of the old data for reuse and make sure we don't overwrite data
+   * that is still not on the GPU.
+   */
+  const void *data() const
+  {
+    return data_;
+  }
+
+  /**
+   * Modify a push constant.
+   *
+   * location: ShaderInput.location of the push constant to update.
+   * comp_len: number of components has the data type that is being updated.
+   * array_size: number of elements when an array to update. (0=no array)
+   * input_data: packed source data to use.
+   */
+  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);
+    BLI_assert(push_constant_layout);
+
+    uint8_t *bytes = static_cast<uint8_t *>(data_);
+    T *dst = static_cast<T *>(static_cast<void *>(&bytes[push_constant_layout->offset]));
+    const bool is_tightly_std140_packed = (comp_len % 4) == 0;
+    if (layout_->storage_type_get() == StorageType::PUSH_CONSTANTS || array_size == 0 ||
+        is_tightly_std140_packed) {
+      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.");
+      memcpy(dst, input_data, comp_len * array_size * sizeof(T));
+      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;
+    for (const int i : IndexRange(array_size)) {
+      UNUSED_VARS(i);
+      memcpy(dst, src, comp_len * sizeof(T));
+      src += comp_len;
+      dst += 4;
+    }
+  }
+
+  /**
+   * Update the GPU resources with the latest push constants.
+   */
+  void update(VKContext &context);
+
+ private:
+  /**
+   * When storage type = StorageType::UNIFORM_BUFFER use this method to update the uniform
+   * buffer.
+   *
+   * It must be called just before adding a draw/compute command to the command queue.
+   */
+  void update_uniform_buffer();
+
+  /**
+   * Get a reference to the uniform buffer.
+   *
+   * Only valid when storage type = StorageType::UNIFORM_BUFFER.
+   */
+  VKUniformBuffer &uniform_buffer_get();
+};
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_shader.cc

@@ -676,7 +676,7 @@ bool VKShader::finalize(const shader::ShaderCreateInfo *info)
   if (!finalize_descriptor_set_layouts(vk_device, *vk_interface, *info)) {
     return false;
   }
-  if (!finalize_pipeline_layout(vk_device, *info)) {
+  if (!finalize_pipeline_layout(vk_device, *vk_interface)) {
     return false;
   }
 
@@ -695,7 +695,11 @@ 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_,
+        vk_interface->push_constants_layout_get());
     result = compute_pipeline_.is_valid();
   }
 
@@ -739,17 +743,29 @@ bool VKShader::finalize_graphics_pipeline(VkDevice /*vk_device */)
 }
 
 bool VKShader::finalize_pipeline_layout(VkDevice vk_device,
-                                        const shader::ShaderCreateInfo & /*info*/)
+                                        const VKShaderInterface &shader_interface)
 {
   VK_ALLOCATION_CALLBACKS
 
   const uint32_t layout_count = layout_ == VK_NULL_HANDLE ? 0 : 1;
   VkPipelineLayoutCreateInfo pipeline_info = {};
+  VkPushConstantRange push_constant_range = {};
   pipeline_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
   pipeline_info.flags = 0;
   pipeline_info.setLayoutCount = layout_count;
   pipeline_info.pSetLayouts = &layout_;
 
+  /* Setup push constants. */
+  const VKPushConstants::Layout &push_constants_layout =
+      shader_interface.push_constants_layout_get();
+  if (push_constants_layout.storage_type_get() == VKPushConstants::StorageType::PUSH_CONSTANTS) {
+    push_constant_range.offset = 0;
+    push_constant_range.size = push_constants_layout.size_in_bytes();
+    push_constant_range.stageFlags = VK_SHADER_STAGE_ALL;
+    pipeline_info.pushConstantRangeCount = 1;
+    pipeline_info.pPushConstantRanges = &push_constant_range;
+  }
+
   if (vkCreatePipelineLayout(
           vk_device, &pipeline_info, vk_allocation_callbacks, &pipeline_layout_) != VK_SUCCESS) {
     return false;
@@ -868,6 +884,21 @@ static VkDescriptorSetLayoutBinding create_descriptor_set_layout_binding(
   return binding;
 }
 
+static VkDescriptorSetLayoutBinding create_descriptor_set_layout_binding(
+    const VKPushConstants::Layout &push_constants_layout)
+{
+  BLI_assert(push_constants_layout.storage_type_get() ==
+             VKPushConstants::StorageType::UNIFORM_BUFFER);
+  VkDescriptorSetLayoutBinding binding = {};
+  binding.binding = push_constants_layout.descriptor_set_location_get();
+  binding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
+  binding.descriptorCount = 1;
+  binding.stageFlags = VK_SHADER_STAGE_ALL;
+  binding.pImmutableSamplers = nullptr;
+
+  return binding;
+}
+
 static void add_descriptor_set_layout_bindings(
     const VKShaderInterface &interface,
     const Vector<shader::ShaderCreateInfo::Resource> &resources,
@@ -877,6 +908,12 @@ static void add_descriptor_set_layout_bindings(
     const VKDescriptorSet::Location location = interface.descriptor_set_location(resource);
     r_bindings.append(create_descriptor_set_layout_binding(location, resource));
   }
+
+  /* Add push constants to the descriptor when push constants are stored in an uniform buffer.*/
+  const VKPushConstants::Layout &push_constants_layout = interface.push_constants_layout_get();
+  if (push_constants_layout.storage_type_get() == VKPushConstants::StorageType::UNIFORM_BUFFER) {
+    r_bindings.append(create_descriptor_set_layout_binding(push_constants_layout));
+  }
 }
 
 static VkDescriptorSetLayoutCreateInfo create_descriptor_set_layout(
@@ -894,11 +931,19 @@ static VkDescriptorSetLayoutCreateInfo create_descriptor_set_layout(
   return set_info;
 }
 
+static bool descriptor_sets_needed(const VKShaderInterface &shader_interface,
+                                   const shader::ShaderCreateInfo &info)
+{
+  return !info.pass_resources_.is_empty() || !info.batch_resources_.is_empty() ||
+         shader_interface.push_constants_layout_get().storage_type_get() ==
+             VKPushConstants::StorageType::UNIFORM_BUFFER;
+}
+
 bool VKShader::finalize_descriptor_set_layouts(VkDevice vk_device,
                                                const VKShaderInterface &shader_interface,
                                                const shader::ShaderCreateInfo &info)
 {
-  if (info.pass_resources_.is_empty() && info.batch_resources_.is_empty()) {
+  if (!descriptor_sets_needed(shader_interface, info)) {
     return true;
   }
 
@@ -958,17 +1003,14 @@ 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*/,
-                           const int * /*data*/)
+
+void VKShader::uniform_int(int location, int comp_len, int array_size, const int *data)
 {
+  pipeline_get().push_constants_get().push_constant_set(location, comp_len, array_size, data);
 }
 
 std::string VKShader::resources_declare(const shader::ShaderCreateInfo &info) const
@@ -993,9 +1035,19 @@ std::string VKShader::resources_declare(const shader::ShaderCreateInfo &info) co
     print_resource_alias(ss, res);
   }
 
-  if (!info.push_constants_.is_empty()) {
+  /* Push constants. */
+  const VKPushConstants::Layout &push_constants_layout = interface.push_constants_layout_get();
+  const VKPushConstants::StorageType push_constants_storage =
+      push_constants_layout.storage_type_get();
+  if (push_constants_storage != VKPushConstants::StorageType::NONE) {
     ss << "\n/* Push Constants. */\n";
-    ss << "layout(push_constant) uniform constants\n";
+    if (push_constants_storage == VKPushConstants::StorageType::PUSH_CONSTANTS) {
+      ss << "layout(push_constant) uniform constants\n";
+    }
+    else if (push_constants_storage == VKPushConstants::StorageType::UNIFORM_BUFFER) {
+      ss << "layout(binding = " << push_constants_layout.descriptor_set_location_get()
+         << ", std140) uniform constants\n";
+    }
     ss << "{\n";
     for (const ShaderCreateInfo::PushConst &uniform : info.push_constants_) {
       ss << "  " << to_string(uniform.type) << " pc_" << uniform.name;

source/blender/gpu/vulkan/vk_shader.hh

@@ -11,6 +11,7 @@
 
 #include "vk_backend.hh"
 #include "vk_context.hh"
+#include "vk_pipeline.hh"
 
 #include "BLI_string_ref.hh"
 
@@ -78,7 +79,7 @@ class VKShader : public Shader {
   bool finalize_descriptor_set_layouts(VkDevice vk_device,
                                        const VKShaderInterface &shader_interface,
                                        const shader::ShaderCreateInfo &info);
-  bool finalize_pipeline_layout(VkDevice vk_device, const shader::ShaderCreateInfo &info);
+  bool finalize_pipeline_layout(VkDevice vk_device, const VKShaderInterface &shader_interface);
   bool finalize_graphics_pipeline(VkDevice vk_device);
 
   bool is_graphics_shader() const

source/blender/gpu/vulkan/vk_shader_interface.cc

@@ -6,15 +6,19 @@
  */
 
 #include "vk_shader_interface.hh"
+#include "vk_context.hh"
 
 namespace blender::gpu {
 
 void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
 {
+  static char PUSH_CONSTANTS_FALLBACK_NAME[] = "push_constants_fallback";
+  static size_t PUSH_CONSTANTS_FALLBACK_NAME_LEN = strlen(PUSH_CONSTANTS_FALLBACK_NAME);
+
   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;
@@ -40,6 +44,17 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
         break;
     }
   }
+
+  /* Reserve 1 uniform buffer for push constants fallback. */
+  size_t names_size = info.interface_names_size_;
+  VKContext &context = *VKContext::get();
+  const VKPushConstants::StorageType push_constants_storage_type =
+      VKPushConstants::Layout::determine_storage_type(info, context.physical_device_limits_get());
+  if (push_constants_storage_type == VKPushConstants::StorageType::UNIFORM_BUFFER) {
+    ubo_len_++;
+    names_size += PUSH_CONSTANTS_FALLBACK_NAME_LEN + 1;
+  }
+
   /* Make sure that the image slots don't overlap with the sampler slots. */
   image_offset_++;
 
@@ -48,7 +63,7 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
       MEM_calloc_arrayN(input_tot_len, sizeof(ShaderInput), __func__));
   ShaderInput *input = inputs_;
 
-  name_buffer_ = (char *)MEM_mallocN(info.interface_names_size_, "name_buffer");
+  name_buffer_ = (char *)MEM_mallocN(names_size, "name_buffer");
   uint32_t name_buffer_offset = 0;
 
   /* Uniform blocks */
@@ -59,6 +74,13 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
       input++;
     }
   }
+  /* Add push constant when using uniform buffer as fallback. */
+  int32_t push_constants_fallback_location = -1;
+  if (push_constants_storage_type == VKPushConstants::StorageType::UNIFORM_BUFFER) {
+    copy_input_name(input, PUSH_CONSTANTS_FALLBACK_NAME, name_buffer_, name_buffer_offset);
+    input->location = input->binding = -1;
+    input++;
+  }
 
   /* Images, Samplers and buffers. */
   for (const ShaderCreateInfo::Resource &res : all_resources) {
@@ -74,6 +96,15 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
     }
   }
 
+  /* Push constants. */
+  int32_t push_constant_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) {
@@ -106,6 +137,17 @@ void VKShaderInterface::init(const shader::ShaderCreateInfo &info)
     const ShaderInput *input = shader_input_get(res);
     descriptor_set_location_update(input, descriptor_set_location++);
   }
+
+  /* Post initializing push constants.*/
+  /* Determine the binding location of push constants fallback buffer.*/
+  int32_t push_constant_descriptor_set_location = -1;
+  if (push_constants_storage_type == VKPushConstants::StorageType::UNIFORM_BUFFER) {
+    push_constant_descriptor_set_location = descriptor_set_location++;
+    const ShaderInput *push_constant_input = ubo_get(PUSH_CONSTANTS_FALLBACK_NAME);
+    descriptor_set_location_update(push_constant_input, push_constants_fallback_location);
+  }
+  push_constants_layout_.init(
+      info, *this, push_constants_storage_type, push_constant_descriptor_set_location);
 }
 
 static int32_t shader_input_index(const ShaderInput *shader_inputs,

source/blender/gpu/vulkan/vk_shader_interface.hh

@@ -12,7 +12,9 @@
 #include "gpu_shader_create_info.hh"
 #include "gpu_shader_interface.hh"
 
-#include "vk_descriptor_set.hh"
+#include "BLI_array.hh"
+
+#include "vk_push_constants.hh"
 
 namespace blender::gpu {
 class VKShaderInterface : public ShaderInterface {
@@ -27,6 +29,9 @@ class VKShaderInterface : public ShaderInterface {
   uint32_t image_offset_ = 0;
   Array<VKDescriptorSet::Location> descriptor_set_locations_;
 
+  VKPushConstants::Layout push_constants_layout_;
+
+
  public:
   VKShaderInterface() = default;
 
@@ -37,6 +42,12 @@ class VKShaderInterface : public ShaderInterface {
   const VKDescriptorSet::Location descriptor_set_location(
       const shader::ShaderCreateInfo::Resource::BindType &bind_type, int binding) const;
 
+  /** Get the Layout of the shader.*/
+  const VKPushConstants::Layout &push_constants_layout_get() const
+  {
+    return push_constants_layout_;
+  }
+
  private:
   /**
    * Retrieve the shader input for the given resource.