From 738c765557df281f4980b543172a3363a485b592 Mon Sep 17 00:00:00 2001
From: Anton Ljungdahl <anton.ljungdahl@fysik.su.se>
Date: Sun, 11 May 2025 14:42:33 +0200
Subject: [PATCH] fastish build of SAH BVH

---
 src/rayt_base.odin |  34 ++-------
 src/rayt_bvh.odin  | 186 +++++++++++++++++++++++++++++++++++++++------
 2 files changed, 169 insertions(+), 51 deletions(-)

diff --git a/src/rayt_base.odin b/src/rayt_base.odin
index bbaa3a9..756f99b 100644
--- a/src/rayt_base.odin
+++ b/src/rayt_base.odin
@@ -51,7 +51,8 @@ Viewport :: struct {
 
 Ray :: struct {
   origin : Vec3,
-  direction : Vec3
+  direction : Vec3,
+  inv_dir : Vec3
 }
 
 HitRecord :: struct {
@@ -108,6 +109,7 @@ ray_get :: proc(x : f32, y : f32) -> Ray {
 
   out.direction = ray_direction
   out.origin = camera.center
+  out.inv_dir = Vec3{1.0/out.direction.x, 1.0/out.direction.y, 1.0/out.direction.z}
 
   return out
 }
@@ -116,29 +118,6 @@ rayt_cpu_main :: proc() {
   rand.reset(RAND_SEED)
    
   load_triangles()
-   
-  // Random triangles
-  //{
-  //  center_shift := Vec3{5.0, 5.0, 5.0} 
-  //  // Generate triangles inside a box 
-  //  for i in 0..<MAX_NUM_ENTITIES {
-  //    r0 := vec3_rand_uniform();
-  //    r1 := vec3_rand_uniform();
-  //    r2 := vec3_rand_uniform();
-//
-  //    // Put the first vertex within a 10x10x10 cube centered at the origin
-  //    v0 := 9.0*r0
-  //    v0 = v0 - center_shift
-  //    
-  //    entities[i].kind = EntityKind.Tri
-  //    entities[i].v0 = v0
-  //    entities[i].v1 = v0 + r1
-  //    entities[i].v2 = v0 + r2
-  //    entities[i].center = 0.3333*(entities[i].v0 + entities[i].v1 + entities[i].v2)
-  //    tri_indices[i] = cast(u32)i
-  //    //fmt.printf("Triangle idx %i \n", tri_indices[i])
-  //  }
-  //}
 
   // Set up scene globals
   {
@@ -178,7 +157,11 @@ rayt_cpu_main :: proc() {
   
   // build bvh
   fmt.println("Building BVH")
+  time.stopwatch_start(&stopwatch)
   bvh_build()
+  time.stopwatch_stop(&stopwatch)
+  elapsed_bvh_ms := elapsed_time_ms()
+  fmt.printf("Build BVH in %.4f ms \n", elapsed_bvh_ms)
   bvh_stats()
   
   fmt.printf("Starting CPU raytracing")
@@ -251,8 +234,7 @@ triangle_intersection :: proc(ray : ^Ray, rec : ^HitRecord, triangle : ^Entity)
     }
   }
   
-
-
+  // If we have an intersection closer than the last, we fill out the hit record
   if rec.t < closest_so_far {
     intersection_point := ray_point(rec.t, ray)
     v0_normal := linalg.cross(edge1, edge2)
diff --git a/src/rayt_bvh.odin b/src/rayt_bvh.odin
index 9bc21f4..ab1cbbf 100644
--- a/src/rayt_bvh.odin
+++ b/src/rayt_bvh.odin
@@ -5,7 +5,7 @@ import "core:math"
 import "core:math/linalg"
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
-
+// Struct definitions
 BVHNode :: struct {
   aabb_min : Vec3,
   aabb_max : Vec3,
@@ -13,37 +13,52 @@ BVHNode :: struct {
   tri_count : u32
 }
 
+// Helper structure used to define a number of primtiives inside a given non-leaf BVH node,
+// and their bounding box. This is used to compute the cost of a given split of that node.
+BVHBin :: struct {
+  aabb_min : Vec3,
+  aabb_max : Vec3, 
+  tri_count : u32,
+}
+
 BVH :: struct {
   nodes : []BVHNode,
   used_nodes : u32,
   root_index : u32,
   max_num_nodes : u32,
   num_leaf_nodes : u32,
-  num_leaf_entities : u32
+  num_leaf_entities : u32,
 }
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
-
+// Globals
 bvh : BVH
-
+BVH_NUM_BINS :: 8
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+bvh_bins_init :: proc(bins : []BVHBin) {
+  for i in 0..<BVH_NUM_BINS {
+    bins[i].aabb_min = Vec3{math.F32_MAX, math.F32_MAX, math.F32_MAX}
+    bins[i].aabb_max = Vec3{-math.F32_MAX, -math.F32_MAX, -math.F32_MAX}
+  }
+}
+
 intersect_aabb :: proc(ray : ^Ray, bmin : Vec3, bmax : Vec3, closest_so_far : f32) -> b32 {
-  tx1 : f32 = (bmin.x - ray.origin.x) / ray.direction.x
-  tx2 : f32 = (bmax.x - ray.origin.x) / ray.direction.x
+  tx1 : f32 = (bmin.x - ray.origin.x) * ray.inv_dir.x
+  tx2 : f32 = (bmax.x - ray.origin.x) * ray.inv_dir.x
   tmin : f32 = math.min(tx1, tx2)
   tmax : f32 = math.max(tx1, tx2)
 
-  ty1 : f32 = (bmin.y - ray.origin.y) / ray.direction.y
-  ty2 : f32 = (bmax.y - ray.origin.y) / ray.direction.y
+  ty1 : f32 = (bmin.y - ray.origin.y) * ray.inv_dir.y
+  ty2 : f32 = (bmax.y - ray.origin.y) * ray.inv_dir.y
   tmin = math.max(tmin, math.min(ty1, ty2))
   tmax = math.min(tmax, math.max(ty1, ty2))
 
-  tz1 : f32 = (bmin.z - ray.origin.z) / ray.direction.z
-  tz2 : f32 = (bmax.z - ray.origin.z) / ray.direction.z
+  tz1 : f32 = (bmin.z - ray.origin.z) * ray.inv_dir.z
+  tz2 : f32 = (bmax.z - ray.origin.z) * ray.inv_dir.z
   tmin = math.max(tmin, math.min(tz1, tz2))
   tmax = math.min(tmax, math.max(tz1, tz2))
 
@@ -51,6 +66,40 @@ intersect_aabb :: proc(ray : ^Ray, bmin : Vec3, bmax : Vec3, closest_so_far : f3
   return out
 }
 
+
+aabb_area :: proc(aabb_min : Vec3, aabb_max : Vec3) -> f32 {
+  e := aabb_max - aabb_min
+  return e.x * e.y + e.y * e.z + e.z * e.x
+}
+
+aabb_grow :: proc(aabb_min : Vec3, aabb_max : Vec3, p_min : Vec3, p_max : Vec3) -> (Vec3, Vec3) {
+  out_aabb_min := aabb_min
+  out_aabb_max := aabb_max
+  if(p_min.x < math.F32_MAX) {
+    out_aabb_min = linalg.min(out_aabb_min, p_min)
+    out_aabb_max = linalg.max(out_aabb_max, p_min)
+    out_aabb_min = linalg.min(out_aabb_min, p_max)
+    out_aabb_max = linalg.max(out_aabb_max, p_max)
+  }
+  return out_aabb_min, out_aabb_max
+}
+
+aabb_min_triangle :: proc(aabb_min : Vec3, tri : ^Entity) -> Vec3 {
+  out_aabb_min := aabb_min
+  out_aabb_min = linalg.min(out_aabb_min, tri.v0)
+  out_aabb_min = linalg.min(out_aabb_min, tri.v1)
+  out_aabb_min = linalg.min(out_aabb_min, tri.v2)
+  return out_aabb_min
+}
+
+aabb_max_triangle :: proc(aabb_max : Vec3, tri : ^Entity) -> Vec3 {
+  out_aabb_max := aabb_max
+  out_aabb_max = linalg.max(out_aabb_max, tri.v0)
+  out_aabb_max = linalg.max(out_aabb_max, tri.v1)
+  out_aabb_max = linalg.max(out_aabb_max, tri.v2)
+  return out_aabb_max
+}
+
 bvh_update_bounds :: proc(node_idx : u32) {
   node : ^BVHNode = &bvh.nodes[node_idx]
 
@@ -61,28 +110,113 @@ bvh_update_bounds :: proc(node_idx : u32) {
   for i in 0..<node.tri_count {
     leaf_tri_idx := tri_indices[first_tri_idx + i]
     triangle : ^Entity = &entities[leaf_tri_idx]
-    node.aabb_min = linalg.min(node.aabb_min, triangle.v0)
-    node.aabb_min = linalg.min(node.aabb_min, triangle.v1)
-    node.aabb_min = linalg.min(node.aabb_min, triangle.v2)
-    node.aabb_max = linalg.max(node.aabb_max, triangle.v0)
-    node.aabb_max = linalg.max(node.aabb_max, triangle.v1)
-    node.aabb_max = linalg.max(node.aabb_max, triangle.v2)
+    node.aabb_min = aabb_min_triangle(node.aabb_min, triangle)
+    node.aabb_max = aabb_max_triangle(node.aabb_max, triangle)
   }
 }
 
+find_best_split_plane :: proc(node : ^BVHNode, out_axis : ^u32, out_split_pos : ^f32) -> f32 {
+  best_cost : f32 = math.F32_MAX
+
+  // Loop over each axis
+  for axis in 0..<3  {
+    bounds_min : f32 = math.F32_MAX
+    bounds_max : f32 = -math.F32_MAX
+    // Find the bounds of all the primitive centers in the node
+    for i in 0..<node.tri_count {
+      tri : ^Entity = &entities[tri_indices[node.left_first + i]]
+      bounds_min = math.min(bounds_min, tri.center[axis])
+      bounds_max = math.max(bounds_max, tri.center[axis])
+    }
+
+    if bounds_min == bounds_max { continue }
+
+    bins : [BVH_NUM_BINS]BVHBin
+    bvh_bins_init(bins[:])
+    bin_scale : f32 = cast(f32)BVH_NUM_BINS / (bounds_max - bounds_min)
+
+    // We put all the primitives in the node in any one of the binds,
+    // depending on the primitive's centroid pos. What we are doing is really 
+    // just splitting the node with BVH_NUM_BINS-1 number of planes.
+
+    for i in 0..<node.tri_count {
+      tri : ^Entity = &entities[tri_indices[node.left_first + i]]
+      primitive_index : u32 = cast(u32)((tri.center[axis] - bounds_min) * bin_scale)
+      bin_idx : u32 = math.min(BVH_NUM_BINS - 1, primitive_index)
+      bins[bin_idx].tri_count += 1
+      bins[bin_idx].aabb_min = aabb_min_triangle(bins[bin_idx].aabb_min, tri)
+      bins[bin_idx].aabb_max = aabb_max_triangle(bins[bin_idx].aabb_min, tri)
+    }
+
+    // Gather data for all BVH_NUM_BINS-1 planes
+    left_area : [BVH_NUM_BINS - 1]f32
+    right_area : [BVH_NUM_BINS - 1]f32
+    left_count : [BVH_NUM_BINS - 1]u32
+    right_count : [BVH_NUM_BINS - 1]u32
+    left_box_aabb_min : Vec3 = Vec3{math.F32_MAX, math.F32_MAX, math.F32_MAX}
+    left_box_aabb_max : Vec3 = Vec3{-math.F32_MAX, -math.F32_MAX, -math.F32_MAX}
+    right_box_aabb_min : Vec3 = Vec3{math.F32_MAX, math.F32_MAX, math.F32_MAX}
+    right_box_aabb_max : Vec3 = Vec3{-math.F32_MAX, -math.F32_MAX, -math.F32_MAX}
+    left_sum : u32 = 0
+    right_sum : u32 = 0
+    // Loop from both sides simultaneously
+    for i in 0..<BVH_NUM_BINS-1 {
+      left_sum += bins[i].tri_count
+      left_count[i] = left_sum
+      left_box_aabb_min, left_box_aabb_max = aabb_grow(left_box_aabb_min, left_box_aabb_max, 
+                                                       bins[i].aabb_min, bins[i].aabb_max)
+      left_area[i] = aabb_area(left_box_aabb_min, left_box_aabb_max)
+      right_idx := BVH_NUM_BINS - 1 - i
+      right_sum += bins[right_idx].tri_count
+      right_count[right_idx-1] = right_sum
+      right_box_aabb_min, right_box_aabb_max = aabb_grow(right_box_aabb_min, 
+                                                         right_box_aabb_max, 
+                                                         bins[right_idx].aabb_min, 
+                                                         bins[right_idx].aabb_max)
+
+      right_area[right_idx-1] = aabb_area(right_box_aabb_min, right_box_aabb_max)
+    }
+
+    plane_scale : f32 = (bounds_max - bounds_min) / cast(f32)BVH_NUM_BINS
+    // Compute the Surface area heuristic (SAH) cost for each plane
+    for i in 0..<BVH_NUM_BINS-1 {
+      plane_cost : f32 = 0.0
+      plane_cost += cast(f32)left_count[i] * left_area[i] 
+      plane_cost += cast(f32)right_count[i] * right_area[i]
+      
+      if(plane_cost < best_cost) {
+        out_axis^ = cast(u32)axis
+        out_split_pos^ = bounds_min + plane_scale * cast(f32)(i + 1)
+        best_cost = plane_cost
+      }
+    }
+
+  }
+
+  return best_cost
+}
+
 bvh_subdivide :: proc(node_idx : u32) {
   node : ^BVHNode = &bvh.nodes[node_idx]
+  if node.tri_count < bvh.num_leaf_entities { 
+    return
+  }
 
-  if(node.tri_count <= bvh.num_leaf_entities) {
+  axis : u32 = 0
+  split_pos : f32
+  split_cost := find_best_split_plane(node, &axis, &split_pos)
+  node_split_cost : f32 = 0.0
+  {
+    e : Vec3 = node.aabb_max - node.aabb_min
+    surface_area := e.x * e.y + e.y * e.z + e.z * e.x
+    node_split_cost = cast(f32)node.tri_count * surface_area
+  }
+
+  if(split_cost >= node_split_cost) {
     return;
   }
 
-  extent := node.aabb_max -  node.aabb_min
-  axis : u32 = 0
-  if(extent.y > extent.x) {axis = 1}
-  if(extent.z > extent[axis]) {axis = 2}
-  split_pos : f32 = node.aabb_min[axis] + extent[axis] * 0.5
-
+  
   i : u32 = node.left_first
   j : u32 = node.tri_count + i - 1
 
@@ -150,8 +284,9 @@ bvh_build :: proc() {
   num_triangles : u32 = cast(u32)len(tri_indices)
   bvh.max_num_nodes = 2 * num_triangles - 1
   bvh.nodes = make([]BVHNode, bvh.max_num_nodes)
-  bvh.used_nodes = 2 // We skip first two nodes, for some reason. TODO comment this, read the tutorial
-  bvh.num_leaf_entities = 2
+  bvh.used_nodes = 2 // We skip first two nodes, for some reason. 
+                     //TODO comment this, read the tutorial
+  bvh.num_leaf_entities = 4
   bvh.root_index = 0
 
   // Init root node
@@ -184,5 +319,6 @@ bvh_stats :: proc() {
     fmt.printf("Total number of leaf nodes: %i \n", num_leaf_nodes)
     fmt.printf("Total number of triangles in BVH: %i \n", total_triangles_in_bvh)
   }
+  assert(cast(int)total_triangles_in_bvh == len(tri_indices))
   bvh.num_leaf_nodes = num_leaf_nodes
 }
\ No newline at end of file