Continuing, ever onwards

Cargo.toml

@@ -12,6 +12,7 @@ cuda = ["dep:cudarc", "dep:cust", "dep:cust_raw", "dep:cuda-runtime-sys"]
 vulkan = ["ash", "gpu-allocator"]
 ash = ["dep:ash"]
 gpu-allocator = ["dep:gpu-allocator"]
+jemalloc = ["jemallocator"]
 
 [dependencies]
 anyhow = "1.0.100"
@@ -52,12 +53,16 @@ cust = { git = "https://github.com/Rust-GPU/Rust-CUDA", branch = "main", default
 cust_raw = { git = "https://github.com/Rust-GPU/Rust-CUDA", branch = "main", default-features = false, optional = true }
 cuda-runtime-sys = { version = "0.3.0-alpha.1", optional = true}
 cudarc = { version = "0.18.2", features = ["cuda-13000"], optional = true }
+jemallocator = { version = "0.5", optional = true }
 
 [build-dependencies]
 spirv-builder = { git = "https://github.com/rust-gpu/rust-gpu", branch = "main", optional = true }
 cuda_builder = { git = "https://github.com/Rust-GPU/Rust-CUDA", branch = "main", optional = true }
 cc = "1.2.53"
 
+[dev-dependencies]
+sysinfo = "0.30"
+
 [workspace]
 members = ["shared"]
 exclude = ["crates/ptex-filter", "kernels"]

shared/src/core/primitive.rs

@@ -6,9 +6,10 @@ use crate::core::medium::{Medium, MediumInterface};
 use crate::core::pbrt::Float;
 use crate::core::shape::{Shape, ShapeIntersection, ShapeTrait};
 use crate::core::texture::{GPUFloatTexture, TextureEvalContext};
-use crate::utils::Ptr;
 use crate::utils::hash::hash_float;
 use crate::utils::transform::{AnimatedTransform, Transform};
+use crate::utils::Ptr;
+use alloc::boxed::Box;
 
 use enum_dispatch::enum_dispatch;
 
@@ -96,22 +97,24 @@ pub struct SimplePrimitive {
 
 impl PrimitiveTrait for SimplePrimitive {
     fn bounds(&self) -> Bounds3f {
-        todo!()
+        self.shape.bounds()
     }
 
-    fn intersect(&self, _r: &Ray, _t_max: Option<Float>) -> Option<ShapeIntersection> {
+    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
-        todo!()
+        let mut si = self.shape.intersect(r, t_max)?;
+        si.set_intersection_properties(self.material, Ptr::null(), MediumInterface::default(), r.medium);
+        Some(si)
     }
 
-    fn intersect_p(&self, _r: &Ray, _t_max: Option<Float>) -> bool {
+    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
-        todo!()
+        self.shape.intersect_p(r, t_max)
     }
 }
 
 #[derive(Debug, Clone)]
 pub struct TransformedPrimitive {
     pub primitive: Ptr<Primitive>,
-    pub render_from_primitive: Transform,
+    pub render_from_primitive: Ptr<Transform>,
 }
 
 impl PrimitiveTrait for TransformedPrimitive {
@@ -137,8 +140,9 @@ impl PrimitiveTrait for TransformedPrimitive {
         Some(si)
     }
 
-    fn intersect_p(&self, _r: &Ray, _t_max: Option<Float>) -> bool {
+    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
-        todo!()
+        let (ray, t_max) = self.render_from_primitive.apply_inverse_ray(r, t_max);
+        self.primitive.intersect_p(&ray, Some(t_max))
     }
 }
 
@@ -146,7 +150,7 @@ impl PrimitiveTrait for TransformedPrimitive {
 #[derive(Debug, Copy, Clone)]
 pub struct AnimatedPrimitive {
     primitive: Ptr<Primitive>,
-    render_from_primitive: AnimatedTransform,
+    render_from_primitive: Ptr<AnimatedTransform>,
 }
 
 impl PrimitiveTrait for AnimatedPrimitive {
@@ -243,3 +247,31 @@ pub enum Primitive {
     BVH(BVHAggregatePrimitive),
     KdTree(KdTreeAggregate),
 }
+
+// impl PrimitiveTrait for Box<TransformedPrimitive> {
+//     fn bounds(&self) -> Bounds3f {
+//         self.as_ref().bounds()
+//     }
+//
+//     fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
+//         self.as_ref().intersect(r, t_max)
+//     }
+//
+//     fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
+//         self.as_ref().intersect_p(r, t_max)
+//     }
+// }
+//
+// impl PrimitiveTrait for Box<AnimatedPrimitive> {
+//     fn bounds(&self) -> Bounds3f {
+//         self.as_ref().bounds()
+//     }
+//
+//     fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
+//         self.as_ref().intersect(r, t_max)
+//     }
+//
+//     fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
+//         self.as_ref().intersect_p(r, t_max)
+//     }
+// }

shared/src/lib.rs

@@ -1,6 +1,7 @@
 #![allow(unused_imports, dead_code)]
 #![feature(associated_type_defaults)]
 #![no_std]
+extern crate alloc;
 
 pub mod bxdfs;
 pub mod cameras;

src/core/aggregates.rs

@@ -66,9 +66,9 @@ impl BVHPrimitiveInfo {
 #[derive(Clone, Debug)]
 pub enum BVHBuildNode {
     Leaf {
-        first_prim_offset: usize,
         n_primitives: usize,
         bounds: Bounds3f,
+        primitive_indices: Vec<usize>,
     },
     Interior {
         split_axis: u8,
@@ -80,19 +80,19 @@ pub enum BVHBuildNode {
 impl Default for BVHBuildNode {
     fn default() -> Self {
         BVHBuildNode::Leaf {
-            first_prim_offset: 0,
             n_primitives: 0,
             bounds: Bounds3f::default(),
+            primitive_indices: Vec::new(),
         }
     }
 }
 
 impl BVHBuildNode {
-    pub fn new_leaf(first_prim_offset: usize, n_primitives: usize, bounds: Bounds3f) -> Self {
+    pub fn new_leaf(n_primitives: usize, bounds: Bounds3f, indices: Vec<usize>) -> Self {
         Self::Leaf {
             bounds,
-            first_prim_offset,
             n_primitives,
+            primitive_indices: indices,
         }
     }
 
@@ -186,50 +186,44 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
             .map(|(i, p)| BVHPrimitiveInfo::new(i, p.bounds()))
             .collect();
 
-        let ordered_prims: Vec<P>;
         let total_nodes_count: usize;
         let root: Box<BVHBuildNode>;
 
         match split_method {
             SplitMethod::Hlbvh => {
                 let nodes_counter = AtomicUsize::new(0);
-                let ordered_prims_offset = AtomicUsize::new(0);
+                root = Self::build_hlbvh(
-                let mut local_ordered = vec![primitives[0].clone(); primitives.len()];
+                    &primitive_info,
-                let shared_buffer =
+                    &nodes_counter,
-                    SharedPrimitiveBuffer::new(&mut local_ordered, &ordered_prims_offset);
+                    &primitives,
-
+                    max_prims_in_node,
-                root =
+                );
-                    Self::build_hlbvh(&primitive_info, &nodes_counter, &shared_buffer, &primitives);
-
-                ordered_prims = local_ordered;
                 total_nodes_count = nodes_counter.load(AtomicOrdering::Relaxed);
             }
             _ => {
                 let nodes_counter = AtomicUsize::new(0);
-                let ordered_prims_offset = AtomicUsize::new(0);
-                let mut local_ordered = vec![primitives[0].clone(); primitives.len()];
-                let shared_buffer =
-                    SharedPrimitiveBuffer::new(&mut local_ordered, &ordered_prims_offset);
-
                 root = Self::build_recursive(
                     &mut primitive_info,
                     &nodes_counter,
-                    &shared_buffer,
                     &primitives,
                     max_prims_in_node,
                     split_method,
                 );
-
-                ordered_prims = local_ordered;
                 total_nodes_count = nodes_counter.load(AtomicOrdering::Relaxed);
             }
         };
 
-        primitives = ordered_prims;
+        // Walk the tree and collect primitive indices in the exact order
+        // the linear layout will visit them (left-to-right, depth-first)
+        let mut leaf_order = Vec::with_capacity(primitives.len());
+        Self::leaf_order(&root, &mut leaf_order);
+        Self::reorder(&mut primitives, &leaf_order);
+        drop(leaf_order);
 
         let mut nodes = vec![LinearBVHNode::default(); total_nodes_count];
         let mut offset = 0;
-        Self::flatten_bvh(&root, &mut nodes, &mut offset);
+        let mut prim_offset = 0;
+        Self::flatten(&root, &mut nodes, &mut offset, &mut prim_offset);
 
         Self {
             max_prims_in_node,
@@ -239,21 +233,65 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
         }
     }
 
-    fn flatten_bvh(node: &BVHBuildNode, nodes: &mut [LinearBVHNode], offset: &mut usize) -> usize {
+    fn reorder(primitives: &mut [P], order: &[usize]) {
+        let n = primitives.len();
+        assert_eq!(n, order.len());
+
+        let mut done = vec![false; n];
+        for i in 0..n {
+            if done[i] || order[i] == i {
+                done[i] = true;
+                continue;
+            }
+
+            let mut prev = i;
+            let mut curr = order[i];
+            while curr != i {
+                primitives.swap(prev, curr);
+                done[prev] = true;
+                prev = curr;
+                curr = order[prev];
+            }
+            done[prev] = true;
+        }
+    }
+
+    fn leaf_order(node: &BVHBuildNode, out: &mut Vec<usize>) {
+        match node {
+            BVHBuildNode::Leaf {
+                primitive_indices, ..
+            } => {
+                out.extend_from_slice(primitive_indices);
+            }
+            BVHBuildNode::Interior { children, .. } => {
+                Self::leaf_order(&children[0], out);
+                Self::leaf_order(&children[1], out);
+            }
+        }
+    }
+
+    fn flatten(
+        node: &BVHBuildNode,
+        nodes: &mut [LinearBVHNode],
+        offset: &mut usize,
+        prim_offset: &mut usize,
+    ) -> usize {
         let local_offset = *offset;
         *offset += 1;
 
         match node {
             BVHBuildNode::Leaf {
-                first_prim_offset,
                 n_primitives,
                 bounds,
+                ..
             } => {
+                let n = *n_primitives;
                 let linear_node = &mut nodes[local_offset];
                 linear_node.bounds = *bounds;
-                linear_node.n_primitives = *n_primitives as u16;
+                linear_node.n_primitives = n as u16;
-                linear_node.primitives_offset = *first_prim_offset;
+                linear_node.primitives_offset = *prim_offset;
                 linear_node.axis = 0; // Irrelevant for leaves
+                *prim_offset += n;
             }
 
             BVHBuildNode::Interior {
@@ -265,8 +303,8 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
                 nodes[local_offset].axis = *split_axis;
                 nodes[local_offset].n_primitives = 0;
 
-                Self::flatten_bvh(&children[0], nodes, offset);
+                Self::flatten(&children[0], nodes, offset, prim_offset);
-                let second_child_offset = Self::flatten_bvh(&children[1], nodes, offset);
+                let second_child_offset = Self::flatten(&children[1], nodes, offset, prim_offset);
                 nodes[local_offset].primitives_offset = second_child_offset;
             }
         }
@@ -277,8 +315,8 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
     pub fn build_hlbvh(
         bvh_primitives: &[BVHPrimitiveInfo],
         total_nodes: &AtomicUsize,
-        ordered_prims: &SharedPrimitiveBuffer,
+        original_primitives: &[P],
-        original_primitives: &[Arc<dyn PrimitiveTrait>],
+        max_prims_in_node: usize,
     ) -> Box<BVHBuildNode> {
         let bounds = bvh_primitives
             .iter()
@@ -337,15 +375,12 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
                 let mut nodes_created = 0;
 
                 const FIRST_BIT_INDEX: i32 = 29 - 12;
-
                 let root = Self::emit_lbvh(
                     bvh_primitives,
                     &morton_prims[tr.start_index..tr.start_index + tr.n_primitives],
                     &mut nodes_created,
-                    ordered_prims,
-                    original_primitives,
                     FIRST_BIT_INDEX,
-                    4,
+                    max_prims_in_node,
                 );
 
                 total_nodes.fetch_add(nodes_created, AtomicOrdering::Relaxed);
@@ -366,8 +401,6 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
         bvh_primitives: &[BVHPrimitiveInfo],
         morton_prims: &[MortonPrimitive],
         total_nodes: &mut usize,
-        ordered_prims: &SharedPrimitiveBuffer,
-        original_primitives: &[Arc<dyn PrimitiveTrait>],
         bit_index: i32,
         max_prims_in_node: usize,
     ) -> Box<BVHBuildNode> {
@@ -375,23 +408,18 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
         if bit_index == -1 || n_primitives <= max_prims_in_node {
             *total_nodes += 1;
 
-            // Calculate bounds while collecting indices
             let mut bounds = Bounds3f::default();
             let mut indices = Vec::with_capacity(n_primitives);
 
             for mp in morton_prims {
                 let info = &bvh_primitives[mp.primitive_index];
                 bounds = bounds.union(info.bounds);
-                indices.push(info.clone());
+                indices.push(mp.primitive_index);
             }
 
-            let first_prim_offset = ordered_prims.append(original_primitives, &indices);
+            return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices));
-            return Box::new(BVHBuildNode::new_leaf(
-                first_prim_offset,
-                n_primitives,
-                bounds,
-            ));
         }
+
         let mask = 1 << bit_index;
         let first_code = morton_prims[0].morton_code;
         let last_match_index = find_interval(n_primitives.try_into().unwrap(), |index| {
@@ -405,23 +433,18 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
                 bvh_primitives,
                 morton_prims,
                 total_nodes,
-                ordered_prims,
-                original_primitives,
                 bit_index - 1,
                 max_prims_in_node,
             );
         }
 
         let (left_morton, right_morton) = morton_prims.split_at(split_offset);
-
         *total_nodes += 1;
 
         let child0 = Self::emit_lbvh(
             bvh_primitives,
             left_morton,
             total_nodes,
-            ordered_prims,
-            original_primitives,
             bit_index - 1,
             max_prims_in_node,
         );
@@ -430,8 +453,6 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
             bvh_primitives,
             right_morton,
             total_nodes,
-            ordered_prims,
-            original_primitives,
             bit_index - 1,
             max_prims_in_node,
         );
@@ -578,8 +599,7 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
     fn build_recursive(
         bvh_primitives: &mut [BVHPrimitiveInfo],
         total_nodes: &AtomicUsize,
-        ordered_prims: &SharedPrimitiveBuffer,
+        original_primitives: &[P],
-        original_primitives: &[Arc<dyn PrimitiveTrait>],
         max_prims_in_node: usize,
         split_method: SplitMethod,
     ) -> Box<BVHBuildNode> {
@@ -590,28 +610,17 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
 
         let n_primitives = bvh_primitives.len();
         if bounds.surface_area() == 0.0 || n_primitives == 1 || n_primitives <= max_prims_in_node {
-            let first_prim_offset = ordered_prims.append(original_primitives, bvh_primitives);
+            let indices: Vec<usize> = bvh_primitives.iter().map(|p| p.primitive_number).collect();
-
+            return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices));
-            return Box::new(BVHBuildNode::new_leaf(
-                first_prim_offset,
-                n_primitives,
-                bounds,
-            ));
         }
-
         let centroid_bounds = bvh_primitives.iter().fold(Bounds3f::default(), |b, p| {
             b.union_point(p.bounds.centroid())
         });
 
         let dim = centroid_bounds.max_dimension();
         if centroid_bounds.p_max[dim] == centroid_bounds.p_min[dim] {
-            let first_prim_offset = ordered_prims.append(original_primitives, bvh_primitives);
+            let indices: Vec<usize> = bvh_primitives.iter().map(|p| p.primitive_number).collect();
-
+            return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices));
-            return Box::new(BVHBuildNode::new_leaf(
-                first_prim_offset,
-                n_primitives,
-                bounds,
-            ));
         }
 
         let mut mid: usize;
@@ -699,68 +708,42 @@ impl<P: PrimitiveTrait + Clone + Send + Sync> BVHAggregate<P> {
                             }
                             b <= min_cost_split_bucket
                         });
+                        if mid == 0 || mid == n_primitives {
+                            mid = n_primitives / 2;
+                            bvh_primitives.select_nth_unstable_by(mid, |a, b| {
+                                a.centroid[dim]
+                                    .partial_cmp(&b.centroid[dim])
+                                    .unwrap_or(Ordering::Equal)
+                            });
+                        }
                     } else {
-                        let first_prim_offset =
+                        let indices: Vec<usize> =
-                            ordered_prims.append(original_primitives, bvh_primitives);
+                            bvh_primitives.iter().map(|p| p.primitive_number).collect();
-                        return Box::new(BVHBuildNode::new_leaf(
+                        return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices));
-                            first_prim_offset,
-                            n_primitives,
-                            bounds,
-                        ));
                     }
                 }
             }
         };
 
         let (left_prims, right_prims) = bvh_primitives.split_at_mut(mid);
