Use 8x8 work groups in compute shader

4 files changed, 27 insertions, 20 deletions

diff --git a/include/vk-renderer.hpp b/include/vk-renderer.hpp
index 311b3e5..c1510c8 100644
--- a/include/vk-renderer.hpp
+++ b/include/vk-renderer.hpp
@@ -7,9 +7,9 @@
 #include "vec3.hpp"
 
 enum MATERIAL_TYPE: uint32_t {
-    LAMBERTIAN,
-    METAL,
-    DIELECTRIC,
+    LAMBERTIAN  = 0,
+    METAL       = 1,
+    DIELECTRIC  = 2,
 };
 
 struct material {
@@ -51,7 +51,9 @@ struct input_data {
     uint32_t max_bounce;
     uint32_t samples_per_pixel;
 
-    float padding[2];
+    // Output image resolution
+    uint32_t width;
+    uint32_t height;
 };
 
 struct path_tracer_data {
diff --git a/shaders/compute.comp b/shaders/compute.comp
index 214b24e..c1c4163 100644
--- a/shaders/compute.comp
+++ b/shaders/compute.comp
@@ -1,3 +1,5 @@
+layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
+
 struct ray {
     vec3 origin;
     vec3 direction;
@@ -54,7 +56,8 @@ layout(set = 0, binding = 0) buffer Inputs {
 
     uint max_bounce;
     uint samples_per_pixel;
-    float [2] padding;
+    uint width;
+    uint height;
 
     vec4[] image;
 };
@@ -112,7 +115,7 @@ bool hit_spheres(ray r, float min_t, float max_t, out hit_info info) {
     float closest_t = max_t;
     hit_info temp_info;
     bool hit = false;
-    for (uint sphere_index = 0; sphere_index < 200; sphere_index++) {
+    for (uint sphere_index = 0; sphere_index < 512; sphere_index++) {
         if (spheres[sphere_index].radius > 0.0 && hit_sphere(spheres[sphere_index], r, min_t, closest_t, temp_info)) {
             info = temp_info;
             closest_t = temp_info.t;
@@ -202,15 +205,13 @@ vec4 ray_color(ray r, uint sample_index) {
 void main() {
     uint x = gl_GlobalInvocationID.x;
     uint y = gl_GlobalInvocationID.y;
-    uint width = gl_NumWorkGroups.x;
-    uint height = gl_NumWorkGroups.y;
 
     vec2 uv = vec2(float(x) / (width - 1), 1.0 - (float(y) / (height - 1)));
     vec3 horizontal = cam.focus_distance * cam.right.xyz * cam.viewport_width;
     vec3 vertical = cam.focus_distance * cam.up.xyz * cam.viewport_height;
     vec3 first = cam.position.xyz - horizontal / 2.0 - vertical / 2.0 + cam.focus_distance * cam.forward.xyz;
 
-    image[y * width + x] = vec4(0.0);
+    vec4 out_color = vec4(0.0);
 
     for (uint sample_index = 0; sample_index < samples_per_pixel; sample_index++) {
         vec2 jittered_uvs = uv + vec2(random_disk[sample_index].x / (width - 1), random_disk[sample_index].y / (height - 1));
@@ -218,9 +219,9 @@ void main() {
         vec3 offset = lens_disk.x * cam.right.xyz + lens_disk.y * cam.up.xyz;
         vec3 proj_plane_pos = first + jittered_uvs.x * horizontal + jittered_uvs.y * vertical;
 
-        image[y * width + x] += ray_color(ray(cam.position.xyz + offset, proj_plane_pos - cam.position.xyz - offset), sample_index);
+        out_color += ray_color(ray(cam.position.xyz + offset, proj_plane_pos - cam.position.xyz - offset), sample_index);
     }
 
     float spp_scale = (1.0 / samples_per_pixel);
-    image[y * width + x] = sqrt(image[y * width + x] * spp_scale);
+    image[y * width + x] = sqrt(out_color * spp_scale);
 }
diff --git a/src/main.cpp b/src/main.cpp
index d38545a..f86529f 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -80,12 +80,11 @@ int main() {
     vkrenderer renderer {};
 
     // Image dimensions
-    // const float aspect_ratio = 16.0 / 9.0;
-    const float aspect_ratio = 1.f;
-    const size_t width = 400;
+    const float aspect_ratio = 16.0 / 9.0;
+    const size_t width = 1920;
     const size_t height = width / aspect_ratio;
-    const uint32_t samples_per_pixel = 10;
-    const uint32_t max_depth = 100;
+    const uint32_t samples_per_pixel = 25;
+    const uint32_t max_depth = 50;
 
     const vec3 camera_position{ 13.0, 2.0, 3.0 };
     const vec3 camera_target{ 0.0, 0.0, 0.0 };
@@ -95,12 +94,18 @@ int main() {
     struct input_data inputs = {
         .sky_color = sky_color,
         .ground_color = ground_color,
+
         .cam = { camera_position, camera.up, camera.right, camera.forward, camera.viewport_width, camera.viewport_height, camera.lens_radius, distance_to_focus },
+
         .spheres = {},
+
         .max_bounce = max_depth,
         .samples_per_pixel = samples_per_pixel,
+        .width = (uint32_t)width,
+        .height = (uint32_t)height
     };
 
+    // Random numbers pools
     for (size_t rand_number_index = 0; rand_number_index < samples_per_pixel * 2; rand_number_index += 2) {
         inputs.random_offset[rand_number_index] = randd();
         inputs.random_offset[rand_number_index + 1] = randd();
@@ -117,7 +122,6 @@ int main() {
     }
 
     // World hittable objects
-    std::memset(inputs.spheres, 0, sizeof(inputs.spheres));
     random_scene(inputs.spheres);
 
     // PPM format header
diff --git a/src/vk-renderer.cpp b/src/vk-renderer.cpp
index 4b33a98..d26fe86 100644
--- a/src/vk-renderer.cpp
+++ b/src/vk-renderer.cpp
@@ -46,7 +46,7 @@ void vkrenderer::compute(const input_data& inputs, size_t width, size_t height)
 
     vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_COMPUTE, compute_pipeline);
 
-    vkCmdDispatch(command_buffer, width, height, 1);
+    vkCmdDispatch(command_buffer, width / 8, height / 8, 1);
 
     vkEndCommandBuffer(command_buffer);
 
@@ -60,9 +60,9 @@ void vkrenderer::compute(const input_data& inputs, size_t width, size_t height)
     submit_info.commandBufferCount      = 1;
     submit_info.pCommandBuffers         = &command_buffer;
 
-    vkQueueSubmit(compute_queue, 1, &submit_info, VK_NULL_HANDLE);
+    VKRESULT(vkQueueSubmit(compute_queue, 1, &submit_info, VK_NULL_HANDLE));
 
-    vkDeviceWaitIdle(device);
+    VKRESULT(vkDeviceWaitIdle(device));
 }
 
 const color* const vkrenderer::output_image() {