diff --git a/shared/src/core/aggregates.rs b/shared/src/core/aggregates.rs new file mode 100644 index 0000000..94f3154 --- /dev/null +++ b/shared/src/core/aggregates.rs @@ -0,0 +1,189 @@ +use crate::core::geometry::{Bounds3f, Ray, Vector3f}; +use crate::core::pbrt::Float; +use crate::core::primitive::{Primitive, PrimitiveTrait}; +use crate::core::shape::ShapeIntersection; +use crate::utils::Ptr; + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct LinearBVHNode { + pub bounds: Bounds3f, + pub primitives_offset: usize, + pub n_primitives: u16, + pub axis: u8, + pub pad: u8, +} + +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct DeviceBVHAggregate { + pub max_prims_in_node: u32, + pub primitives: Ptr, + pub primitive_count: u32, + pub nodes: Ptr, + pub node_count: u32, +} + +impl DeviceBVHAggregate { + pub const fn empty() -> Self { + Self { + max_prims_in_node: 0, + primitives: Ptr::null(), + primitive_count: 0, + nodes: Ptr::null(), + node_count: 0, + } + } + + #[inline(always)] + fn node(&self, i: usize) -> &LinearBVHNode { + unsafe { self.nodes.at(i) } + } + + #[inline(always)] + fn primitive(&self, i: usize) -> &Primitive { + unsafe { self.primitives.at(i) } + } +} + +impl PrimitiveTrait for DeviceBVHAggregate { + fn bounds(&self) -> Bounds3f { + if self.nodes.is_null() || self.node_count == 0 { + Bounds3f::default() + } else { + self.node(0).bounds + } + } + + fn intersect(&self, r: &Ray, t_max: Option) -> Option { + if self.nodes.is_null() { + return None; + } + + let mut hit_t = t_max.unwrap_or(Float::INFINITY); + let mut best_si: Option = None; + + let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z()); + let dir_is_neg = [ + if inv_dir.x() < 0.0 { 1 } else { 0 }, + if inv_dir.y() < 0.0 { 1 } else { 0 }, + if inv_dir.z() < 0.0 { 1 } else { 0 }, + ]; + + let mut stack = [0usize; 64]; + let mut stack_ptr = 0; + let mut node_idx = 0usize; + + loop { + let node = self.node(node_idx); + + if node + .bounds + .intersect_p(r.o, hit_t, inv_dir, &dir_is_neg) + .is_none() + { + if stack_ptr == 0 { + break; + } + stack_ptr -= 1; + node_idx = stack[stack_ptr]; + continue; + } + + if node.n_primitives > 0 { + // Leaf: test all primitives + for i in 0..node.n_primitives { + let prim_idx = node.primitives_offset + i as usize; + let prim = self.primitive(prim_idx); + if let Some(si) = prim.intersect(r, Some(hit_t)) { + hit_t = si.t_hit(); + best_si = Some(si); + } + } + + if stack_ptr == 0 { + break; + } + stack_ptr -= 1; + node_idx = stack[stack_ptr]; + } else { + // Interior: push far, visit near + if dir_is_neg[node.axis as usize] == 1 { + stack[stack_ptr] = node_idx + 1; + stack_ptr += 1; + node_idx = node.primitives_offset; // second child + } else { + stack[stack_ptr] = node.primitives_offset; + stack_ptr += 1; + node_idx += 1; // first child + } + } + } + + best_si + } + + fn intersect_p(&self, r: &Ray, t_max: Option) -> bool { + if self.nodes.is_null() || self.node_count == 0 { + return false; + } + + let t_max = t_max.unwrap_or(Float::INFINITY); + + let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z()); + let dir_is_neg = [ + if inv_dir.x() < 0.0 { 1 } else { 0 }, + if inv_dir.y() < 0.0 { 1 } else { 0 }, + if inv_dir.z() < 0.0 { 1 } else { 0 }, + ]; + + let mut stack = [0usize; 64]; + let mut stack_ptr = 0; + let mut node_idx = 0usize; + + loop { + let node = self.node(node_idx); + + if node + .bounds + .intersect_p(r.o, t_max, inv_dir, &dir_is_neg) + .is_none() + { + if stack_ptr == 0 { + break; + } + stack_ptr -= 1; + node_idx = stack[stack_ptr]; + continue; + } + + if node.n_primitives > 0 { + for i in 0..node.n_primitives { + let prim_idx = node.primitives_offset + i as usize; + let prim = self.primitive(prim_idx); + if prim.intersect_p(r, Some(t_max)) { + return true; + } + } + + if stack_ptr == 0 { + break; + } + stack_ptr -= 1; + node_idx = stack[stack_ptr]; + } else { + if dir_is_neg[node.axis as usize] == 1 { + stack[stack_ptr] = node_idx + 1; + stack_ptr += 1; + node_idx = node.primitives_offset; + } else { + stack[stack_ptr] = node.primitives_offset; + stack_ptr += 1; + node_idx += 1; + } + } + } + + false + } +} diff --git a/shared/src/core/mod.rs b/shared/src/core/mod.rs index f5db8f2..c9c0747 100644 --- a/shared/src/core/mod.rs +++ b/shared/src/core/mod.rs @@ -1,3 +1,4 @@ +pub mod aggregates; pub mod bsdf; pub mod bssrdf; pub mod bxdf; diff --git a/shared/src/core/primitive.rs b/shared/src/core/primitive.rs index ae8b908..6cbb18d 100644 --- a/shared/src/core/primitive.rs +++ b/shared/src/core/primitive.rs @@ -1,4 +1,5 @@ use crate::core::geometry::{Bounds3f, Ray}; +use crate::core::aggregates::DeviceBVHAggregate; use crate::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction}; use crate::core::light::Light; use crate::core::material::Material; @@ -10,6 +11,7 @@ use crate::utils::hash::hash_float; use crate::utils::transform::{AnimatedTransform, Transform}; use crate::utils::Ptr; use alloc::boxed::Box; +use alloc::sync::Arc; use enum_dispatch::enum_dispatch; @@ -111,7 +113,7 @@ impl PrimitiveTrait for SimplePrimitive { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Copy)] pub struct TransformedPrimitive { pub primitive: Ptr, pub render_from_primitive: Ptr, @@ -149,8 +151,8 @@ impl PrimitiveTrait for TransformedPrimitive { #[repr(C)] #[derive(Debug, Copy, Clone)] pub struct AnimatedPrimitive { - primitive: Ptr, - render_from_primitive: Ptr, + pub primitive: Ptr, + pub render_from_primitive: Ptr, } impl PrimitiveTrait for AnimatedPrimitive { @@ -186,41 +188,7 @@ pub struct LinearBVHNode { bounds: Bounds3f, } -#[repr(C)] #[derive(Debug, Clone, Copy)] -pub struct BVHAggregatePrimitive { - max_prims_in_node: u32, - primitives: Ptr<[Primitive]>, - nodes: Ptr, -} - -impl PrimitiveTrait for BVHAggregatePrimitive { - fn bounds(&self) -> Bounds3f { - if !self.nodes.is_null() { - self.nodes.bounds - } else { - Bounds3f::default() - } - } - - fn intersect(&self, _r: &Ray, _t_max: Option) -> Option { - if !self.nodes.is_null() { - return None; - } - todo!() - // self.intersect(r, t_max) - } - - fn intersect_p(&self, _r: &Ray, _t_max: Option) -> bool { - if !self.nodes.is_null() { - return false; - } - todo!() - // self.intersect_p(r, t_max) - } -} - -#[derive(Debug, Clone)] pub struct KdTreeAggregate; impl PrimitiveTrait for KdTreeAggregate { @@ -237,41 +205,13 @@ impl PrimitiveTrait for KdTreeAggregate { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Copy)] #[enum_dispatch(PrimitiveTrait)] pub enum Primitive { Simple(SimplePrimitive), Geometric(GeometricPrimitive), Transformed(TransformedPrimitive), Animated(AnimatedPrimitive), - BVH(BVHAggregatePrimitive), + BVH(DeviceBVHAggregate), 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/src/core/aggregates.rs b/src/core/aggregates.rs index 7ca1b1f..4cb8dea 100644 --- a/src/core/aggregates.rs +++ b/src/core/aggregates.rs @@ -1,9 +1,11 @@ use rayon::prelude::*; +use shared::core::aggregates::DeviceBVHAggregate; use shared::core::geometry::{Bounds3f, Point3f, Ray, Vector3f}; -use shared::core::primitive::PrimitiveTrait; +use shared::core::primitive::{Primitive, PrimitiveTrait}; use shared::core::shape::ShapeIntersection; use shared::utils::math::encode_morton_3; use shared::utils::{find_interval, partition_slice}; +use crate::Arena; use shared::Float; use std::cmp::Ordering; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; @@ -893,3 +895,29 @@ impl BVHAggregate

{ false } } + +impl BVHAggregate { + pub fn to_device(&self, arena: &mut Arena) -> DeviceBVHAggregate { + let (prims_ptr, _) = arena.alloc_slice(&self.primitives); + + let shared_nodes: Vec = self.nodes + .iter() + .map(|n| shared::core::aggregates::LinearBVHNode { + bounds: n.bounds, + primitives_offset: n.primitives_offset, + n_primitives: n.n_primitives, + axis: n.axis, + pad: 0, + }) + .collect(); + let (nodes_ptr, _) = arena.alloc_slice(&shared_nodes); + + DeviceBVHAggregate { + max_prims_in_node: self.max_prims_in_node as u32, + primitives: prims_ptr, + primitive_count: self.primitives.len() as u32, + nodes: nodes_ptr, + node_count: self.nodes.len() as u32, + } + } +} diff --git a/src/core/scene/scene.rs b/src/core/scene/scene.rs index 103954c..e6bb424 100644 --- a/src/core/scene/scene.rs +++ b/src/core/scene/scene.rs @@ -15,7 +15,6 @@ use crate::utils::{resolve_filename, Upload}; use crate::{Arena, FileLoc}; use anyhow::{anyhow, Result}; use parking_lot::Mutex; -use rayon::prelude::*; use shared::core::camera::{Camera, CameraTransform}; use shared::core::color::LINEAR; use shared::core::film::Film; @@ -23,7 +22,7 @@ use shared::core::filter::Filter; use shared::core::light::Light; use shared::core::material::Material; use shared::core::medium::{Medium, MediumInterface}; -use shared::core::primitive::{GeometricPrimitive, Primitive, SimplePrimitive}; +use shared::core::primitive::{AnimatedPrimitive, GeometricPrimitive, Primitive, SimplePrimitive}; use shared::core::sampler::Sampler; use shared::core::shape::Shape; use shared::core::texture::SpectrumType; @@ -575,7 +574,7 @@ impl BasicScene { let al_entity = &light_state.area_lights[light_idx]; - let alpha_tex = self.get_alpha_texture( + let alpha_tex = Self::get_alpha_texture( &entity.base.parameters, &entity.base.loc, &textures.float_textures, @@ -628,99 +627,271 @@ impl BasicScene { &self, textures: &NamedTextures, named_materials: &HashMap, - materials: &Vec, - shape_lights: &HashMap>, + materials: &[Material], arena: &mut Arena, - ) -> Vec { - let shapes = self.shapes.lock(); - let animated_shapes = self.animated_shapes.lock(); + ) -> (Vec, Vec>) { + let entities = self.shapes.lock(); + let animated = self.animated_shapes.lock(); + let light_state = self.light_state.lock(); let media = self.media_state.lock(); - - let lookup = SceneLookup { - textures, - media: &media.map, - named_materials, - materials, - shape_lights, - }; + let film_cs = self.film_colorspace.lock(); let mut primitives = Vec::new(); + let mut area_lights = Vec::new(); - let loaded = self.load_shapes_parallel(&shapes, &lookup, arena); - primitives.extend(self.upload_shapes(arena, &shapes, loaded, &lookup)); - - let loaded_anim = self.load_animated_shapes_parallel(&animated_shapes, &lookup, arena); - primitives.extend(self.upload_animated_shapes( - arena, - &animated_shapes, - loaded_anim, - &lookup, - )); - - primitives - } - - pub fn load_shapes_parallel( - &self, - entities: &[ShapeSceneEntity], - lookup: &SceneLookup, - arena: &mut Arena, - ) -> Vec { - // Flat vector with context - let mut shapes_with_context = Vec::new(); - - for (entity_index, entity) in entities.iter().enumerate() { - let shapes = Shape::create( - &entity.base.name, - *entity.render_from_object.as_ref(), - *entity.object_from_render.as_ref(), - entity.reverse_orientation, - entity.base.parameters.clone(), - &lookup.textures.float_textures, - entity.base.loc.clone(), + for entity in entities.iter() { + Self::build_primitives_for_entity( + entity, + textures, + named_materials, + materials, + &light_state, + &media, + film_cs.as_ref().map(|v| &**v), arena, - ) - .unwrap_or_else(|e| { - eprintln!("Shape '{}' failed: {}", entity.base.name, e); - Vec::new() - }); - - for (shape_index, shape) in shapes.into_iter().enumerate() { - shapes_with_context.push(ShapeWithContext { - shape, - entity_index, - shape_index, - }); - } + &mut primitives, + &mut area_lights, + ); } - shapes_with_context + for entity in animated.iter() { + Self::build_animated_primitives_for_entity( + entity, + textures, + named_materials, + materials, + &light_state, + &media, + film_cs.as_ref().map(|v| &**v), + arena, + &mut primitives, + &mut area_lights, + ); + } + + (primitives, area_lights) } - fn load_animated_shapes_parallel( - &self, - entities: &[AnimatedShapeSceneEntity], - lookup: &SceneLookup, - arena: &Arena, - ) -> Vec> { - entities - .par_iter() - .flat_map(|sh| { - Shape::create( - &sh.transformed_base.base.name, - *sh.identity.as_ref(), - *sh.identity.as_ref(), - sh.reverse_orientation, - sh.transformed_base.base.parameters.clone(), - &lookup.textures.float_textures, - sh.transformed_base.base.loc.clone(), + + fn build_primitives_for_entity( + entity: &ShapeSceneEntity, + textures: &NamedTextures, + named_materials: &HashMap, + materials: &[Material], + light_state: &LightState, + media: &MediaState, + film_cs: Option<&RGBColorSpace>, + arena: &mut Arena, + primitives: &mut Vec, + area_lights: &mut Vec>, + ) { + let shapes = Shape::create( + &entity.base.name, + *entity.render_from_object.as_ref(), + *entity.object_from_render.as_ref(), + entity.reverse_orientation, + entity.base.parameters.clone(), + &textures.float_textures, + entity.base.loc.clone(), + arena, + ) + .unwrap_or_else(|e| { + eprintln!("Shape '{}' failed: {}", entity.base.name, e); + Vec::new() + }); + + let mtl = match &entity.material { + MaterialRef::Name(name) => named_materials.get(name).copied(), + MaterialRef::Index(idx) => materials.get(*idx).copied(), + MaterialRef::None => None, + } + .unwrap_or_else(|| crate::core::material::default_diffuse_material(arena)); + + let alpha_tex = Self::get_alpha_texture( + &entity.base.parameters, + &entity.base.loc, + &textures.float_textures, + ); + + let mi = MediumInterface { + inside: media + .map + .get(&entity.inside_medium) + .map(|m| Ptr::from(m.as_ref())) + .unwrap_or(Ptr::null()), + outside: media + .map + .get(&entity.outside_medium) + .map(|m| Ptr::from(m.as_ref())) + .unwrap_or(Ptr::null()), + }; + + let al_params = entity.light_index.map(|idx| &light_state.area_lights[idx]); + + for shape in shapes { + let area_light = al_params.and_then(|al_entity| { + let colorspace_ref = al_entity.parameters.color_space.as_deref().or(film_cs); + let default_alpha = Arc::new(FloatTexture::default()); + let alpha_ref = alpha_tex.as_ref().unwrap_or(&default_alpha); + + crate::core::light::create_area_light( + *entity.render_from_object, + None, + &al_entity.parameters, + &al_entity.loc, + &shape, + alpha_ref, + colorspace_ref, arena, ) - .unwrap_or_else(|e| { - eprintln!("Shape '{}' failed: {}", sh.transformed_base.base.name, e); - Vec::new() - }) - }) - .collect() + .ok() + }); + + let uploaded_light = area_light + .as_ref() + .map(|l| l.upload(arena)) + .unwrap_or(Ptr::null()); + + if let Some(ref light) = area_light { + area_lights.push(Arc::new(light.clone())); + } + + let shape_ptr = shape.upload(arena); + + let prim = + if uploaded_light.is_null() && !mi.is_medium_transition() && alpha_tex.is_none() { + Primitive::Simple(SimplePrimitive::new(shape_ptr, Ptr::from(&mtl))) + } else { + Primitive::Geometric(GeometricPrimitive::new( + shape_ptr, + mtl.upload(arena), + uploaded_light, + mi.clone(), + alpha_tex + .as_ref() + .map(|t| t.upload(arena)) + .unwrap_or(Ptr::null()), + )) + }; + + primitives.push(prim); + } + } + + fn build_animated_primitives_for_entity( + entity: &AnimatedShapeSceneEntity, + textures: &NamedTextures, + named_materials: &HashMap, + materials: &[Material], + light_state: &LightState, + media: &MediaState, + film_cs: Option<&RGBColorSpace>, + arena: &mut Arena, + primitives: &mut Vec, + area_lights: &mut Vec>, + ) { + let shapes = Shape::create( + &entity.transformed_base.base.name, + entity.transformed_base.render_from_object.start_transform, + entity + .transformed_base + .render_from_object + .start_transform + .inverse(), + entity.reverse_orientation, + entity.transformed_base.base.parameters.clone(), + &textures.float_textures, + entity.transformed_base.base.loc.clone(), + arena, + ) + .unwrap_or_else(|e| { + eprintln!( + "Animated shape '{}' failed: {}", + entity.transformed_base.base.name, e + ); + Vec::new() + }); + + let mtl = match &entity.material { + MaterialRef::Name(name) => named_materials.get(name).copied(), + MaterialRef::Index(idx) => materials.get(*idx).copied(), + MaterialRef::None => None, + } + .unwrap_or_else(|| crate::core::material::default_diffuse_material(arena)); + + let alpha_tex = Self::get_alpha_texture( + &entity.transformed_base.base.parameters, + &entity.transformed_base.base.loc, + &textures.float_textures, + ); + + let mi = MediumInterface { + inside: media + .map + .get(&entity.inside_medium) + .map(|m| Ptr::from(m.as_ref())) + .unwrap_or(Ptr::null()), + outside: media + .map + .get(&entity.outside_medium) + .map(|m| Ptr::from(m.as_ref())) + .unwrap_or(Ptr::null()), + }; + + let al_params = entity.light_index.map(|idx| &light_state.area_lights[idx]); + + for shape in shapes { + let area_light = al_params.and_then(|al_entity| { + let colorspace_ref = al_entity.parameters.color_space.as_deref().or(film_cs); + let default_alpha = Arc::new(FloatTexture::default()); + let alpha_ref = alpha_tex.as_ref().unwrap_or(&default_alpha); + + crate::core::light::create_area_light( + entity.transformed_base.render_from_object.start_transform, + None, + &al_entity.parameters, + &al_entity.loc, + &shape, + alpha_ref, + colorspace_ref, + arena, + ) + .ok() + }); + + let uploaded_light = area_light + .as_ref() + .map(|l| l.upload(arena)) + .unwrap_or(Ptr::null()); + + if let Some(ref light) = area_light { + area_lights.push(Arc::new(light.clone())); + } + + let shape_ptr = shape.upload(arena); + + let base_prim = + if uploaded_light.is_null() && !mi.is_medium_transition() && alpha_tex.is_none() { + Primitive::Simple(SimplePrimitive::new(shape_ptr, Ptr::from(&mtl))) + } else { + Primitive::Geometric(GeometricPrimitive::new( + shape_ptr, + mtl.upload(arena), + uploaded_light, + mi.clone(), + alpha_tex + .as_ref() + .map(|t| t.upload(arena)) + .unwrap_or(Ptr::null()), + )) + }; + + let base_ptr = arena.alloc(base_prim); + + primitives.push(Primitive::Animated(AnimatedPrimitive { + primitive: base_ptr, + render_from_primitive: arena.alloc(entity.transformed_base.render_from_object), + })); + } } fn upload_shapes( @@ -731,11 +902,16 @@ impl BasicScene { lookup: &SceneLookup, ) -> Vec { let mut primitives = Vec::new(); + let mut count = 0; for shape_ctx in loaded { + count += 1; + if count % 500_000 == 0 { + eprintln!(" processed {} primitives", count); + } let entity = &entities[shape_ctx.entity_index]; - let alpha_tex = self.get_alpha_texture( + let alpha_tex = Self::get_alpha_texture( &entity.base.parameters, &entity.base.loc, &lookup.textures.float_textures, @@ -885,7 +1061,6 @@ impl BasicScene { } fn get_alpha_texture( - &self, params: &ParameterDictionary, loc: &FileLoc, textures: &HashMap>, diff --git a/src/core/scene/state.rs b/src/core/scene/state.rs index 76c5617..570dbc9 100644 --- a/src/core/scene/state.rs +++ b/src/core/scene/state.rs @@ -3,7 +3,6 @@ use crate::core::image::Image; use crate::core::texture::{FloatTexture, SpectrumTexture}; use crate::utils::parallel::AsyncJob; use anyhow::Result; -use shared::core::light::Light; use shared::core::medium::Medium; use std::collections::{HashMap, HashSet}; use std::sync::Arc; diff --git a/src/integrators/mod.rs b/src/integrators/mod.rs index eb1853d..1ea4497 100644 --- a/src/integrators/mod.rs +++ b/src/integrators/mod.rs @@ -1,8 +1,8 @@ -mod base; -mod constants; -mod path; -mod pipeline; -mod state; +pub mod base; +pub mod constants; +pub mod path; +pub mod pipeline; +pub mod state; pub use path::PathIntegrator; diff --git a/src/integrators/path.rs b/src/integrators/path.rs index 5b5fe19..0ad517a 100644 --- a/src/integrators/path.rs +++ b/src/integrators/path.rs @@ -64,7 +64,7 @@ impl PathConfig { } pub struct PathIntegrator { - base: IntegratorBase, + pub base: IntegratorBase, camera: Arc, sampler: LightSampler, config: PathConfig, diff --git a/src/lights/sampler.rs b/src/lights/sampler.rs index 1e66946..7d4159e 100644 --- a/src/lights/sampler.rs +++ b/src/lights/sampler.rs @@ -1,5 +1,7 @@ use crate::utils::sampling::AliasTableHost; +use crate::Arena; use shared::core::light::{Light, LightTrait}; +use shared::lights::sampler::PowerLightSampler; use shared::spectra::{SampledSpectrum, SampledWavelengths}; use std::collections::HashMap; use std::sync::Arc; @@ -44,4 +46,16 @@ impl PowerSamplerHost { alias_table, } } + + pub fn to_device(&self, arena: &Arena) -> PowerLightSampler { + let device_lights: Vec = self.lights.iter().map(|l| (**l).clone()).collect(); + let (lights_ptr, _) = arena.alloc_slice(&device_lights); + let alias_device = self.alias_table.to_device(arena); + + PowerLightSampler { + lights: lights_ptr, + lights_len: self.lights.len() as u32, + alias_table: alias_device, + } + } } diff --git a/src/utils/arena.rs b/src/utils/arena.rs index d996ac0..f6267f2 100644 --- a/src/utils/arena.rs +++ b/src/utils/arena.rs @@ -22,6 +22,7 @@ use shared::utils::sampling::{ use shared::utils::Ptr; use std::alloc::Layout; use std::collections::HashMap; +use std::panic::Location; use std::slice::from_raw_parts; use std::sync::Arc; @@ -31,16 +32,26 @@ pub struct Arena { } struct ArenaInner { - buffer: Vec<(*mut u8, Layout)>, + blocks: Vec<(*mut u8, Layout)>, + current_block: *mut u8, + current_offset: usize, + current_capacity: usize, + current_align: usize, texture_cache: HashMap, } +const DEFAULT_BLOCK_SIZE: usize = 256 * 1024; + impl Arena { pub fn new(allocator: A) -> Self { Self { allocator, inner: Mutex::new(ArenaInner { - buffer: Vec::new(), + blocks: Vec::new(), + current_block: std::ptr::null_mut(), + current_offset: 0, + current_capacity: 0, + current_align: 1, texture_cache: HashMap::new(), }), } @@ -48,13 +59,89 @@ impl Arena { pub fn alloc(&self, value: T) -> Ptr { let layout = Layout::new::(); - let ptr = unsafe { self.allocator.alloc(layout) } as *mut T; - unsafe { ptr.write(value) }; + let mut inner = self.inner.lock(); - self.inner.lock().buffer.push((ptr as *mut u8, layout)); + let aligned = (inner.current_offset + layout.align() - 1) & !(layout.align() - 1); + + // Checking if current block alignment is sufficient + if aligned + layout.size() > inner.current_capacity || inner.current_align < layout.align() + { + let block_size = DEFAULT_BLOCK_SIZE.max(layout.size() * 2); + let block_layout = Layout::from_size_align(block_size, layout.align().max(16)).unwrap(); + let block = unsafe { self.allocator.alloc(block_layout) }; + + // null check + if block.is_null() { + panic!( + "Arena alloc failed at {}:{} — size={} align={}", + caller.file(), + caller.line(), + block_layout.size(), + block_layout.align() + ); + } + + inner.blocks.push((block, block_layout)); + inner.current_block = block; + inner.current_offset = 0; + inner.current_capacity = block_size; + inner.current_align = block_layout.align(); // NEW + + let aligned = (0 + layout.align() - 1) & !(layout.align() - 1); + let ptr = unsafe { inner.current_block.add(aligned) as *mut T }; + unsafe { ptr.write(value) }; + inner.current_offset = aligned + layout.size(); + return Ptr::from_raw(ptr); + } + + let ptr = unsafe { inner.current_block.add(aligned) as *mut T }; + unsafe { ptr.write(value) }; + inner.current_offset = aligned + layout.size(); Ptr::from_raw(ptr) } + pub fn alloc_slice(&self, values: &[T]) -> (Ptr, usize) { + if values.is_empty() { + return (Ptr::null(), 0); + } + let layout = Layout::array::(values.len()).unwrap(); + let mut inner = self.inner.lock(); + + let aligned = (inner.current_offset + layout.align() - 1) & !(layout.align() - 1); + + if aligned + layout.size() > inner.current_capacity || inner.current_align < layout.align() + { + let block_size = DEFAULT_BLOCK_SIZE.max(layout.size() * 2); + let block_layout = Layout::from_size_align(block_size, layout.align().max(16)).unwrap(); + let block = unsafe { self.allocator.alloc(block_layout) }; + + if block.is_null() { + panic!( + "Arena allocation failed: size={} align={}", + block_layout.size(), + block_layout.align() + ); + } + + inner.blocks.push((block, block_layout)); + inner.current_block = block; + inner.current_offset = 0; + inner.current_capacity = block_size; + inner.current_align = block_layout.align(); // NEW + + let aligned = 0; + let ptr = unsafe { inner.current_block.add(aligned) as *mut T }; + unsafe { std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len()) }; + inner.current_offset = aligned + layout.size(); + return (Ptr::from_raw(ptr), values.len()); + } + + let ptr = unsafe { inner.current_block.add(aligned) as *mut T }; + unsafe { std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len()) }; + inner.current_offset = aligned + layout.size(); + (Ptr::from_raw(ptr), values.len()) + } + pub fn alloc_opt(&self, value: Option) -> Ptr { match value { Some(v) => self.alloc(v), @@ -62,19 +149,6 @@ impl Arena { } } - pub fn alloc_slice(&self, values: &[T]) -> (Ptr, usize) { - if values.is_empty() { - return (Ptr::null(), 0); - } - - let layout = Layout::array::(values.len()).unwrap(); - let ptr = unsafe { self.allocator.alloc(layout) } as *mut T; - unsafe { std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len()) }; - - self.inner.lock().buffer.push((ptr as *mut u8, layout)); - (Ptr::from_raw(ptr), values.len()) - } - pub fn get_texture_object(&self, mipmap: &Arc) -> u64 { let key = Arc::as_ptr(mipmap) as usize; let mut inner = self.inner.lock(); @@ -102,7 +176,7 @@ impl Default for Arena { impl Drop for Arena { fn drop(&mut self) { let inner = self.inner.get_mut(); - for (ptr, layout) in inner.buffer.drain(..) { + for (ptr, layout) in inner.blocks.drain(..) { unsafe { self.allocator.dealloc(ptr, layout) }; } } @@ -174,7 +248,12 @@ impl Upload for SpectrumTexture { scale: tex.base.scale, invert: tex.base.invert, is_single_channel: tex.base.mipmap.is_single_channel(), - color_space: tex.base.mipmap.color_space.clone().unwrap_or_else(crate::spectra::default_colorspace), + color_space: tex + .base + .mipmap + .color_space + .clone() + .unwrap_or_else(crate::spectra::default_colorspace), spectrum_type: tex.spectrum_type, }; GPUSpectrumTexture::Image(gpu_img) diff --git a/src/utils/sampling.rs b/src/utils/sampling.rs index e9ed86d..591d287 100644 --- a/src/utils/sampling.rs +++ b/src/utils/sampling.rs @@ -2,13 +2,13 @@ use crate::core::image::Image; use crate::utils::arena::Arena; use crate::utils::backend::GpuAllocator; use crate::utils::containers::Array2D; -use shared::Float; use shared::core::geometry::{Bounds2f, Point2i, Vector2f, Vector2i}; use shared::utils::sampling::{ AliasTable, Bin, DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceSummedAreaTable, DeviceWindowedPiecewiseConstant2D, PiecewiseLinear2D, }; -use shared::utils::{Ptr, gpu_array_from_fn}; +use shared::utils::{gpu_array_from_fn, Ptr}; +use shared::Float; use std::sync::Arc; #[derive(Debug, Clone)] @@ -443,6 +443,20 @@ impl AliasTableHost { _storage: bins, } } + + pub fn to_device(&self, arena: &Arena) -> AliasTable { + if self._storage.is_empty() { + return AliasTable { + bins: Ptr::null(), + size: 0, + }; + } + let (bins_ptr, _) = arena.alloc_slice(&self._storage); + AliasTable { + bins: bins_ptr, + size: self._storage.len() as u32, + } + } } #[derive(Clone, Debug)]