From b73c6fb94b1a12095943b8335f32beb5a3991769 Mon Sep 17 00:00:00 2001
From: Stephen Seo <seo.disparate@gmail.com>
Date: Tue, 19 Mar 2024 14:58:35 +0900
Subject: [PATCH] WIP Vulkan compute: compute descriptors, glsl

---
 src/blue_noise.cpp  | 50 +++++++++++++++++++++++++++++++++++++++++++
 src/blue_noise.glsl | 52 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+)
 create mode 100644 src/blue_noise.glsl

diff --git a/src/blue_noise.cpp b/src/blue_noise.cpp
index a9303fc..a570d8f 100644
--- a/src/blue_noise.cpp
+++ b/src/blue_noise.cpp
@@ -266,6 +266,56 @@ image::Bl dither::blue_noise(int width, int height, int threads,
     utility::Cleanup device_cleanup(
         [](void *ptr) { vkDestroyDevice(*((VkDevice *)ptr), nullptr); },
         &device);
+
+    VkQueue compute_queue;
+    vkGetDeviceQueue(device, indices.computeFamily.value(), 0, &compute_queue);
+
+    std::array<VkDescriptorSetLayoutBinding, 4> compute_layout_bindings{};
+    compute_layout_bindings[0].binding = 0;
+    compute_layout_bindings[0].descriptorCount = 1;
+    compute_layout_bindings[0].descriptorType =
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    compute_layout_bindings[0].pImmutableSamplers = nullptr;
+    compute_layout_bindings[0].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+
+    compute_layout_bindings[1].binding = 0;
+    compute_layout_bindings[1].descriptorCount = 1;
+    compute_layout_bindings[1].descriptorType =
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    compute_layout_bindings[1].pImmutableSamplers = nullptr;
+    compute_layout_bindings[1].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+
+    compute_layout_bindings[2].binding = 0;
+    compute_layout_bindings[2].descriptorCount = 1;
+    compute_layout_bindings[2].descriptorType =
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    compute_layout_bindings[2].pImmutableSamplers = nullptr;
+    compute_layout_bindings[2].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+
+    compute_layout_bindings[3].binding = 0;
+    compute_layout_bindings[3].descriptorCount = 1;
+    compute_layout_bindings[3].descriptorType =
+        VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+    compute_layout_bindings[3].pImmutableSamplers = nullptr;
+    compute_layout_bindings[3].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+
+    VkDescriptorSetLayoutCreateInfo layout_info{};
+    layout_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
+    layout_info.bindingCount = compute_layout_bindings.size();
+    layout_info.pBindings = compute_layout_bindings.data();
+
+    VkDescriptorSetLayout compute_desc_set_layout;
+    if (vkCreateDescriptorSetLayout(device, &layout_info, nullptr,
+                                    &compute_desc_set_layout) != VK_SUCCESS) {
+      std::clog << "WARNING: Failed to create compute descriptor set layout!\n";
+      goto ENDOF_VULKAN;
+    }
+    utility::Cleanup compute_desc_set_layout_cleanup(
+        [device](void *ptr) {
+          vkDestroyDescriptorSetLayout(device, *((VkDescriptorSetLayout *)ptr),
+                                       nullptr);
+        },
+        &compute_desc_set_layout);
   }
 ENDOF_VULKAN:
   std::clog << "TODO: Remove this once Vulkan support is implemented.\n";
diff --git a/src/blue_noise.glsl b/src/blue_noise.glsl
new file mode 100644
index 0000000..2b5212e
--- /dev/null
+++ b/src/blue_noise.glsl
@@ -0,0 +1,52 @@
+#version 450
+
+int twoToOne(int x, int y, int width, int height) {
+  while (x < 0) {
+    x += width;
+  }
+  while (y < 0) {
+    y += height;
+  }
+  x = x % width;
+  y = y % height;
+  return x + y * width;
+}
+
+layout(std140, binding = 0) readonly buffer PreComputed {
+  float precomputed[];
+};
+
+layout(std140, binding = 1) writeonly buffer FilterOut { float filter_out[]; };
+
+layout(std140, binding = 2) readonly buffer PBP { int pbp[]; };
+
+layout(std140, binding = 3) readonly buffer Other {
+  int width;
+  int height;
+  int filter_size;
+};
+
+layout(local_size_x = 256) in;
+
+void main() {
+  uint index = gl_GlobalInvocationID.x;
+  if (index >= width * height) {
+    return;
+  }
+
+  int x = index % width;
+  int y = index / width;
+
+  float sum = 0.0F;
+  for (int q = 0; q < filter_size; ++q) {
+    int q_prime = height - filter_size / 2 + y + q;
+    for (int p = 0; p < filter_size; ++p) {
+      int p_prime = width - filter_size / 2 + x + p;
+      if (pbp[twoToOne(p_prime, q_prime, width, height)] != 0) {
+        sum += precomputed[twoToOne(p, q, filter_size, filter_size)];
+      }
+    }
+  }
+
+  filter_out[index] = sum;
+}