Vulkan: Indirect Compute

This PR adds support for indirect compute.
Indirect compute is almost the same as regular compute. The
only difference is that the parameters for the compute dispatch
isn't passed as a parameter, but that these parameters are part
of a buffer.

Pull Request: https://projects.blender.org/blender/blender/pulls/108879

source/blender/gpu/CMakeLists.txt

@@ -847,6 +847,7 @@ if(WITH_GTESTS)
       tests/gpu_testing.cc
 
       tests/buffer_texture_test.cc
+      tests/compute_test.cc
       tests/framebuffer_test.cc
       tests/immediate_test.cc
       tests/index_buffer_test.cc

source/blender/gpu/tests/compute_test.cc

@@ -0,0 +1,112 @@
+/* SPDX-FileCopyrightText: 2023 Blender Foundation
+ *
+ * SPDX-License-Identifier: Apache-2.0 */
+
+#include "gpu_testing.hh"
+
+#include "MEM_guardedalloc.h"
+
+#include "BLI_math_vector_types.hh"
+
+#include "GPU_capabilities.h"
+#include "GPU_compute.h"
+#include "GPU_storage_buffer.h"
+#include "GPU_texture.h"
+
+namespace blender::gpu::tests {
+static void test_compute_direct()
+{
+  if (!GPU_compute_shader_support()) {
+    /* We can't test as a the platform does not support compute shaders. */
+    GTEST_SKIP() << "Skipping test: platform not supported";
+    return;
+  }
+
+  static constexpr uint SIZE = 32;
+
+  /* Build compute shader. */
+  GPUShader *shader = GPU_shader_create_from_info_name("gpu_compute_2d_test");
+  EXPECT_NE(shader, nullptr);
+
+  /* Create texture to store result and attach to shader. */
+  GPUTexture *texture = GPU_texture_create_2d(
+      "gpu_shader_compute_2d", SIZE, SIZE, 1, GPU_RGBA32F, GPU_TEXTURE_USAGE_GENERAL, nullptr);
+  EXPECT_NE(texture, nullptr);
+
+  GPU_shader_bind(shader);
+  GPU_texture_image_bind(texture, GPU_shader_get_sampler_binding(shader, "img_output"));
+
+  /* Dispatch compute task. */
+  GPU_compute_dispatch(shader, SIZE, SIZE, 1);
+
+  /* Check if compute has been done. */
+  GPU_memory_barrier(GPU_BARRIER_TEXTURE_UPDATE);
+  float4 *data = static_cast<float4 *>(GPU_texture_read(texture, GPU_DATA_FLOAT, 0));
+  const float4 expected_result(1.0f, 0.5f, 0.2f, 1.0f);
+  EXPECT_NE(data, nullptr);
+  for (int index = 0; index < SIZE * SIZE; index++) {
+    EXPECT_EQ(data[index], expected_result);
+  }
+  MEM_freeN(data);
+
+  /* Cleanup. */
+  GPU_shader_unbind();
+  GPU_texture_unbind(texture);
+  GPU_texture_free(texture);
+  GPU_shader_free(shader);
+}
+GPU_TEST(compute_direct)
+
+static void test_compute_indirect()
+{
+  if (!GPU_compute_shader_support()) {
+    /* We can't test as a the platform does not support compute shaders. */
+    GTEST_SKIP() << "Skipping test: platform not supported";
+    return;
+  }
+
+  static constexpr uint SIZE = 32;
+
+  /* Build compute shader. */
+  GPUShader *shader = GPU_shader_create_from_info_name("gpu_compute_2d_test");
+  EXPECT_NE(shader, nullptr);
+
+  /* Create texture to store result and attach to shader. */
+  GPUTexture *texture = GPU_texture_create_2d(
+      "gpu_shader_compute_2d", SIZE, SIZE, 1, GPU_RGBA32F, GPU_TEXTURE_USAGE_GENERAL, nullptr);
+  EXPECT_NE(texture, nullptr);
+  GPU_texture_clear(texture, GPU_DATA_FLOAT, float4(0.0f));
+
+  GPU_shader_bind(shader);
+  GPU_texture_image_bind(texture, GPU_shader_get_sampler_binding(shader, "img_output"));
+
+  /* Generate compute tasks. */
+  uint4 commands[1] = {
+      {SIZE, SIZE, 1, 0},
+  };
+  GPUStorageBuf *compute_commands = GPU_storagebuf_create_ex(
+      sizeof(commands), &commands, GPU_USAGE_STATIC, __func__);
+
+  /* Dispatch compute task. */
+  GPU_compute_dispatch_indirect(shader, compute_commands);
+
+  /* Check if compute has been done. */
+  GPU_memory_barrier(GPU_BARRIER_TEXTURE_UPDATE);
+  float4 *data = static_cast<float4 *>(GPU_texture_read(texture, GPU_DATA_FLOAT, 0));
+  const float4 expected_result(1.0f, 0.5f, 0.2f, 1.0f);
+  EXPECT_NE(data, nullptr);
+  for (int index = 0; index < SIZE * SIZE; index++) {
+    EXPECT_EQ(data[index], expected_result);
+  }
+  MEM_freeN(data);
+
+  /* Cleanup. */
+  GPU_storagebuf_free(compute_commands);
+  GPU_shader_unbind();
+  GPU_texture_unbind(texture);
+  GPU_texture_free(texture);
+  GPU_shader_free(shader);
+}
+GPU_TEST(compute_indirect);
+
+}  // namespace blender::gpu::tests

source/blender/gpu/vulkan/vk_backend.cc

