Use 8x8 work groups in compute shader

1 files changed, 8 insertions, 7 deletions

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);
 }