diff --git a/src/lib.rs b/src/lib.rs index 824c8c1..0318440 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,7 @@ mod scene; extern crate nalgebra_glm as glm; mod object; + struct MyUserEvent; struct State<'a> { @@ -156,7 +157,7 @@ fn init( #[cfg_attr(target_arch = "wasm32", wasm_bindgen(start))] pub async fn run() { info!("Starting up"); - let scale = 2.2; + let scale = 2.0; let width = 900 * scale as u32; let height = 450 * scale as u32; let (window, event_loop) = init(width, height); @@ -168,11 +169,11 @@ pub async fn run() { last_time: instant::Instant::now(), render_context: RenderContext::new( &window, - &Scene::teapot_scene( + &Scene::raytracing_scene_oneweek( scene::RenderParam { samples_per_pixel: 1, - max_depth: 2, - samples_max_per_pixel: 1, + max_depth: 7, + samples_max_per_pixel: 1000, total_samples: 0, clear_samples: 0, }, diff --git a/src/object/aabb.rs b/src/object/aabb.rs new file mode 100644 index 0000000..afc882b --- /dev/null +++ b/src/object/aabb.rs @@ -0,0 +1,54 @@ +use glm::Vec3; + +#[repr(C)] +#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] +pub struct Bounds { + pub min: Vec3, + pub max: Vec3, +} + +impl Bounds { + pub fn new(min: Vec3, max: Vec3) -> Self { + Self { min, max } + } + + pub fn empty() -> Self { + Self { + min: Vec3::new(0.0, 0.0, 0.0), + max: Vec3::new(0.0, 0.0, 0.0), + } + } + + pub fn union(b1: Bounds, b2: Bounds) -> Bounds { + let min = Vec3::new( + b1.min.x.min(b2.min.x), + b1.min.y.min(b2.min.y), + b1.min.z.min(b2.min.z), + ); + let max = Vec3::new( + b1.max.x.max(b2.max.x), + b1.max.y.max(b2.max.y), + b1.max.z.max(b2.max.z), + ); + Bounds { min, max } + } + + pub fn maximum_extent(&self) -> usize { + let diag = self.max - self.min; + if diag.x > diag.y && diag.x > diag.z { + 0 + } else if diag.y > diag.z { + 1 + } else { + 2 + } + } +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] +pub struct AABB { + pub bounds: Bounds, + pub centroid: Vec3, + pub type_: u32, +} diff --git a/src/object/bvh.rs b/src/object/bvh.rs new file mode 100644 index 0000000..da5c1ed --- /dev/null +++ b/src/object/bvh.rs @@ -0,0 +1,206 @@ +#![allow(dead_code)] +use std::rc::Rc; + +use glm::{Vec3, Vec4}; + +use super::{Bounds, Mesh, Object, ObjectType, Sphere, AABB}; + +#[derive(Debug, Clone)] +struct BVHBuildNode { + bounds: Bounds, + left: Option>, + right: Option>, + split_axis: u32, + first_obj_offset: u32, + n_obj: u32, + obj_type: u32, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)] +pub struct LinearBVHNode { + min: Vec4, + max: Vec3, + offset: u32, + n_obj: u32, + axis: u32, + obj_type: u32, + obj_offset: u32, +} + +#[derive(Debug, Clone)] +pub struct BVH { + root: Rc, + pub total_nodes: u32, +} + +impl BVHBuildNode { + pub fn default() -> Self { + BVHBuildNode { + bounds: Bounds { + min: Vec3::new(0.0, 0.0, 0.0), + max: Vec3::new(0.0, 0.0, 0.0), + }, + left: None, + right: None, + split_axis: 0, + first_obj_offset: 0, + n_obj: 0, + obj_type: 0, + } + } + pub fn init_leaf(first: u32, n: u32, bounds: Bounds) -> BVHBuildNode { + BVHBuildNode { + bounds, + first_obj_offset: first, + n_obj: n, + left: None, + right: None, + split_axis: 0, + obj_type: 0, + } + } + + pub fn init_interior(axis: u32, c0: Rc, c1: Rc) -> BVHBuildNode { + let c0_bounds = c0.bounds; + let c1_bounds = c1.bounds; + + let bounds = Bounds::union(c0_bounds, c1_bounds); + + BVHBuildNode { + bounds, + first_obj_offset: 0, + n_obj: 0, + left: Some(c0), + right: Some(c1), + split_axis: axis, + obj_type: 0, + } + } +} + +fn recursive_build( + aabb: &Vec, + start: usize, + end: usize, + total_nodes: &mut u32, + ordered_objects: &mut Vec, +) -> Rc { + *total_nodes += 1; + let mut bounds = Bounds::empty(); + for i in start..end { + bounds = Bounds::union(bounds, aabb[i].bounds); + } + + let n_obj = end - start; + if n_obj == 1 { + // create leaf node + let first_obj_offset = ordered_objects.len() as u32; + for i in start..end { + ordered_objects.push(aabb[i].type_); + } + return Rc::new(BVHBuildNode::init_leaf( + first_obj_offset, + n_obj as u32, + bounds, + )); + } else { + let mut centroid_bounds = Bounds::empty(); + for i in start..end { + centroid_bounds = Bounds::union(centroid_bounds, aabb[i].bounds); + } + let dim = centroid_bounds.maximum_extent(); + let mid = start + (end - start) / 2; + if centroid_bounds.max[dim] == centroid_bounds.min[dim] { + // create leaf node + let first_obj_offset = ordered_objects.len() as u32; + for i in start..end { + ordered_objects.push(aabb[i].type_); + } + return Rc::new(BVHBuildNode::init_leaf( + first_obj_offset, + n_obj as u32, + bounds, + )); + } else { + // partition objects based on the midpoint + return Rc::new(BVHBuildNode::init_interior( + dim as u32, + recursive_build(aabb, start, mid, total_nodes, ordered_objects), + recursive_build(aabb, mid, end, total_nodes, ordered_objects), + )); + } + } +} + +fn build_aabbs(objects: &Vec, spheres: &Vec, meshes: &Vec) -> Vec { + let mut aabbs = Vec::new(); + for object in objects { + let aabb = match ObjectType::from_u32(object.obj_type) { + ObjectType::Sphere => spheres[object.id as usize].get_aabb(), + ObjectType::Mesh => meshes[object.id as usize].get_aabb(), + }; + aabbs.push(aabb); + } + aabbs +} + +fn print_tree(node: &Rc, depth: u32) { + for _ in 0..depth { + print!(" "); + } + println!( + "Bounds: {:?}, n_obj: {}, split_axis: {}, first_obj_offset: {}, n_obj: {}", + node.bounds, node.n_obj, node.split_axis, node.first_obj_offset, node.n_obj, + ); + if node.n_obj == 0 { + print_tree(&node.left.as_ref().unwrap(), depth + 1); + print_tree(&node.right.as_ref().unwrap(), depth + 1); + } +} + +fn linearize_bvh(node: &Rc, linear_nodes: &mut Vec, offset: &mut u32) { + let linear_node = LinearBVHNode { + min: Vec4::new(node.bounds.min.x, node.bounds.min.y, node.bounds.min.z, 0.0), + max: node.bounds.max, + offset: *offset, + n_obj: node.n_obj, + axis: node.split_axis, + obj_type: node.obj_type, + obj_offset: node.first_obj_offset, + }; + *offset += 1; + linear_nodes.push(linear_node); + if node.n_obj == 0 { + linearize_bvh(node.left.as_ref().unwrap(), linear_nodes, offset); + linearize_bvh(node.right.as_ref().unwrap(), linear_nodes, offset); + } +} +// https://pbr-book.org/3ed-2018/Primitives_and_Intersection_Acceleration/Bounding_Volume_Hierarchies#LinearBVHNode::secondChildOffset +pub fn get_bvh( + objects: &Vec, + spheres: &Vec, + meshes: &Vec, +) -> Vec { + let aabbs = build_aabbs(objects, spheres, meshes); + + let mut ordered_objects = Vec::new(); + let mut total_nodes = 0; + let root = recursive_build( + &aabbs, + 0, + aabbs.len(), + &mut total_nodes, + &mut ordered_objects, + ); + + // print_tree(&root, 0); + let mut linear_nodes = Vec::new(); + let mut offset = 0; + linearize_bvh(&root, &mut linear_nodes, &mut offset); + + for node in &linear_nodes { + println!("{:?}", node); + } + linear_nodes +} diff --git a/src/object/mesh.rs b/src/object/mesh.rs index 88b5dca..18613c8 100644 --- a/src/object/mesh.rs +++ b/src/object/mesh.rs @@ -1,3 +1,5 @@ +use super::AABB; + #[repr(C)] #[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable, PartialEq)] // TODO: For the moment, vec4 for padding, include manually @@ -63,4 +65,21 @@ impl Mesh { }); indices.collect() } + + pub fn get_aabb(&self) -> AABB { + let mut min = glm::vec3(self.vertices[0].x, self.vertices[0].y, self.vertices[0].z); + let mut max = glm::vec3(self.vertices[0].x, self.vertices[0].y, self.vertices[0].z); + + for vertex in &self.vertices { + let vertex_pos = glm::vec3(vertex.x, vertex.y, vertex.z); + min = glm::min2(&min, &vertex_pos); + max = glm::max2(&max, &vertex_pos); + } + + AABB { + bounds: super::Bounds { min, max }, + centroid: (min + max) / 2.0, + type_: 1, + } + } } diff --git a/src/object/mod.rs b/src/object/mod.rs index 605c50f..159c709 100644 --- a/src/object/mod.rs +++ b/src/object/mod.rs @@ -4,6 +4,12 @@ pub use sphere::Sphere; mod mesh; pub use mesh::Mesh; +mod aabb; +pub use aabb::{Bounds, AABB}; + +mod bvh; +pub use bvh::get_bvh; + #[repr(C)] #[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable, PartialEq)] pub struct Object { @@ -26,3 +32,13 @@ pub enum ObjectType { Sphere = 0, Mesh = 1, } + +impl ObjectType { + pub fn from_u32(value: u32) -> Self { + match value { + 0 => ObjectType::Sphere, + 1 => ObjectType::Mesh, + _ => panic!("Unknown object type"), + } + } +} diff --git a/src/object/sphere.rs b/src/object/sphere.rs index bb692eb..314ebc9 100644 --- a/src/object/sphere.rs +++ b/src/object/sphere.rs @@ -1,3 +1,5 @@ +use super::aabb; + #[repr(C)] #[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable, PartialEq)] pub struct Sphere { @@ -27,4 +29,18 @@ impl Sphere { _padding: [0; 2], } } + + pub fn get_aabb(&self) -> aabb::AABB { + let radius = glm::vec3(self.radius, self.radius, self.radius); + let center = self.center.xyz(); + + aabb::AABB { + bounds: aabb::Bounds { + min: center - radius, + max: center + radius, + }, + centroid: center, + type_: 0, + } + } } diff --git a/src/render_context.rs b/src/render_context.rs index e81e8b2..d1a12a2 100644 --- a/src/render_context.rs +++ b/src/render_context.rs @@ -6,6 +6,7 @@ use winit::{ }; use crate::{ + object::get_bvh, scene::{GpuCamera, GpuMaterial, Scene}, utils::{EguiRenderer, StorageBuffer, UniformBuffer, Vertex}, }; @@ -218,6 +219,15 @@ impl<'a> RenderContext<'a> { Some("surfaces buffer"), ); + let aabb_buffer = StorageBuffer::new_from_bytes( + &device, + bytemuck::cast_slice( + get_bvh(&scene.objects, &scene.spheres, &scene.meshes).as_slice(), + ), + 5_u32, + Some("aabb buffer"), + ); + let scene_bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { entries: &[ @@ -226,6 +236,7 @@ impl<'a> RenderContext<'a> { material_buffer.layout(wgpu::ShaderStages::FRAGMENT, true), texture_buffer.layout(wgpu::ShaderStages::FRAGMENT, true), surfaces_buffer.layout(wgpu::ShaderStages::FRAGMENT, true), + aabb_buffer.layout(wgpu::ShaderStages::FRAGMENT, true), ], label: Some("scene layout"), }); @@ -238,6 +249,7 @@ impl<'a> RenderContext<'a> { material_buffer.binding(), texture_buffer.binding(), surfaces_buffer.binding(), + aabb_buffer.binding(), ], label: Some("scene bind group"), }); diff --git a/src/shader/raytracing.wgsl b/src/shader/raytracing.wgsl index d1c9cd8..2024280 100644 --- a/src/shader/raytracing.wgsl +++ b/src/shader/raytracing.wgsl @@ -27,7 +27,7 @@ const MAX_T = 1000f; @group(1) @binding(3) var textures: array>; // TODO: for now, surfaces will represent a single Mesh @group(1) @binding(4) var surfaces: array; - +@group(1) @binding(5) var aabbs: array; @vertex fn vs_main( @@ -166,6 +166,21 @@ struct Ray { direction: vec3, }; +struct Bouding { + min: vec4, + max: vec3, +}; + +struct AABB { + min: vec4, + max: vec3, + offset: u32, + n_obj: u32, + axis: u32, + obj_type: u32, + obj_offset: u32, +}; + struct Sphere { center: vec4, radius: f32, @@ -295,35 +310,107 @@ fn hit_triangle( return false; } +fn hit_object( + object_index: u32, + ray: Ray, + ray_min: f32, + ray_max: f32, + hit: ptr, +) -> bool { + switch objects[object_index].obj_type { + case OBJECT_SPHERE: { + return hit_sphere(object_index, ray, ray_min, ray_max, hit); + } + case OBJECT_MESHES: { + return hit_triangle(object_index, ray, ray_min, ray_max, hit); + } + default: { + return false; + } + } +} + fn check_intersection(ray: Ray, intersection: ptr) -> bool { var closest_so_far = MAX_T; var hit_anything = false; var tmp_rec = HitRecord(); for (var i = 0u; i < arrayLength(&objects); i += 1u) { - switch (objects[i].obj_type) { - case OBJECT_SPHERE: { - if hit_sphere(objects[i].id, ray, MIN_T, closest_so_far, &tmp_rec) { + if hit_object(i, ray, MIN_T, closest_so_far, &tmp_rec) { + hit_anything = true; + closest_so_far = tmp_rec.t; + *intersection = tmp_rec; + } + } + + return hit_anything; +} + +// https://pbr-book.org/3ed-2018/Primitives_and_Intersection_Acceleration/Bounding_Volume_Hierarchies#LinearBVHNode::secondChildOffset +fn check_intersection_bvh(ray: Ray, intersection: ptr) -> bool { + var closest_so_far = MAX_T; + var hit_anything = false; + var tmp_rec = HitRecord(); + + let invD = 1.0 / ray.direction; + + var to_visit_offset = 0u; + var current_node_index = 0u; + var nodes_to_visit = array(); + + var ray_ = ray; + + var node = aabbs[10]; + let bounds = Bouding(node.min, node.max); + if rayIntersectBV(&ray_, bounds) { + if node.n_obj > 0 { + for (var j = 0u; j < node.n_obj; j += 1u) { + if hit_object(node.obj_offset - 1 + j, ray_, MIN_T, closest_so_far, &tmp_rec) { hit_anything = true; closest_so_far = tmp_rec.t; *intersection = tmp_rec; } } - case OBJECT_MESHES: { - for (var j = 0u; j < arrayLength(&surfaces); j += 1u) { - if hit_triangle(j, ray, MIN_T, closest_so_far, &tmp_rec) { - hit_anything = true; - closest_so_far = tmp_rec.t; - *intersection = tmp_rec; - } - } - } - default: { - // Do nothing - } } } + var max_iterations = 64u; + // for (var i = 0u; i < max_iterations; i += 1u) { + // var node = aabbs[current_node_index]; + // let bounds = Bouding(node.min, node.max); + // if rayIntersectBV(&ray_, bounds) { + // if node.n_obj > 0 { + // for (var j = 0u; j < node.n_obj; j += 1u) { + // if hit_object(node.obj_offset - 1 + j, ray_, MIN_T, closest_so_far, &tmp_rec) { + // hit_anything = true; + // closest_so_far = tmp_rec.t; + // *intersection = tmp_rec; + // } + // } + // if to_visit_offset == 0u { + // break; + // } + // to_visit_offset -= 1u; + // current_node_index = nodes_to_visit[to_visit_offset]; + // } else { + // if ray_.direction[node.axis] < 0.0 { + // nodes_to_visit[to_visit_offset] = current_node_index + 1u; + // current_node_index = node.offset; + // } else { + // nodes_to_visit[to_visit_offset] = node.offset; + // current_node_index = current_node_index + 1u; + // } + // to_visit_offset += 1u; + // } + // } else { + // if to_visit_offset == 0u { + // break; + // } + // to_visit_offset -= 1u; + // current_node_index = nodes_to_visit[to_visit_offset]; + // } + // } + return hit_anything; } @@ -350,22 +437,6 @@ fn get_ray(rngState: ptr, x: f32, y: f32) -> Ray { -// pub fn from_linear_rgb(c: [f32; 3]) -> Color { -// let f = |x: f32| -> u32 { -// let y = if x > 0.0031308 { -// let a = 0.055; -// (1.0 + a) * x.powf(-2.4) - a -// } else { -// 12.92 * x -// }; -// (y * 255.0).round() as u32 -// }; -// f(c[0]) << 16 | f(c[1]) << 8 | f(c[2]) -// } - - - - fn ray_color(first_ray: Ray, rngState: ptr) -> vec3 { var ray = first_ray; var sky_color = vec3(0.0); @@ -373,7 +444,7 @@ fn ray_color(first_ray: Ray, rngState: ptr) -> vec3 { for (var i = 0u; i < render_param.max_depth; i += 1u) { var intersection = HitRecord(); - if check_intersection(ray, &intersection) { + if check_intersection_bvh(ray, &intersection) { let material = materials[intersection.material_index]; let scattered = scatter(ray, intersection, material, rngState); color *= scattered.attenuation; @@ -390,6 +461,15 @@ fn ray_color(first_ray: Ray, rngState: ptr) -> vec3 { } +fn rayIntersectBV(ray: ptr, aabb: Bouding) -> bool { + let t0 = ((aabb).min.xyz - (*ray).origin) / (*ray).direction; + let t1 = ((aabb).max.xyz - (*ray).origin) / (*ray).direction; + let tmin = min(t0, t1); + let tmax = max(t0, t1); + let maxMinT = max(tmin.x, max(tmin.y, tmin.z)); + let minMaxT = min(tmax.x, min(tmax.y, tmax.z)); + return maxMinT < minMaxT; +} fn scatter(ray: Ray, hit: HitRecord, material: Material, rngState: ptr) -> Scatter { switch (material.id)