pbrt/src/core/aggregates.rs

725 lines
24 KiB
Rust

use rayon::prelude::*;
use shared::core::aggregates::{BVHAggregate, LinearBVHNode, SplitMethod};
use shared::core::geometry::{Bounds3f, Point3f};
use shared::core::primitive::{Primitive, PrimitiveTrait};
use shared::utils::math::encode_morton_3;
use shared::utils::{find_interval, partition_slice};
use shared::{gvec, gvec_from_slice, Float};
use std::cmp::Ordering;
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
#[repr(C)]
#[derive(Debug, Default, Clone, Copy, PartialEq)]
struct BVHSplitBucket {
pub count: usize,
pub bounds: Bounds3f,
}
#[derive(Debug, Clone, Copy, Default)]
struct MortonPrimitive {
primitive_index: usize,
morton_code: u32,
}
struct LBVHTreelet {
start_index: usize,
n_primitives: usize,
}
#[derive(Debug, Clone)]
pub struct BVHPrimitiveInfo {
primitive_number: usize,
bounds: Bounds3f,
centroid: Point3f,
}
impl BVHPrimitiveInfo {
fn new(primitive_number: usize, bounds: Bounds3f) -> Self {
Self {
primitive_number,
bounds,
centroid: bounds.centroid(),
}
}
}
#[derive(Clone, Debug)]
pub enum BVHBuildNode {
Leaf {
n_primitives: usize,
bounds: Bounds3f,
primitive_indices: Vec<usize>,
},
Interior {
split_axis: u8,
children: [Box<BVHBuildNode>; 2],
bounds: Bounds3f,
},
}
impl Default for BVHBuildNode {
fn default() -> Self {
BVHBuildNode::Leaf {
n_primitives: 0,
bounds: Bounds3f::default(),
primitive_indices: Vec::new(),
}
}
}
impl BVHBuildNode {
pub fn new_leaf(n_primitives: usize, bounds: Bounds3f, indices: Vec<usize>) -> Self {
Self::Leaf {
bounds,
n_primitives,
primitive_indices: indices,
}
}
pub fn new_interior(axis: u8, c0: Box<BVHBuildNode>, c1: Box<BVHBuildNode>) -> Self {
let bounds = c0.bounds().union(c1.bounds());
Self::Interior {
bounds,
children: [c0, c1],
split_axis: axis,
}
}
pub fn bounds(&self) -> Bounds3f {
match self {
Self::Leaf { bounds, .. } => *bounds,
Self::Interior { bounds, .. } => *bounds,
}
}
pub fn split_axis(&self) -> Option<u8> {
match self {
Self::Interior { split_axis, .. } => Some(*split_axis),
_ => None,
}
}
}
pub trait CreateBVH {
fn new(primitives: Vec<Primitive>, max_prims_in_node: usize, split_method: SplitMethod) -> Self;
fn build_hlbvh(
bvh_primitives: &[BVHPrimitiveInfo],
total_nodes: &AtomicUsize,
_original_primitives: &[Primitive],
max_prims_in_node: usize,
) -> Box<BVHBuildNode>;
fn emit_lbvh(
bvh_primitives: &[BVHPrimitiveInfo],
morton_prims: &[MortonPrimitive],
total_nodes: &mut usize,
bit_index: i32,
max_prims_in_node: usize,
) -> Box<BVHBuildNode>;
fn build_upper_sah(nodes: &mut [BVHBuildNode], total_nodes: &AtomicUsize) -> Box<BVHBuildNode>;
}
impl CreateBVH for BVHAggregate {
fn new(
mut primitives: Vec<Primitive>,
max_prims_in_node: usize,
split_method: SplitMethod,
) -> Self {
let max_prims_in_node = std::cmp::min(255, max_prims_in_node);
if primitives.is_empty() {
return Self {
max_prims_in_node: max_prims_in_node.try_into().unwrap(),
node_count: 0,
primitives: gvec_from_slice(&primitives),
split_method,
nodes: gvec(),
};
}
let mut primitive_info: Vec<BVHPrimitiveInfo> = primitives
.iter()
.enumerate()
.map(|(i, p)| BVHPrimitiveInfo::new(i, p.bounds()))
.collect();
let node_count: usize;
let root: Box<BVHBuildNode>;
match split_method {
SplitMethod::Hlbvh => {
let nodes_counter = AtomicUsize::new(0);
root = Self::build_hlbvh(
&primitive_info,
&nodes_counter,
&primitives,
max_prims_in_node,
);
node_count = nodes_counter.load(AtomicOrdering::Relaxed);
}
_ => {
let nodes_counter = AtomicUsize::new(0);
root = build_recursive(
&mut primitive_info,
&nodes_counter,
&primitives,
max_prims_in_node,
split_method,
);
node_count = nodes_counter.load(AtomicOrdering::Relaxed);
}
};
// 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_vec = Vec::with_capacity(primitives.len());
leaf_order(&root, &mut leaf_vec);
reorder(&mut primitives, &leaf_vec);
drop(leaf_vec);
let mut nodes = vec![LinearBVHNode::default(); node_count];
let mut offset = 0;
let mut prim_offset = 0;
flatten(&root, &mut nodes, &mut offset, &mut prim_offset);
Self {
node_count: node_count.try_into().unwrap(),
max_prims_in_node: max_prims_in_node.try_into().unwrap(),
split_method,
primitives: gvec_from_slice(&primitives),
nodes: gvec_from_slice(&nodes),
}
}
fn build_hlbvh(
bvh_primitives: &[BVHPrimitiveInfo],
total_nodes: &AtomicUsize,
_original_primitives: &[Primitive],
max_prims_in_node: usize,
) -> Box<BVHBuildNode> {
let bounds = bvh_primitives
.iter()
.fold(Bounds3f::default(), |b, p| b.union(p.bounds));
let mut morton_prims: Vec<MortonPrimitive> = bvh_primitives
.par_iter()
.map(|prim| {
const MORTON_BITS: i32 = 10;
const MORTON_SCALE: i32 = 1 << MORTON_BITS;
let centroid_offset = bounds.offset(&prim.centroid);
let offset = centroid_offset * (MORTON_SCALE as Float);
MortonPrimitive {
primitive_index: prim.primitive_number,
morton_code: encode_morton_3(offset.x(), offset.y(), offset.z()),
}
})
.collect();
morton_prims.par_sort_unstable_by_key(|p| p.morton_code);
const TREELET_MASK: u32 = 0b00111111111111000000000000000000;
let mut split_indices: Vec<usize> = morton_prims
.par_windows(2) // Iterates over overlapping pairs [i, i+1]
.enumerate()
.filter_map(|(i, w)| {
let m1 = w[0].morton_code & TREELET_MASK;
let m2 = w[1].morton_code & TREELET_MASK;
// If mask changes, the split is at index i + 1
if m1 != m2 {
Some(i + 1)
} else {
None
}
})
.collect();
let mut boundaries = Vec::with_capacity(split_indices.len() + 2);
boundaries.push(0);
boundaries.append(&mut split_indices);
boundaries.push(morton_prims.len());
let treelets_to_build: Vec<LBVHTreelet> = boundaries
.windows(2)
.map(|w| LBVHTreelet {
start_index: w[0],
n_primitives: w[1] - w[0],
})
.collect();
let treelet_roots: Vec<Box<BVHBuildNode>> = treelets_to_build
.par_iter()
.map(|tr| {
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,
FIRST_BIT_INDEX,
max_prims_in_node,
);
total_nodes.fetch_add(nodes_created, AtomicOrdering::Relaxed);
root
})
.collect();
let mut contiguous_nodes: Vec<BVHBuildNode> = treelet_roots
.into_iter()
.map(|node_box| *node_box)
.collect();
Self::build_upper_sah(&mut contiguous_nodes, total_nodes)
}
fn emit_lbvh(
bvh_primitives: &[BVHPrimitiveInfo],
morton_prims: &[MortonPrimitive],
total_nodes: &mut usize,
bit_index: i32,
max_prims_in_node: usize,
) -> Box<BVHBuildNode> {
let n_primitives = morton_prims.len();
if bit_index == -1 || n_primitives <= max_prims_in_node {
*total_nodes += 1;
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(mp.primitive_index);
}
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| {
let current_code = morton_prims[index as usize].morton_code;
(current_code & mask) == (first_code & mask)
});
let split_offset = (last_match_index + 1) as usize;
if split_offset >= n_primitives {
return Self::emit_lbvh(
bvh_primitives,
morton_prims,
total_nodes,
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,
bit_index - 1,
max_prims_in_node,
);
let child1 = Self::emit_lbvh(
bvh_primitives,
right_morton,
total_nodes,
bit_index - 1,
max_prims_in_node,
);
let axis = (bit_index % 3) as u8;
Box::new(BVHBuildNode::new_interior(axis, child0, child1))
}
fn build_upper_sah(nodes: &mut [BVHBuildNode], total_nodes: &AtomicUsize) -> Box<BVHBuildNode> {
let n_nodes = nodes.len();
if n_nodes == 1 {
return Box::new(nodes[0].clone());
}
total_nodes.fetch_add(1, AtomicOrdering::Relaxed);
let bounds = nodes
.iter()
.fold(Bounds3f::default(), |b, node| b.union(node.bounds()));
let centroid_bounds = nodes.iter().fold(Bounds3f::default(), |b, node| {
b.union_point(node.bounds().centroid())
});
let dim = centroid_bounds.max_dimension();
if centroid_bounds.p_max[dim] == centroid_bounds.p_min[dim] {
let mid = n_nodes / 2;
let (left_part, right_part) = nodes.split_at_mut(mid);
return Box::new(BVHBuildNode::new_interior(
dim as u8,
Self::build_upper_sah(left_part, total_nodes),
Self::build_upper_sah(right_part, total_nodes),
));
}
const N_BUCKETS: usize = 12;
#[derive(Copy, Clone, Default)]
struct Bucket {
count: usize,
bounds: Bounds3f,
}
let mut buckets = [Bucket::default(); N_BUCKETS];
let get_bucket_idx = |node: &BVHBuildNode| -> usize {
let offset = centroid_bounds.offset(&node.bounds().centroid())[dim];
let mut b = (N_BUCKETS as Float * offset) as usize;
if b == N_BUCKETS {
b = N_BUCKETS - 1;
}
b
};
// Initialize _Bucket_ for HLBVH SAH partition buckets
for node in nodes.iter() {
let b = get_bucket_idx(node);
buckets[b].count += 1;
buckets[b].bounds = buckets[b].bounds.union(node.bounds());
}
// Compute costs for splitting after each bucket
let mut cost = [0.0; N_BUCKETS - 1];
// Forward Pass: Accumulate Left side (0 -> N-1)
let mut left_area = [0.0; N_BUCKETS];
let mut left_count = [0; N_BUCKETS];
let mut b_left = Bounds3f::default();
let mut c_left = 0;
for i in 0..N_BUCKETS {
b_left = b_left.union(buckets[i].bounds);
c_left += buckets[i].count;
left_area[i] = b_left.surface_area();
left_count[i] = c_left;
}
// Backward Pass: Accumulate Right side (N-1 -> 0) and compute cost
let mut b_right = Bounds3f::default();
let mut c_right = 0;
let inv_total_sa = 1.0 / bounds.surface_area();
for i in (0..N_BUCKETS - 1).rev() {
b_right = b_right.union(buckets[i + 1].bounds);
c_right += buckets[i + 1].count;
let count_left = left_count[i];
let sa_left = left_area[i];
let sa_right = b_right.surface_area();
cost[i] = 0.125
+ (count_left as Float * sa_left + c_right as Float * sa_right) * inv_total_sa;
}
// Find bucket to split at that minimizes SAH metric
let mut min_cost = cost[0];
let mut min_cost_split_bucket = 0;
for (i, &c) in cost.iter().enumerate().skip(1) {
if c < min_cost {
min_cost = c;
min_cost_split_bucket = i;
}
}
// Split nodes and create interior HLBVH SAH node
let mid = {
let mut left = 0;
for i in 0..n_nodes {
let b = get_bucket_idx(&nodes[i]);
if b <= min_cost_split_bucket {
nodes.swap(left, i);
left += 1;
}
}
left
};
if mid == 0 || mid == n_nodes {
let mid = n_nodes / 2;
// Partially sort so the median is in the middle and elements are partitioned around it
nodes.select_nth_unstable_by(mid, |a, b| {
a.bounds().centroid()[dim]
.partial_cmp(&b.bounds().centroid()[dim])
.unwrap_or(std::cmp::Ordering::Equal)
});
let (left_part, right_part) = nodes.split_at_mut(mid);
Box::new(BVHBuildNode::new_interior(
dim as u8,
Self::build_upper_sah(left_part, total_nodes),
Self::build_upper_sah(right_part, total_nodes),
))
} else {
// Standard SAH Split
let (left_part, right_part) = nodes.split_at_mut(mid);
Box::new(BVHBuildNode::new_interior(
dim as u8,
Self::build_upper_sah(left_part, total_nodes),
Self::build_upper_sah(right_part, total_nodes),
))
}
}
}
fn build_recursive(
bvh_primitives: &mut [BVHPrimitiveInfo],
total_nodes: &AtomicUsize,
original_primitives: &[Primitive],
max_prims_in_node: usize,
split_method: SplitMethod,
) -> Box<BVHBuildNode> {
total_nodes.fetch_add(1, AtomicOrdering::Relaxed);
let bounds = bvh_primitives
.iter()
.fold(Bounds3f::default(), |b, p| b.union(p.bounds));
let n_primitives = bvh_primitives.len();
if bounds.surface_area() == 0.0 || n_primitives == 1 || n_primitives <= max_prims_in_node {
let indices: Vec<usize> = 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 indices: Vec<usize> = bvh_primitives.iter().map(|p| p.primitive_number).collect();
return Box::new(BVHBuildNode::new_leaf(n_primitives, bounds, indices));
}
let mut mid: usize;
match split_method {
SplitMethod::Middle => {
let pmid = (centroid_bounds.p_min[dim] + centroid_bounds.p_max[dim]) / 2.;
mid = partition_slice(bvh_primitives, |p| p.centroid[dim] < pmid);
if mid != 0 && mid != n_primitives {
} else {
mid = n_primitives / 2;
bvh_primitives.select_nth_unstable_by(mid, |a, b| {
a.centroid[dim].partial_cmp(&b.centroid[dim]).unwrap()
});
}
}
SplitMethod::EqualCounts => {
mid = n_primitives / 2;
bvh_primitives.select_nth_unstable_by(mid, |a, b| {
a.centroid[dim].partial_cmp(&b.centroid[dim]).unwrap()
});
}
SplitMethod::SAH | _ => {
if n_primitives < 2 {
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 {
const N_BUCKETS: usize = 12;
let mut buckets = [BVHSplitBucket::default(); N_BUCKETS];
for prim in bvh_primitives.iter() {
let mut b =
(N_BUCKETS as Float * centroid_bounds.offset(&prim.centroid)[dim]) as usize;
if b == N_BUCKETS {
b = N_BUCKETS - 1;
}
buckets[b].count += 1;
buckets[b].bounds = buckets[b].bounds.union(prim.bounds);
}
// Compute costs for splitting after each bucket>
const N_SPLITS: usize = N_BUCKETS - 1;
let mut costs = [0.0 as Float; N_SPLITS];
let mut count_below = 0;
let mut bound_below = Bounds3f::default();
for i in 0..N_SPLITS {
bound_below = bound_below.union(buckets[i].bounds);
count_below += buckets[i].count;
costs[i] += count_below as Float * bound_below.surface_area();
}
// Finish initializing costs using a backward scan over splits
let mut count_above = 0;
let mut bound_above = Bounds3f::default();
for i in (0..N_SPLITS).rev() {
bound_above = bound_above.union(buckets[i + 1].bounds);
count_above += buckets[i + 1].count;
costs[i] += count_above as Float * bound_above.surface_area();
}
// Find bucket to split at that minimizes SAH metric>
let mut min_cost = Float::INFINITY;
let mut min_cost_split_bucket = 0;
for (i, &cost) in costs.iter().enumerate().take(N_SPLITS) {
if cost < min_cost {
min_cost = cost;
min_cost_split_bucket = i;
}
}
// Compute leaf cost and SAH split cost for chosen split
let leaf_cost = n_primitives as Float;
min_cost = 0.5 + min_cost / bounds.surface_area();
// Either create leaf or split primitives at selected SAH bucket>
if n_primitives > max_prims_in_node || min_cost < leaf_cost {
mid = partition_slice(bvh_primitives, |bp| {
let mut b = (N_BUCKETS as Float * centroid_bounds.offset(&bp.centroid)[dim])
as usize;
if b == N_BUCKETS {
b = N_BUCKETS - 1;
}
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 indices: Vec<usize> =
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);
let build_leaf = |prims: &mut [BVHPrimitiveInfo]| -> Box<BVHBuildNode> {
build_recursive(
prims,
total_nodes,
original_primitives,
max_prims_in_node,
split_method,
)
};
let (child0, child1) = if n_primitives > 128 * 1024 {
rayon::join(|| build_leaf(left_prims), || build_leaf(right_prims))
} else {
(build_leaf(left_prims), build_leaf(right_prims))
};
let axis = dim as u8;
Box::new(BVHBuildNode::new_interior(axis, child0, child1))
}
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 {
n_primitives,
bounds,
..
} => {
let n = *n_primitives;
let linear_node = &mut nodes[local_offset];
linear_node.bounds = *bounds;
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 {
split_axis,
children,
bounds,
} => {
nodes[local_offset].bounds = *bounds;
nodes[local_offset].axis = *split_axis;
nodes[local_offset].n_primitives = 0;
flatten(&children[0], nodes, offset, prim_offset);
let second_child_offset = flatten(&children[1], nodes, offset, prim_offset);
nodes[local_offset].primitives_offset = second_child_offset;
}
}
local_offset
}
fn reorder(primitives: &mut [Primitive], 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, .. } => {
leaf_order(&children[0], out);
leaf_order(&children[1], out);
}
}
}
// BVHAggregate *BVHAggregate::Create(std::vector<Primitive> prims,
// const ParameterDictionary &parameters) {
// std::string splitMethodName = parameters.GetOneString("splitmethod", "sah");
// BVHAggregate::SplitMethod splitMethod;
// if (splitMethodName == "sah")
// splitMethod = BVHAggregate::SplitMethod::SAH;
// else if (splitMethodName == "hlbvh")
// splitMethod = BVHAggregate::SplitMethod::HLBVH;
// else if (splitMethodName == "middle")
// splitMethod = BVHAggregate::SplitMethod::Middle;
// else if (splitMethodName == "equal")
// splitMethod = BVHAggregate::SplitMethod::EqualCounts;
// else {
// Warning(R"(BVH split method "%s" unknown. Using "sah".)", splitMethodName);
// splitMethod = BVHAggregate::SplitMethod::SAH;
// }
//
// int maxPrimsInNode = parameters.GetOneInt("maxnodeprims", 4);
// return new BVHAggregate(std::move(prims), maxPrimsInNode, splitMethod);
// }