diff --git a/Cargo.toml b/Cargo.toml index 91bd11c..da79e13 100644 --- a/Cargo.toml +++ b/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"] diff --git a/shared/src/core/primitive.rs b/shared/src/core/primitive.rs index 2325848..ae8b908 100644 --- a/shared/src/core/primitive.rs +++ b/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) -> Option { - todo!() + fn intersect(&self, r: &Ray, t_max: Option) -> Option { + 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) -> bool { - todo!() + fn intersect_p(&self, r: &Ray, t_max: Option) -> bool { + self.shape.intersect_p(r, t_max) } } #[derive(Debug, Clone)] pub struct TransformedPrimitive { pub primitive: Ptr, - pub render_from_primitive: Transform, + pub render_from_primitive: Ptr, } impl PrimitiveTrait for TransformedPrimitive { @@ -137,8 +140,9 @@ impl PrimitiveTrait for TransformedPrimitive { Some(si) } - fn intersect_p(&self, _r: &Ray, _t_max: Option) -> bool { - todo!() + fn intersect_p(&self, r: &Ray, t_max: Option) -> bool { + 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, - render_from_primitive: AnimatedTransform, + render_from_primitive: Ptr, } impl PrimitiveTrait for AnimatedPrimitive { @@ -243,3 +247,31 @@ pub enum Primitive { BVH(BVHAggregatePrimitive), KdTree(KdTreeAggregate), } + +// impl PrimitiveTrait for Box { +// fn bounds(&self) -> Bounds3f { +// self.as_ref().bounds() +// } +// +// fn intersect(&self, r: &Ray, t_max: Option) -> Option { +// self.as_ref().intersect(r, t_max) +// } +// +// fn intersect_p(&self, r: &Ray, t_max: Option) -> bool { +// self.as_ref().intersect_p(r, t_max) +// } +// } +// +// impl PrimitiveTrait for Box { +// fn bounds(&self) -> Bounds3f { +// self.as_ref().bounds() +// } +// +// fn intersect(&self, r: &Ray, t_max: Option) -> Option { +// self.as_ref().intersect(r, t_max) +// } +// +// fn intersect_p(&self, r: &Ray, t_max: Option) -> bool { +// self.as_ref().intersect_p(r, t_max) +// } +// } diff --git a/shared/src/lib.rs b/shared/src/lib.rs index e4fdadc..d36af1a 100644 --- a/shared/src/lib.rs +++ b/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; diff --git a/src/core/aggregates.rs b/src/core/aggregates.rs index cd3bb53..2277cb5 100644 --- a/src/core/aggregates.rs +++ b/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, }, 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) -> Self { Self::Leaf { bounds, - first_prim_offset, n_primitives, + primitive_indices: indices, } } @@ -186,50 +186,44 @@ impl BVHAggregate

{ .map(|(i, p)| BVHPrimitiveInfo::new(i, p.bounds())) .collect(); - let ordered_prims: Vec

; let total_nodes_count: usize; let root: Box; match split_method { SplitMethod::Hlbvh => { 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_hlbvh(&primitive_info, &nodes_counter, &shared_buffer, &primitives); - - ordered_prims = local_ordered; + root = Self::build_hlbvh( + &primitive_info, + &nodes_counter, + &primitives, + max_prims_in_node, + ); 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 BVHAggregate

{ } } - 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) { + 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.primitives_offset = *first_prim_offset; + linear_node.n_primitives = n as u16; + linear_node.primitives_offset = *prim_offset; linear_node.axis = 0; // Irrelevant for leaves + *prim_offset += n; } BVHBuildNode::Interior { @@ -265,8 +303,8 @@ impl BVHAggregate

{ nodes[local_offset].axis = *split_axis; nodes[local_offset].n_primitives = 0; - Self::flatten_bvh(&children[0], nodes, offset); - let second_child_offset = Self::flatten_bvh(&children[1], nodes, offset); + Self::flatten(&children[0], nodes, offset, prim_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 BVHAggregate

{ pub fn build_hlbvh( bvh_primitives: &[BVHPrimitiveInfo], total_nodes: &AtomicUsize, - ordered_prims: &SharedPrimitiveBuffer, - original_primitives: &[Arc], + original_primitives: &[P], + max_prims_in_node: usize, ) -> Box { let bounds = bvh_primitives .iter() @@ -337,15 +375,12 @@ impl BVHAggregate

{ 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 BVHAggregate

{ bvh_primitives: &[BVHPrimitiveInfo], morton_prims: &[MortonPrimitive], total_nodes: &mut usize, - ordered_prims: &SharedPrimitiveBuffer, - original_primitives: &[Arc], bit_index: i32, max_prims_in_node: usize, ) -> Box { @@ -375,23 +408,18 @@ impl BVHAggregate

{ 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( - first_prim_offset, - n_primitives, - bounds, - )); + return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices)); } + 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 BVHAggregate

{ 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 BVHAggregate

{ bvh_primitives, right_morton, total_nodes, - ordered_prims, - original_primitives, bit_index - 1, max_prims_in_node, ); @@ -578,8 +599,7 @@ impl BVHAggregate

{ fn build_recursive( bvh_primitives: &mut [BVHPrimitiveInfo], total_nodes: &AtomicUsize, - ordered_prims: &SharedPrimitiveBuffer, - original_primitives: &[Arc], + original_primitives: &[P], max_prims_in_node: usize, split_method: SplitMethod, ) -> Box { @@ -590,28 +610,17 @@ impl BVHAggregate

{ 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); - - return Box::new(BVHBuildNode::new_leaf( - first_prim_offset, - n_primitives, - bounds, - )); + let indices: Vec = bvh_primitives.iter().map(|p| p.primitive_number).collect(); + return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices)); } - 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); - - return Box::new(BVHBuildNode::new_leaf( - first_prim_offset, - n_primitives, - bounds, - )); + let indices: Vec = bvh_primitives.iter().map(|p| p.primitive_number).collect(); + return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices)); } let mut mid: usize; @@ -699,68 +708,42 @@ impl BVHAggregate

{ } 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 = - ordered_prims.append(original_primitives, bvh_primitives); - return Box::new(BVHBuildNode::new_leaf( - first_prim_offset, - n_primitives, - bounds, - )); + let indices: Vec = + bvh_primitives.iter().map(|p| p.primitive_number).collect(); + return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices)); } } } }; let (left_prims, right_prims) = bvh_primitives.split_at_mut(mid); - if n_primitives > 128 * 1024 { - let (child0, child1) = rayon::join( - || { - Self::build_recursive( - left_prims, - total_nodes, - ordered_prims, - 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 build_leaf = |prims: &mut [BVHPrimitiveInfo]| -> Box { + Self::build_recursive( + prims, + total_nodes, + original_primitives, + max_prims_in_node, + split_method, + ) + }; - let axis = dim as u8; - Box::new(BVHBuildNode::new_interior(axis, child0, child1)) + let (child0, child1) = if n_primitives > 128 * 1024 { + rayon::join(|| build_leaf(left_prims), || build_leaf(right_prims)) } else { - let child0 = Self::build_recursive( - left_prims, - total_nodes, - ordered_prims, - original_primitives, - max_prims_in_node, - split_method, - ); + (build_leaf(left_prims), build_leaf(right_prims)) + }; - let child1 = Self::build_recursive( - right_prims, - 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)) - } + let axis = dim as u8; + Box::new(BVHBuildNode::new_interior(axis, child0, child1)) } pub fn intersect(&self, r: &Ray, t_max: Option) -> Option {