Wip BVH + BVH debugging tools

6 files changed, 223 insertions, 89 deletions

diff --git a/include/bvh-node.hpp b/include/bvh-node.hpp
index 78780e7..ee5068b 100644
--- a/include/bvh-node.hpp
+++ b/include/bvh-node.hpp
@@ -3,6 +3,7 @@
 

 #include <vector>

 #include <algorithm>

+#include <fstream>

 

 #include "aabb.hpp"

 #include "sphere.hpp"

@@ -11,44 +12,137 @@ class bvh_node {
     public:

         bvh_node() = default;

         bvh_node(sphere& sphere, int32_t sphere_id): bounding_box(sphere.bounding_box), sphere_id(sphere_id){}

-        bvh_node(sphere* spheres, size_t start, size_t end, bvh_node* root_node, int32_t id) {

-            left_node_id = id * 2;

-            right_node_id = id * 2 + 1;

+        // bvh_node(sphere* spheres, size_t start, size_t end, bvh_node* root_node, int32_t id) {

+        //     left_node_id = id * 2;

+        //     right_node_id = id * 2 + 1;

 

-            bvh_node left, right;

+        //     bvh_node left, right;

 

-            // int axis = randi(0, 2);

-            // auto comparator = (axis == 0) ? box_x_comparator : (axis == 1) ? box_y_comparator : box_z_comparator;

+        //     // int axis = randi(0, 2);

+        //     // auto comparator = (axis == 0) ? box_x_comparator : (axis == 1) ? box_y_comparator : box_z_comparator;

 

-            size_t span = end - start;

+        //     size_t span = end - start;

 

-            if (span == 1) {

-                root_node[left_node_id - 1] = root_node[right_node_id - 1] = { spheres[start], (int32_t)start };

-            } else if (span == 2) {

-                // if (comparator(spheres[start], spheres[start + 1]) {

-                    root_node[left_node_id - 1] = { spheres[start], (int32_t)start };

-                    root_node[right_node_id - 1] = { spheres[start], (int32_t)start + 1 };

-                // } else {

-                //     root_node[left_node_id - 1] = { spheres[start + 1], start + 1 };

-                //     root_node[right_node_id - 1] = { spheres[start], start };

-                // }

+        //     if (span == 1) {

+        //         root_node[left_node_id - 1] = root_node[right_node_id - 1] = { spheres[start], (int32_t)start };

+        //     } else if (span == 2) {

+        //         // if (comparator(spheres[start], spheres[start + 1]) {

+        //             root_node[left_node_id - 1] = { spheres[start], (int32_t)start };

+        //             root_node[right_node_id - 1] = { spheres[start], (int32_t)start + 1 };

+        //         // } else {

+        //         //     root_node[left_node_id - 1] = { spheres[start + 1], start + 1 };

+        //         //     root_node[right_node_id - 1] = { spheres[start], start };

+        //         // }

+        //     } else {

+        //         uint32_t mid = start + span / 2;

+        //         root_node[left_node_id - 1] = { spheres, start, mid, root_node, left_node_id };

+        //         root_node[right_node_id - 1] = { spheres, mid, end, root_node, right_node_id };

+        //     }

+

+        //     bounding_box = aabb::surrounding_box(root_node[left_node_id - 1].bounding_box, root_node[right_node_id -1].bounding_box);

+        // }

+

+        aabb bounding_box;

+

+        int32_t left_id    = -1;

+

+        int32_t sphere_id   = -1;

+        uint32_t padding[2];

+};

+

+class bvh {

+public: 

+    // TODO: Fix one level too deep BVH

+    void build(sphere* spheres, uint32_t spheres_count) {

+        nodes = std::vector<bvh_node>(spheres_count * spheres_count);

+        boxes = std::vector<aabb>(spheres_count);

+

+        for (uint32_t sphere_id = 0; sphere_id < spheres_count; sphere_id++) {

+            boxes[sphere_id] = spheres[sphere_id].bounding_box;

+        }

+

+        nodes[0] = bvh_node();

+

+        nodes[0].bounding_box = compute_bounds(0, spheres_count);

+        subdivide(0, 0, spheres_count);

+    }

+

+    //-------------------------

+    // BVH exporter for graphviz

+    //-------------------------

+    void exporter() {

+        std::cout << "digraph bvh {" << std::endl;

+

+        std::cout << "\t r";

+        preorder(nodes[0].left_id, nodes[0].sphere_id);

+

+        std::cout << "\t r";

+        preorder(nodes[0].left_id + 1, nodes[0].sphere_id);

+

+        std::cout << "}" << std::endl;

+    }

+

+    void preorder(int32_t id, int32_t sphere_id) {

+        if (id <= 0) {

+            if (sphere_id < 0) {

+                std::cout << ";" << std::endl;

             } else {

-                uint32_t mid = start + span / 2;

-                root_node[left_node_id - 1] = { spheres, start, mid, root_node, left_node_id };

-                root_node[right_node_id - 1] = { spheres, mid, end, root_node, right_node_id };

+                std::cout << " -> sphere_" << sphere_id << ";" << std::endl;

             }

+            return;

+        }

 

-            bounding_box = aabb::surrounding_box(root_node[left_node_id - 1].bounding_box, root_node[right_node_id -1].bounding_box);

+        std::cout << " -> " << id;

+        preorder(nodes[id].left_id, nodes[id].sphere_id);

+

+        if (nodes[id].left_id > -1) {

+            std::cout << "\t " << id;

+            preorder(nodes[id].left_id + 1, nodes[id].sphere_id);

         }

+    }

 

-        aabb bounding_box;

+private:

+

+    aabb compute_bounds(uint32_t begin, uint32_t end) {

+        aabb global_box = aabb();

+

+        for (uint32_t box_id = begin; box_id < end; box_id++) {

+            global_box = aabb::surrounding_box(global_box, boxes[box_id]);

+        }

+

+        return global_box;

+    }

 

-        // Store sphere id ?

-        int32_t left_node_id    = -1;

-        int32_t right_node_id   = -1;

+    void subdivide(uint32_t parent_id, uint32_t begin, uint32_t end) {

+        if ((end - begin) == 0) {

+            return;

+        }

+

+        auto left_id = parent_id * 2 + 1;

+        auto right_id = parent_id * 2 + 2;

+

+        nodes[parent_id].left_id = left_id;

+

+        nodes[left_id] = bvh_node();

+        nodes[right_id] = bvh_node();

+

+        nodes[left_id].bounding_box = compute_bounds(begin, begin + ((end - begin) / 2));

+        nodes[right_id].bounding_box = compute_bounds(begin + ((end - begin) / 2), end);

+

+        if ((end - begin) == 2) {

+            nodes[left_id].sphere_id = begin;

+            nodes[right_id].sphere_id = begin + 1;

+        } else if ((end - begin) == 1) {

+            nodes[left_id].sphere_id = begin;

+        } else {

+            subdivide(left_id, begin, begin + ((end - begin) / 2));

+            subdivide(right_id, begin + ((end - begin) / 2), end);

+        }

+    }

 

-        int32_t sphere_id       = -1;

-        uint32_t padding;

+public:

+    std::vector<bvh_node> nodes;

+    std::vector<aabb> boxes;

 };

 

 #endif // !__BVH_NODE_HPP_

diff --git a/include/utils.hpp b/include/utils.hpp
index a95c95d..ee48d2c 100644
--- a/include/utils.hpp
+++ b/include/utils.hpp
@@ -28,7 +28,8 @@ struct input_data {
     uint32_t sample_index;
 
     uint32_t spheres_count;
-    uint32_t padding[3];
+    uint32_t debug_bvh;
+    uint32_t padding[2];
 };
 
 std::vector<uint8_t> get_shader_code(const char* path);
diff --git a/include/window.hpp b/include/window.hpp
index 5a8bcc8..d2a8a78 100644
--- a/include/window.hpp
+++ b/include/window.hpp
@@ -38,6 +38,7 @@ enum KEYS: uint32_t {
     SHIFT = 0x10,
     CTRL,
     ALT,
+    SPACE = 0x20,
     A = 0x41,
     B,
     C,
diff --git a/shaders/compute.comp b/shaders/compute.comp
index a6fd9e4..b11a296 100644
--- a/shaders/compute.comp
+++ b/shaders/compute.comp
@@ -61,9 +61,8 @@ struct camera {
 struct bvh_node {
     aabb bb;
     int left;
-    int right;
     int sphere_id;
-    uint padding;
+    uint[2] padding;
 };
 
 layout(set = 0, binding = 0) buffer Inputs {
@@ -80,7 +79,8 @@ layout(set = 0, binding = 0) buffer Inputs {
     uint sample_index;
 
     uint spheres_count;
-    uint[3] padding;
+    bool debug_bvh;
+    uint[2] padding;
 };
 
 layout(set = 0, binding = 1, rgba32f) uniform writeonly image2D output_image;
@@ -103,9 +103,7 @@ uint wang_hash(inout uint seed)
 
 float rand(inout uint seed)
 {
-    float r = wang_hash(seed) * (1.0 / 4294967296.0);
-    // seed += 10;
-    return r;
+    return float(wang_hash(seed)) / 4294967296.0;
 }
 
 vec3 unit_vec(inout uint seed) {
@@ -133,13 +131,13 @@ bool hit_aabb(aabb b, ray r, float min_t, float max_t) {
         if (inv_dir < 0.f) {
             float temp = t0;
             t0 = t1;
-            t1 = t0;
+            t1 = temp;
         }
 
-        min_t = max(t0, min_t);
-        max_t = min(t1, max_t);
+        min_t = t0 > min_t ? t0 : min_t;
+        max_t = t1 < max_t ? t1 : max_t;
 
-        if (max_t < min_t) {
+        if (max_t <= min_t) {
             return false;
         }
     }
@@ -147,45 +145,59 @@ bool hit_aabb(aabb b, ray r, float min_t, float max_t) {
     return true;
 }
 
-bool hit_node(bvh_node node, ray r, float min_t, float max_t, out hit_info info) {
-    if (!hit_aabb(node.bb, r, min_t, max_t)) {
-        return false;
-    }
-
-    int stack_id = -1;
-    int stack[64];
-    bvh_node current_node = node;
-
-    while(true) {
-        bool hit_left = hit_aabb(nodes[current_node.left - 1].bb, r, min_t, max_t);
-        bool hit_right = hit_aabb(nodes[current_node.right - 1].bb, r, min_t, max_t);
-
-        if (hit_left && hit_right) {
-            stack_id++;
-            stack[stack_id] = current_node.left - 1;
+uint hit_aabbs(ray r, float min_t, float max_t) {
+    uint hit_count = 0;
+    for (uint idx = 0; idx <= spheres_count * 2; idx++) {
+        hit_count += uint(hit_aabb(nodes[idx].bb, r, min_t, max_t));
 
-            current_node = nodes[current_node.right - 1];
-        } else if (hit_left) {
-            current_node = nodes[current_node.left - 1];
-        } else if (hit_right) {
-            current_node = nodes[current_node.right - 1];
-        } else {
-            if (stack_id >= 0) {
-                current_node = nodes[stack[stack_id]];
-                stack_id--;
-            } else {
-                return false;
-            }
-        }
-
-        // Reached a leaf
-        if (current_node.left == -1 && current_node.right == -1) {
-            return true;
+        if (nodes[idx].sphere_id >= 0) {
+            hit_count += uint(hit_aabb(spheres[nodes[idx].sphere_id].bb, r, min_t, max_t));
         }
     }
 
+    return hit_count;
 }
 
+// bool hit_node(ray r, float min_t, float max_t, out hit_info info) {
+//     if (!hit_aabb(nodes[0].bb, r, min_t, max_t)) {
+//         return false;
+//     }
+//     return true;
+// 
+//     int stack_id = -1;
+//     int stack[64];
+//     bvh_node current_node = nodes[0];
+// 
+//     while(true) {
+//         bool hit_left = hit_aabb(nodes[current_node.left - 1].bb, r, min_t, max_t);
+//         bool hit_right = hit_aabb(nodes[current_node.right - 1].bb, r, min_t, max_t);
+// 
+//         if (hit_left && hit_right) {
+//             stack_id++;
+//             stack[stack_id] = current_node.left - 1;
+// 
+//             current_node = nodes[current_node.right - 1];
+//         } else if (hit_left) {
+//             current_node = nodes[current_node.left - 1];
+//         } else if (hit_right) {
+//             current_node = nodes[current_node.right - 1];
+//         } else {
+//             if (stack_id >= 0) {
+//                 current_node = nodes[stack[stack_id]];
+//                 stack_id--;
+//             } else {
+//                 return false;
+//             }
+//         }
+// 
+//         // Reached a leaf
+//         if (current_node.left == -1 && current_node.right == -1) {
+//             return true;
+//         }
+//     }
+// 
+// }
+
 bool hit_sphere(sphere s, ray r, float min_t, float max_t, out hit_info info) {
     vec3 origin_center = r.origin - s.position.xyz;
 
@@ -299,20 +311,20 @@ vec4 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.
-            {
-                float p = max(throughput.r, max(throughput.g, throughput.b));
-                if (rand(seed) > p) {
-                    break;
-                }
+            // {
+            //     float p = max(throughput.r, max(throughput.g, throughput.b));
+            //     if (rand(seed) > p) {
+            //         break;
+            //     }
 
-                // Add the energy we 'lose' by randomly terminating paths
-                throughput *= 1.0 / p;
-            }
+            //     // Add the energy we 'lose' by randomly terminating paths
+            //     throughput *= 1.0 / p;
+            // }
         }
         bounces++;
     }
 
-    return final_color;
+    return throughput;
 }
 
 void main() {
@@ -322,16 +334,22 @@ void main() {
     vec2 uv = vec2(float(x) / (width - 1), 1.0 - (float(y) / (height - 1)));
 
     uint seed = uint(x * uint(1973) + y * uint(9277) + sample_index * uint(26699)) | uint(1);
-    vec2 rand_disk = disk_vec(seed);
-    vec2 jittered_uvs = uv + vec2(rand_disk.x / (width - 1), rand_disk.y / (height - 1));
-    vec2 lens_disk = rand_disk * cam.lens_radius;
-    vec3 offset = lens_disk.x * cam.right.xyz + lens_disk.y * cam.up.xyz;
+    // vec2 rand_disk = disk_vec(seed);
+    // vec2 jittered_uvs = uv + vec2(rand(seed) / (width - 1), rand(seed) / (height - 1));
+    vec2 jittered_uvs = uv;
+    // vec2 lens_disk = rand_disk * cam.lens_radius;
+    // vec3 offset = lens_disk.x * cam.right.xyz + lens_disk.y * cam.up.xyz;
+    vec3 offset = vec3(0.0);
     vec4 proj_plane_pos = cam.first_pixel + jittered_uvs.x * cam.horizontal + jittered_uvs.y * cam.vertical;
 
     ray r = ray(cam.position.xyz + offset, proj_plane_pos.xyz - cam.position.xyz - offset);
 
     vec4 out_color = ray_color(r, seed);
 
+    if (debug_bvh) {
+        out_color -= hit_aabbs(r, 0.001, 1e15) * vec4(0.0, 0.05, 0.05, 0.0);
+    }
+
     float spp_scale = (1.0 / sample_index);
     vec4 acc_color = imageLoad(accumulation_image, ivec2(x, y));
     imageStore(output_image, ivec2(x, y), vec4(mix(acc_color, out_color, spp_scale)));
diff --git a/src/main.cpp b/src/main.cpp
index 071903a..19b7b54 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -103,8 +103,8 @@ int main() {
                 }
                 case EVENT_TYPES::KEY_PRESS: {
                     auto move_speed = 0.01f * delta_time;
-                    if (wnd.keyboard[KEYS::LSHIFT] || wnd.keyboard[KEYS::RSHIFT]) {
-                        move_speed *= 2.f;
+                    if (wnd.keyboard[KEYS::LSHIFT] || wnd.keyboard[KEYS::RSHIFT] || wnd.keyboard[KEYS::SHIFT]) {
+                        move_speed *= 10.f;
                     }
 
                     vec3 move_vec {};
@@ -126,6 +126,14 @@ int main() {
                             move_vec = -((input_data*) renderer.compute_shader_buffer.alloc_info.pMappedData)->cam.right.unit() * move_speed;
                             break;
                         }
+                        case KEYS::SPACE: {
+                            move_vec = ((input_data*) renderer.compute_shader_buffer.alloc_info.pMappedData)->cam.up.unit() * move_speed;
+                            break;
+                        }
+                        case KEYS::CTRL: {
+                            move_vec = -((input_data*) renderer.compute_shader_buffer.alloc_info.pMappedData)->cam.up.unit() * move_speed;
+                            break;
+                        }
                         default:
                             break;
                     }
@@ -183,6 +191,10 @@ int main() {
 
         ImGui::SliderInt("bounces", (int32_t*)&((input_data*) renderer.compute_shader_buffer.alloc_info.pMappedData)->max_bounce, 2, 250);
 
+        if (ImGui::Checkbox("debug bvh", (bool*)&((input_data*) renderer.compute_shader_buffer.alloc_info.pMappedData)->debug_bvh)) {
+            reset_accumulation = true;
+        }
+
         ImGui::End();
 
         ImGui::EndFrame();
@@ -226,6 +238,7 @@ int main() {
     //     write_color(std::cout, output_image[index]);
     // }
 
+
     std::cerr << "Done !" << std::endl;
 
     return 0;
diff --git a/src/utils.cpp b/src/utils.cpp
index 6467563..bbf8ea4 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -106,12 +106,12 @@ uint32_t random_scene(sphere *world) {
 
     material material4 = {};
     material4.albedo = color{ 0.0, 0.0, 0.0 };
-    material4.emissive = color{ 1.0, 1.0, 1.0 } * 100.0;
+    material4.emissive = color{ 1.0, 1.0, 1.0 } * 100000.0;
     material4.fuzz = 0.0;
     material4.type = MATERIAL_TYPE::EMISSIVE;
     world[world_sphere_index + 3] = { { 2, 1.5, 2 }, material4, 0.5 };
 
-    return world_sphere_index + 3;
+    return world_sphere_index + 4;
 }
 
 struct input_data create_inputs(uint32_t width, uint32_t height) {
@@ -141,6 +141,7 @@ struct input_data create_inputs(uint32_t width, uint32_t height) {
         .height = height,
 
         .sample_index = 0,
+        .debug_bvh = false,
     };
 
     // World hittable objects
@@ -148,7 +149,13 @@ struct input_data create_inputs(uint32_t width, uint32_t height) {
 
     inputs.spheres_count = spheres_count;
 
-    inputs.nodes[0] = { inputs.spheres, 0, spheres_count, inputs.nodes, 1 };
+    bvh builder;
+    builder.build(inputs.spheres, spheres_count);
+    builder.exporter();
+
+    for (uint32_t idx = 0; idx < builder.nodes.size(); idx++) {
+        inputs.nodes[idx] = builder.nodes[idx];
+    }
 
     return inputs;
 }