-        if n_primitives > 128 * 1024 {
+        let build_leaf = |prims: &mut [BVHPrimitiveInfo]| -> Box<BVHBuildNode> {
-            let (child0, child1) = rayon::join(
+            Self::build_recursive(
-                || {
+                prims,
-                    Self::build_recursive(
+                total_nodes,
-                        left_prims,
+                original_primitives,
-                        total_nodes,
+                max_prims_in_node,
-                        ordered_prims,
+                split_method,
-                        original_primitives,
+            )
-                        max_prims_in_node,
+        };
-                        split_method,
-                    )
-                },
-                || {
-                    Self::build_recursive(
-                        right_prims,
-                        total_nodes,
-                        ordered_prims,
-                        original_primitives,
-                        max_prims_in_node,
-                        split_method,
-                    )
-                },
-            );
 
-            let axis = dim as u8;
+        let (child0, child1) = if n_primitives > 128 * 1024 {
-            Box::new(BVHBuildNode::new_interior(axis, child0, child1))
+            rayon::join(|| build_leaf(left_prims), || build_leaf(right_prims))
         } else {
-            let child0 = Self::build_recursive(
+            (build_leaf(left_prims), build_leaf(right_prims))
-                left_prims,
+        };
-                total_nodes,
-                ordered_prims,
-                original_primitives,
-                max_prims_in_node,
-                split_method,
-            );
 
-            let child1 = Self::build_recursive(
+        let axis = dim as u8;
-                right_prims,
+        Box::new(BVHBuildNode::new_interior(axis, child0, child1))
-                total_nodes,
-                ordered_prims,
-                original_primitives,
-                max_prims_in_node,
-                split_method,
-            );
-
-            let axis = dim as u8;
-            Box::new(BVHBuildNode::new_interior(axis, child0, child1))
-        }
     }
 
     pub fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {