Change shaders push constants to reflect vulkan binding refactor

1 files changed, 48 insertions, 41 deletions

diff --git a/shaders/compute.comp b/shaders/compute.comp
index 14f1238..934ff70 100644
--- a/shaders/compute.comp
+++ b/shaders/compute.comp
@@ -54,7 +54,9 @@ struct camera {
     vec4 first_pixel;
 
     float lens_radius;
-    float[3] padding;
+    float fov;
+    float focus_distance;
+    float aspect_ratio;
 };
 
 struct bvh_node {
@@ -64,20 +66,6 @@ struct bvh_node {
     int primitive_id;
 };
 
-layout(buffer_reference) readonly buffer scene_metadata {
-    camera cam;
-
-    uint max_bounce;
-    uint min_bounce;
-    uint width;
-    uint height;
-    uint sample_index;
-
-    uint enable_dof;
-    uint debug_bvh;
-    int downscale_factor;
-};
-
 layout(buffer_reference) readonly buffer indices_array {
     uint indices[];
 };
@@ -102,12 +90,25 @@ layout(buffer_reference) readonly buffer nodes_array {
     bvh_node[] nodes;
 };
 
-layout(set = 0, binding = 0) uniform sampler samplers[];
-layout(set = 0, binding = 1) uniform texture2D textures[];
-layout(set = 0, binding = 2, rgba32f) uniform image2D images[];
+layout(buffer_reference) readonly buffer global_params {
+    camera cam;
 
-layout(push_constant) uniform buffers {
-    scene_metadata scene;
+    uint max_bounce;
+    uint min_bounce;
+    uint width;
+    uint height;
+
+    uint sample_index;
+    uint enable_dof;
+    uint debug_bvh;
+    int downscale_factor;
+};
+
+layout(buffer_reference) readonly buffer pass_params {
+    uint padding;
+};
+
+layout(buffer_reference) readonly buffer draw_params {
     nodes_array bvh;
     indices_array indices_arr;
     positions_array positions_arr;
@@ -117,7 +118,17 @@ layout(push_constant) uniform buffers {
     uint output_image_index;
     uint accumulation_image_index;
     uint padding[2];
-} bufs;
+};
+
+layout(set = 0, binding = 0) uniform sampler samplers[];
+layout(set = 0, binding = 1) uniform texture2D textures[];
+layout(set = 0, binding = 2, rgba32f) uniform image2D images[];
+
+layout(push_constant) uniform registers {
+    global_params globals;
+    pass_params pass;
+    draw_params draw;
+} regs;
 
 #include "ray.h"
 
@@ -126,7 +137,7 @@ vec3 ray_color(ray r, uint seed) {
     vec3 radiance = vec3(0.0);
     vec3 throughput = vec3(1.0);
 
-    for (uint bounce = 0; bounce < bufs.scene.max_bounce; bounce++) {
+    for (uint bounce = 0; bounce < regs.globals.max_bounce; bounce++) {
         if (!hit_node(r, info))
             return throughput * vec3(1.0);
 
@@ -236,7 +247,7 @@ vec3 ray_color(ray r, uint seed) {
         // Russian Roulette
         // As the throughput gets smaller, the ray is more likely to get terminated early.
         // Survivors have their value boosted to make up for fewer samples being in the average.
-        if (bounce > bufs.scene.min_bounce) {
+        if (bounce > regs.globals.min_bounce) {
             float probability = min(0.95f, luminance(throughput));
             if (probability < rand(seed)) break;
             throughput /= probability;
@@ -247,46 +258,42 @@ vec3 ray_color(ray r, uint seed) {
 }
 
 ray generate_camera_ray(ivec2 coords, uint seed) {
-    vec2 scene_size = vec2(bufs.scene.width - 1, bufs.scene.height - 1);
+    vec2 scene_size = vec2(regs.globals.width - 1, regs.globals.height - 1);
     vec2 uv = vec2(coords.x / scene_size.x, 1.0 - coords.y / scene_size.y);
     vec2 rand_disk = disk_vec(vec2(rand(seed), rand(seed)));
     vec2 jittered_uvs = uv + vec2(rand(seed), rand(seed)) / scene_size;
     vec2 lens_disk = vec2(0.0);
 
-    if (bufs.scene.enable_dof == 1) {
-        lens_disk = rand_disk * bufs.scene.cam.lens_radius;
+    if (regs.globals.enable_dof == 1) {
+        lens_disk = rand_disk * regs.globals.cam.lens_radius;
     }
 
-    vec3 offset = lens_disk.x * bufs.scene.cam.right.xyz + lens_disk.y * bufs.scene.cam.up.xyz;
-    vec3 proj_plane_pos = bufs.scene.cam.first_pixel.xyz + jittered_uvs.x * bufs.scene.cam.horizontal.xyz + jittered_uvs.y * bufs.scene.cam.vertical.xyz;
+    vec3 offset = lens_disk.x * regs.globals.cam.right.xyz + lens_disk.y * regs.globals.cam.up.xyz;
+    vec3 proj_plane_pos = regs.globals.cam.first_pixel.xyz + jittered_uvs.x * regs.globals.cam.horizontal.xyz + jittered_uvs.y * regs.globals.cam.vertical.xyz;
 
-    return ray(bufs.scene.cam.position.xyz + offset, proj_plane_pos.xyz - bufs.scene.cam.position.xyz - offset, 0.001f, 1e15);
+    return ray(regs.globals.cam.position.xyz + offset, proj_plane_pos.xyz - regs.globals.cam.position.xyz - offset, 0.001f, 1e15);
 }
 
 void main() {
-    if (any(greaterThanEqual(gl_GlobalInvocationID.xy, uvec2(bufs.scene.width, bufs.scene.height))))
+    if (any(greaterThanEqual(gl_GlobalInvocationID.xy, uvec2(regs.globals.width, regs.globals.height))))
         return;
 
-    ivec2 coords = ivec2(gl_GlobalInvocationID.xy) * bufs.scene.downscale_factor;
+    ivec2 coords = ivec2(gl_GlobalInvocationID.xy) * regs.globals.downscale_factor;
 
-    uint seed = uint(coords.x * uint(1973) + coords.y * uint(9277) + bufs.scene.sample_index * uint(26699)) | uint(1);
+    uint seed = uint(coords.x * uint(1973) + coords.y * uint(9277) + regs.globals.sample_index * uint(26699)) | uint(1);
     ray r = generate_camera_ray(coords, seed);
 
     vec3 out_color = vec3(0.0);
-    if (bufs.scene.debug_bvh == 1) {
+    if (regs.globals.debug_bvh == 1) {
         out_color = hit_aabbs(r) * vec3(0.001, 0.0, 0.0);
     } else {
         out_color = ray_color(r, seed);
     }
 
-    float spp_scale = (1.0 / bufs.scene.sample_index);
-    vec3 acc_color = srgb_to_linear(imageLoad(images[nonuniformEXT(bufs.accumulation_image_index)], coords).xyz);
+    // TODO: what is going on with srgb here ?
+    float spp_scale = (1.0 / regs.globals.sample_index);
+    vec3 acc_color = srgb_to_linear(imageLoad(images[nonuniformEXT(regs.draw.accumulation_image_index)], coords).xyz);
     out_color = linear_to_srgb(mix(acc_color, out_color, spp_scale));
 
-    // for (uint i = 0; i < bufs.scene.downscale_factor; i++) {
-    //     for (uint j = 0; j < bufs.scene.downscale_factor; j++) {
-    //         imageStore(images[nonuniformEXT(bufs.output_image_index)], coords + ivec2(i, j), vec4(out_color, 1.0));
-    //     }
-    // }
-    imageStore(images[nonuniformEXT(bufs.output_image_index)], coords, vec4(out_color, 1.0));
+    imageStore(images[nonuniformEXT(regs.draw.output_image_index)], coords, vec4(out_color, 1.0));
 }