rtiow: BVHTriangles use SAH for division and leave original triangles untouched.

rtiow/renderer/src/bvh_triangles.rs

@@ -3,6 +3,7 @@
 use std::f32::EPSILON;
 use std::fmt;
 
+use log::info;
 use stl::STL;
 
 use crate::{
@@ -16,15 +17,15 @@ use crate::{
 #[derive(Debug, PartialEq)]
 struct BVHNode {
     aabb: AABB,
-    // When prim_count==0, left_first holds the left child's index in bvh_nodes. When >0 left_first
+    // When tri_count==0, left_first holds the left child's index in bvh_nodes. When >0 left_first
     // holds the index for the first triangle in triangles.
     left_first: u32,
-    prim_count: u32,
+    tri_count: u32,
 }
 
 impl BVHNode {
     fn is_leaf(&self) -> bool {
-        self.prim_count > 0
+        self.tri_count > 0
     }
 }
 
@@ -48,6 +49,7 @@ where
     M: Material,
 {
     pub triangles: Vec<Triangle>,
+    triangle_index: Vec<usize>,
     material: M,
     bvh_nodes: Vec<BVHNode>,
 }
@@ -80,7 +82,7 @@ where
             }
             let n = &self.bvh_nodes[i];
             if n.is_leaf() {
-                for t_idx in n.left_first..(n.left_first + n.prim_count) {
+                for t_idx in n.left_first..(n.left_first + n.tri_count) {
                     if f.alternate() {
                         write!(f, "\t")?;
                     }
@@ -101,6 +103,8 @@ where
     M: Material,
 {
     pub fn new(stl: &STL, material: M) -> BVHTriangles<M> {
+        let now = std::time::Instant::now();
+
         assert_eq!(std::mem::size_of::<BVHNode>(), 32);
         let div3 = 1. / 3.;
         let triangles: Vec<_> = stl
@@ -117,15 +121,21 @@ where
                 }
             })
             .collect();
+        let triangle_index = (0..triangles.len()).collect();
 
         let n = 2 * triangles.len() - 2;
         let bvh_nodes = Vec::with_capacity(n);
         let mut bvh = BVHTriangles {
             triangles,
+            triangle_index,
             bvh_nodes,
             material,
         };
         bvh.build_bvh();
+        info!(
+            "BVHTriangles build time {:0.3}s",
+            now.elapsed().as_secs_f32()
+        );
         bvh
     }
 
@@ -134,7 +144,7 @@ where
         let root = BVHNode {
             aabb: AABB::default(),
             left_first: 0,
-            prim_count: self.triangles.len() as u32,
+            tri_count: self.triangles.len() as u32,
         };
         self.bvh_nodes.push(root);
         self.update_node_bounds(ROOT_NODE_IDX);
@@ -145,8 +155,9 @@ where
         let node = &mut self.bvh_nodes[node_idx];
         let mut aabb_min: Vec3 = f32::MAX.into();
         let mut aabb_max: Vec3 = f32::MIN.into();
-        for i in node.left_first..(node.left_first + node.prim_count) {
+        for i in node.left_first..(node.left_first + node.tri_count) {
-            let leaf_tri = &self.triangles[i as usize];
+            let leaf_tri_idx = self.triangle_index[i as usize];
+            let leaf_tri = &self.triangles[leaf_tri_idx];
             aabb_min = vec3::min(aabb_min, leaf_tri.verts[0]);
             aabb_min = vec3::min(aabb_min, leaf_tri.verts[1]);
             aabb_min = vec3::min(aabb_min, leaf_tri.verts[2]);
@@ -157,36 +168,59 @@ where
         node.aabb = AABB::new(aabb_min, aabb_max);
     }
     fn subdivide(&mut self, idx: usize) {
-        // Early out if we're down to just 2 or less triangles
+        let node = &self.bvh_nodes[idx];
-        if self.bvh_nodes[idx].prim_count <= 2 {
+        let parent_area = node.aabb.area();
-            return;
+        let parent_cost = node.tri_count as f32 * parent_area;
-        }
+
-        let (left_first, prim_count, left_count, i) = {
+        let (left_first, tri_count, left_count, i) = {
             let node = &self.bvh_nodes[idx];
 
-            // Compute split plane and position.
+            let mut best_axis = usize::MAX;
-            let extent = node.aabb.max() - node.aabb.min();
+            let mut best_pos = 0.;
-            let axis = node.aabb.longest_axis();
+            let mut best_cost = f32::MAX;
-            let split_pos = node.aabb.min()[axis] + extent[axis] * 0.5;
+
+            for axis in 0..3 {
+                for i in 0..node.tri_count {
+                    let triangle =
+                        &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
+                    let candidate_pos = triangle.centroid[axis];
+                    let cost = self.evaluate_sah(node, axis, candidate_pos);
+                    if cost <= best_cost {
+                        best_pos = candidate_pos;
+                        best_axis = axis;
+                        best_cost = cost;
+                    }
+                }
+            }
+            // Stop subdividing if it isn't getting any better.
+            if best_cost <= parent_cost {
+                return;
+            }
+
+            let axis = best_axis;
+            let split_pos = best_pos;
 
             // Split the group in two halves.
             let mut i = node.left_first as isize;
-            let mut j = i + node.prim_count as isize - 1;
+            let mut j = i + node.tri_count as isize - 1;
             while i <= j {
-                if self.triangles[i as usize].centroid[axis] < split_pos {
+                if self.triangles[self.triangle_index[i as usize]].centroid[axis] < split_pos {
                     i += 1;
                 } else {
-                    self.triangles.swap(i as usize, j as usize);
+                    self.triangles.swap(
+                        self.triangle_index[i as usize],
+                        self.triangle_index[j as usize],
+                    );
                     j -= 1;
                 }
             }
 
             // Create child nodes for each half.
             let left_count = i as u32 - node.left_first;
-            if left_count == 0 || left_count == node.prim_count {
+            if left_count == 0 || left_count == node.tri_count {
                 return;
             }
-            (node.left_first, node.prim_count, left_count, i)
+            (node.left_first, node.tri_count, left_count, i)
         };
 
         // create child nodes
@@ -195,27 +229,56 @@ where
         let left = BVHNode {
             aabb: AABB::default(),
             left_first,
-            prim_count: left_count as u32,
+            tri_count: left_count as u32,
         };
         let right = BVHNode {
             aabb: AABB::default(),
             left_first: i as u32,
-            prim_count: (prim_count - left_count) as u32,
+            tri_count: (tri_count - left_count) as u32,
         };
         self.bvh_nodes.push(left);
         self.bvh_nodes.push(right);
         let node = &mut self.bvh_nodes[idx];
         node.left_first = left_child_idx;
-        node.prim_count = 0;
+        node.tri_count = 0;
 
         // Recurse
         self.update_node_bounds(left_child_idx as usize);
-        self.update_node_bounds(right_child_idx as usize);
         self.subdivide(left_child_idx as usize);
+        self.update_node_bounds(right_child_idx as usize);
         self.subdivide(right_child_idx as usize);
     }
 
+    fn evaluate_sah(&self, node: &BVHNode, axis: usize, pos: f32) -> f32 {
+        // determine triangle counts and bounds for this split candidate
+        let mut left_box = AABB::infinite();
+        let mut right_box = AABB::infinite();
+        let mut left_count = 0;
+        let mut right_count = 0;
+        for i in 0..node.tri_count {
+            let triangle = &self.triangles[self.triangle_index[(node.left_first + i) as usize]];
+            if triangle.centroid[axis] < pos {
+                left_count += 1;
+                left_box.grow(triangle.verts[0]);
+                left_box.grow(triangle.verts[1]);
+                left_box.grow(triangle.verts[2]);
+            } else {
+                right_count += 1;
+                right_box.grow(triangle.verts[0]);
+                right_box.grow(triangle.verts[1]);
+                right_box.grow(triangle.verts[2]);
+            }
+        }
+        let cost = left_count as f32 * left_box.area() + right_count as f32 * right_box.area();
+        if cost > 0. {
+            cost
+        } else {
+            f32::MAX
+        }
+    }
+
     fn intersect_bvh(&self, r: Ray, node_idx: u32, t_min: f32, t_max: f32) -> Option<HitRecord> {
+        // TODO(wathiede): visit nodes front to back and remove recursion.
         let node = &self.bvh_nodes[node_idx as usize];
         if !node.aabb.hit(r, t_min, t_max) {
             return None;
@@ -367,58 +430,58 @@ mod tests {
     use stl::STL;
 
     /*
-        #[test]
+           #[test]
-        fn build_bvh() {
+           fn build_bvh() {
-            let stl_triangles: Vec<_> = (0..4)
+           let stl_triangles: Vec<_> = (0..4)
-                .flat_map(|y| {
+           .flat_map(|y| {
-                    (0..2).map(move |x| {
+           (0..2).map(move |x| {
-                        let x = x as f32;
+           let x = x as f32;
-                        let y = y as f32;
+           let y = y as f32;
-                        stl::Triangle {
+           stl::Triangle {
-                            normal: [1., 0., 0.].into(),
+           normal: [1., 0., 0.].into(),
-                            verts: [
+           verts: [
-                                [2. * x + 0., 2. * y + 0., 0.].into(),
+           [2. * x + 0., 2. * y + 0., 0.].into(),
-                                [2. * x + 1., 2. * y + 0., 0.].into(),
+           [2. * x + 1., 2. * y + 0., 0.].into(),
-                                [2. * x + 1., 2. * y + 1., 0.].into(),
+           [2. * x + 1., 2. * y + 1., 0.].into(),
-                            ],
+           ],
-                            attr: 0,
+           attr: 0,
-                        }
+           }
-                    })
+           })
-                })
+           })
-                .collect();
+           .collect();
-            let stl = STL {
+           let stl = STL {
-                header: [0; 80],
+           header: [0; 80],
-                triangles: stl_triangles,
+           triangles: stl_triangles,
-            };
+           };
-            /*
+        /*
-            let mut bvh_triangles: Vec<_> = stl_triangles
+        let mut bvh_triangles: Vec<_> = stl_triangles
-            .iter()
+        .iter()
-            .map(|tri| {
+        .map(|tri| {
-            let div3 = 1. / 3.;
+        let div3 = 1. / 3.;
-            let v0 = tri.verts[0];
+        let v0 = tri.verts[0];
-            let v1 = tri.verts[1];
+        let v1 = tri.verts[1];
-            let v2 = tri.verts[2];
+        let v2 = tri.verts[2];
-            let centroid = (v0 + v1 + v2) * div3;
+        let centroid = (v0 + v1 + v2) * div3;
-            Triangle {
+        Triangle {
-            centroid,
+        centroid,
-            verts: tri.verts,
+        verts: tri.verts,
-            }
-            })
-            .collect();
-            bvh_triangles.sort_by(|a, b| a.centroid.y.partial_cmp(&b.centroid.y).unwrap());
-            let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
-            let bvh_nodes = Default::default();
-            let want = BVHTriangles {
-            triangles: bvh_triangles,
-            bvh_nodes,
-            material,
-            };
-            */
-            let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
-            let bvh = BVHTriangles::new(&stl, material);
-            dbg!(&bvh);
-            assert_eq!(bvh.bvh_nodes.len(), 2 * bvh.triangles.len() - 2);
         }
+        })
+        .collect();
+        bvh_triangles.sort_by(|a, b| a.centroid.y.partial_cmp(&b.centroid.y).unwrap());
+        let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
+        let bvh_nodes = Default::default();
+        let want = BVHTriangles {
+        triangles: bvh_triangles,
+        bvh_nodes,
+        material,
+        };
+        */
+        let material = Lambertian::new(ConstantTexture::new([0., 0., 0.]));
+        let bvh = BVHTriangles::new(&stl, material);
+        dbg!(&bvh);
+        assert_eq!(bvh.bvh_nodes.len(), 2 * bvh.triangles.len() - 2);
+    }
     */
 
     #[test]