@@ -108,9 +108,15 @@ void VKBackend::compute_dispatch(int groups_x_len, int groups_y_len, int groups_
   command_buffer.dispatch(groups_x_len, groups_y_len, groups_z_len);
 }
 
-void VKBackend::compute_dispatch_indirect(StorageBuf * /*indirect_buf*/)
+void VKBackend::compute_dispatch_indirect(StorageBuf *indirect_buf)
 {
-  NOT_YET_IMPLEMENTED;
+  BLI_assert(indirect_buf);
+  VKContext &context = *VKContext::get();
+  context.state_manager_get().apply_bindings();
+  context.bind_compute_pipeline();
+  VKStorageBuffer &indirect_buffer = *unwrap(indirect_buf);
+  VKCommandBuffer &command_buffer = context.command_buffer_get();
+  command_buffer.dispatch(indirect_buffer);
 }
 
 Context *VKBackend::context_alloc(void *ghost_window, void *ghost_context)

source/blender/gpu/vulkan/vk_command_buffer.cc

@@ -13,6 +13,7 @@
 #include "vk_index_buffer.hh"
 #include "vk_memory.hh"
 #include "vk_pipeline.hh"
+#include "vk_storage_buffer.hh"
 #include "vk_texture.hh"
 #include "vk_vertex_buffer.hh"
 
@@ -307,6 +308,12 @@ void VKCommandBuffer::dispatch(int groups_x_len, int groups_y_len, int groups_z_
   vkCmdDispatch(vk_command_buffer_, groups_x_len, groups_y_len, groups_z_len);
 }
 
+void VKCommandBuffer::dispatch(VKStorageBuffer &command_buffer)
+{
+  ensure_no_active_framebuffer();
+  vkCmdDispatchIndirect(vk_command_buffer_, command_buffer.vk_handle(), 0);
+}
+
 void VKCommandBuffer::submit()
 {
   ensure_no_active_framebuffer();

source/blender/gpu/vulkan/vk_command_buffer.hh

@@ -21,6 +21,7 @@ class VKFrameBuffer;
 class VKIndexBuffer;
 class VKPipeline;
 class VKPushConstants;
+class VKStorageBuffer;
 class VKTexture;
 class VKVertexBuffer;
 
@@ -160,6 +161,7 @@ class VKCommandBuffer : NonCopyable, NonMovable {
                       const VkPipelineLayout vk_pipeline_layout,
                       const VkShaderStageFlags vk_shader_stages);
   void dispatch(int groups_x_len, int groups_y_len, int groups_z_len);
+  void dispatch(VKStorageBuffer &command_buffer);
   /** Copy the contents of a texture MIP level to the dst buffer. */
   void copy(VKBuffer &dst_buffer, VKTexture &src_texture, Span<VkBufferImageCopy> regions);
   void copy(VKTexture &dst_texture, VKBuffer &src_buffer, Span<VkBufferImageCopy> regions);

source/blender/gpu/vulkan/vk_storage_buffer.cc

@@ -14,28 +14,32 @@
 namespace blender::gpu {
 
 void VKStorageBuffer::update(const void *data)
+{
+  ensure_allocated();
+  buffer_.update(data);
+}
+
+void VKStorageBuffer::ensure_allocated()
 {
   if (!buffer_.is_allocated()) {
     allocate();
   }
-  buffer_.update(data);
 }
 
 void VKStorageBuffer::allocate()
 {
   buffer_.create(size_in_bytes_,
                  usage_,
-                 static_cast<VkBufferUsageFlagBits>(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
+                 static_cast<VkBufferUsageFlagBits>(VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT |
+                                                    VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
                                                     VK_BUFFER_USAGE_TRANSFER_DST_BIT));
   debug::object_label(buffer_.vk_handle(), name_);
 }
 
 void VKStorageBuffer::bind(int slot)
 {
+  ensure_allocated();
   VKContext &context = *VKContext::get();
-  if (!buffer_.is_allocated()) {
-    allocate();
-  }
   VKShader *shader = static_cast<VKShader *>(context.shader);
   const VKShaderInterface &shader_interface = shader->interface_get();
   const std::optional<VKDescriptorSet::Location> location =
@@ -49,10 +53,8 @@ void VKStorageBuffer::unbind() {}
 
 void VKStorageBuffer::clear(uint32_t clear_value)
 {
+  ensure_allocated();
   VKContext &context = *VKContext::get();
-  if (!buffer_.is_allocated()) {
-    allocate();
-  }
   buffer_.clear(context, clear_value);
 }
 
@@ -61,14 +63,12 @@ void VKStorageBuffer::copy_sub(VertBuf * /*src*/,
                                uint /*src_offset*/,
                                uint /*copy_size*/)
 {
+  NOT_YET_IMPLEMENTED;
 }
 
 void VKStorageBuffer::read(void *data)
 {
-  if (!buffer_.is_allocated()) {
-    allocate();
-  }
-
+  ensure_allocated();
   VKContext &context = *VKContext::get();
   VKCommandBuffer &command_buffer = context.command_buffer_get();
   command_buffer.submit();

source/blender/gpu/vulkan/vk_storage_buffer.hh

@@ -45,8 +45,15 @@ class VKStorageBuffer : public StorageBuf {
     return buffer_.size_in_bytes();
   }
 
+  void ensure_allocated();
+
  private:
   void allocate();
 };
 
+static inline VKStorageBuffer *unwrap(StorageBuf *storage_buffer)
+{
+  return static_cast<VKStorageBuffer *>(storage_buffer);
+}
+
 }  // namespace blender::gpu