pbrt/kernels/src/workitem.rs

535 lines
14 KiB
Rust

#![allow(clippy::too_many_arguments)]
use super::Float4;
use crate::Float;
use crate::core::geometry::{Normal3f, Point2f, Point2i, Point3f, Point3fi, Ray, Vector3f};
use crate::lights::LightSampleContext;
use crate::soa_struct;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use cust::memory::{CopyDestination, DeviceMemory};
use cust::prelude::*;
#[macro_export]
macro_rules! soa_struct {
(
$(#[$outer:meta])*
pub struct $name:ident {
$(
pub $field:ident : $type:ty
),* $(,)?
}
) => {
#[cfg(feature = "use_gpu")]
$(#[$outer])*
pub struct $name {
capacity: u32,
pub count: cust::memory::DeviceBuffer<u32>,
$(
pub $field: cust::memory::DeviceBuffer<$type>,
)*
}
#[cfg(feature = "use_gpu")]
impl $name {
pub fn new(capacity: usize) -> cust::error::CudaResult<Self> {
use cust::memory::DeviceBuffer;
Ok(Self {
capacity: capacity as u32,
count: DeviceBuffer::zeroed(1)?,
$(
$field: DeviceBuffer::zeroed(capacity)?,
)*
})
}
pub fn len(&self) -> cust::error::CudaResult<u32> {
let mut host_count = [0u32; 1];
self.count.copy_to(&mut host_count)?;
Ok(host_count[0])
}
pub fn reset(&mut self) -> cust::error::CudaResult<()> {
self.count.copy_from(&[0])
}
// Generate the View name
pub fn as_view(&mut self) -> paste::paste! { [<$name View>] } {
paste::paste! {
[<$name View>] {
capacity: self.capacity,
count: self.count.as_device_ptr().as_mut_ptr(),
$(
$field: self.$field.as_device_ptr().as_raw() as *mut $type,
)*
}
}
}
}
paste::paste! {
#[repr(C)]
#[derive(Clone, Copy)]
pub struct [<$name View>] {
pub capacity: u32,
pub count: *mut u32,
$(
pub $field: *mut $type,
)*
}
unsafe impl cust::memory::DeviceCopy for [<$name View>] {}
impl [<$name View>] {
// The raw push that fills every field
#[cfg(feature = "use_gpu")]
pub unsafe fn push(&self, $( $field : $type ),* ) -> Option<u32> {
use core::sync::atomic::{AtomicU32, Ordering};
let index = unsafe {
let counter_ptr = self.count as *mut AtomicU32;
(*counter_ptr).fetch_add(1, Ordering::Relaxed)
};
if index >= self.capacity {
return None;
}
unsafe {
$(
*self.$field.add(index as usize) = $field;
)*
}
Some(index)
}
#[cfg(feature = "use_gpu")]
pub unsafe fn size(&self) -> u32 {
use core::sync::atomic::{AtomicU32, Ordering};
unsafe {
(*(self.count as *const AtomicU32)).load(Ordering::Relaxed)
}
}
$(
#[cfg(feature = "use_gpu")]
pub fn [<$field _ptr>](&self) -> *mut $type {
self.$field
}
)*
}
}
};
}
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesDirect {
pub u: Point2f,
pub uc: Float,
}
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesIndirect {
pub uc: Float,
pub rr: Float,
pub u: Point2f,
}
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamplesSubsurface {
pub uc: Float,
pub u: Point2f,
}
#[repr(C)]
#[derive(Clone, Copy, Default)]
pub struct RaySamples {
pub direct: RaySamplesDirect,
pub indirect: RaySamplesIndirect,
pub have_subsurface: bool,
pub subsurface: RaySamplesSubsurface,
}
soa_struct! {
pub struct RayQueue {
pub ray_o: Point3f,
pub ray_d: Vector3f,
pub depth: i32,
pub lambda: SampledWavelengths,
pub pixel_index: u32,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub ctx_pi: Point3f,
pub ctx_n: Normal3f,
pub ctx_ns: Normal3f,
pub eta_scale: Float,
pub specular_bounce: u32,
pub any_non_specular_bounces: u32,
}
}
soa_struct! {
pub struct PixelSampleStateStorage {
pub p_pixel: Point2i,
pub l: SampledSpectrum,
pub lambda: SampledWavelengths,
pub filter_weight: Float,
pub visible_surface: u32,
pub camera_ray_weight: SampledSpectrum,
pub rs_direct_packed: Float4,
pub rs_indirect_packed: Float4,
pub rs_subsurface_packed: Float4,
}
}
soa_struct! {
pub struct EscapedRayQueue {
pub ray_o: Point3f,
pub ray_d: Vector3f,
pub depth: i32,
pub lambda: SampledWavelengths,
pub pixel_index: u32,
pub beta: SampledSpectrum,
pub specular_bounce: u32,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub ctx_pi: Point3f,
pub ctx_n: Normal3f,
pub ctx_ns: Normal3f,
}
}
soa_struct! {
pub struct HitAreaLightQueue {
pub area_light_id: u32, // Light ID
pub p: Point3f,
pub n: Normal3f,
pub uv: Point2f,
pub wo: Vector3f,
pub lambda: SampledWavelengths,
pub depth: i32,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub ctx_pi: Point3f,
pub ctx_n: Normal3f,
pub ctx_ns: Normal3f,
pub specular_bounce: u32,
pub pixel_index: u32,
}
}
soa_struct! {
pub struct ShadowRayQueue {
pub ray_o: Point3f,
pub ray_d: Vector3f,
pub t_max: Float,
pub lambda: SampledWavelengths,
pub ld: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub pixel_index: u32,
}
}
soa_struct! {
pub struct GetBSSRDFAndProbeRayQueue {
pub material_id: u32,
pub lambda: SampledWavelengths,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub p: Point3f,
pub wo: Vector3f,
pub n: Normal3f,
pub ns: Normal3f,
pub dpdus: Vector3f,
pub uv: Point2f,
pub depth: i32,
pub mi_inside: u32,
pub mi_outside: u32,
pub eta_scale: Float,
pub pixel_index: u32,
}
}
soa_struct! {
pub struct SubsurfaceScatterQueue {
pub p0: Point3f,
pub p1: Point3f,
pub depth: i32,
pub material_id: u32,
pub lambda: SampledWavelengths,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub mi_inside: u32,
pub mi_outside: u32,
pub eta_scale: Float,
pub pixel_index: u32,
}
}
soa_struct! {
pub struct MediumSampleQueue {
pub ray_o: Point3f,
pub ray_d: Vector3f,
pub t_max: Float,
pub lambda: SampledWavelengths,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub pixel_index: u32,
pub ctx_pi: Point3f,
pub ctx_n: Normal3f,
pub ctx_ns: Normal3f,
pub specular_bounce: u32,
pub any_non_specular_bounces: u32,
pub eta_scale: Float,
pub area_light_id: u32,
pub pi: Point3fi,
pub n: Normal3f,
pub dpdu: Vector3f,
pub dpdv: Vector3f,
pub wo: Vector3f,
pub uv: Point2f,
pub material_id: u32,
pub ns: Normal3f,
pub dpdus: Vector3f,
pub dpdvs: Vector3f,
pub dndus: Normal3f,
pub dndvs: Normal3f,
pub face_index: i32,
pub mi_inside: u32,
pub mi_outside: u32,
}
}
soa_struct! {
pub struct MaterialEvalQueue {
pub material_id: u32,
pub pi: Point3fi,
pub n: Normal3f,
pub dpdu: Vector3f,
pub dpdv: Vector3f,
pub time: Float,
pub depth: i32,
pub ns: Normal3f,
pub dpdus: Vector3f,
pub dpdvs: Vector3f,
pub dndus: Normal3f,
pub dndvs: Normal3f,
pub uv: Point2f,
pub face_index: i32,
pub lambda: SampledWavelengths,
pub pixel_index: u32,
pub any_non_specular_bounces: u32,
pub wo: Vector3f,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub eta_scale: Float,
pub mi_inside: u32,
pub mi_outside: u32,
}
}
soa_struct! {
pub struct MediumScatterQueue {
pub p: Point3f,
pub depth: usize,
pub lambda: SampledWavelengths,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub wo: Vector3f,
pub time: Float,
pub eta_scale: Float,
pub pixel_index: usize,
// ID
pub phase_function: u32,
pub medium: u32,
}
}
#[repr(C)]
#[derive(Clone, Copy)]
pub struct RayWorkItem {
pub ray: Ray,
pub depth: i32,
pub lambda: SampledWavelengths,
pub pixel_index: u32,
pub beta: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub prev_intr_ctx: LightSampleContext,
pub eta_scale: Float,
pub specular_bounce: bool,
pub any_non_specular_bounces: bool,
}
#[repr(C)]
#[derive(Clone, Copy)]
pub struct EscapedRayWorkItem {
pub ray_o: Point3f,
pub ray_d: Vector3f,
pub depth: i32,
pub lambda: SampledWavelengths,
pub pixel_index: u32,
pub beta: SampledSpectrum,
pub specular_bounce: bool,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub prev_intr_ctx: LightSampleContext,
}
#[repr(C)]
#[derive(Clone, Copy)]
pub struct ShadowRayWorkItem {
pub ray: Ray,
pub t_max: Float,
pub lambda: SampledWavelengths,
pub ld: SampledSpectrum,
pub r_u: SampledSpectrum,
pub r_l: SampledSpectrum,
pub pixel_index: u32,
}
impl RayQueueView {
#[cfg(feature = "use_gpu")]
pub unsafe fn push_work_item(&self, item: RayWorkItem) -> Option<u32> {
unsafe {
self.push(
item.ray.o,
item.ray.d,
item.depth,
item.lambda,
item.pixel_index,
item.beta,
item.r_u,
item.r_l,
item.prev_intr_ctx.pi.into(),
item.prev_intr_ctx.n,
item.prev_intr_ctx.ns,
item.eta_scale,
if item.specular_bounce { 1 } else { 0 },
if item.any_non_specular_bounces { 1 } else { 0 },
)
}
}
}
impl EscapedRayQueueView {
#[cfg(feature = "use_gpu")]
pub unsafe fn push_work_item(&self, r: &RayWorkItem) -> Option<u32> {
unsafe {
self.push(
r.ray.o,
r.ray.d,
r.depth,
r.lambda,
r.pixel_index,
r.beta,
if r.specular_bounce { 1 } else { 0 },
r.r_u,
r.r_l,
r.prev_intr_ctx.pi.into(),
r.prev_intr_ctx.n,
r.prev_intr_ctx.ns,
)
}
}
}
impl PixelSampleStateStorageView {
#[cfg(feature = "use_gpu")]
pub unsafe fn get_samples(&self, index: u32) -> RaySamples {
let i = index as usize;
let (dir, ind, ss) = unsafe {
(
*self.rs_direct_packed.add(i),
*self.rs_indirect_packed.add(i),
*self.rs_subsurface_packed.add(i),
)
};
let direct_u = Point2f::new(dir.v[0], dir.v[1]);
let direct_uc = dir.v[2];
let flags = dir.v[3] as i32;
let have_subsurface = (flags & 1) != 0;
let indirect_uc = ind.v[0];
let indirect_rr = ind.v[1];
let indirect_u = Point2f::new(ind.v[2], ind.v[3]);
let subsurface_uc = ss.v[0];
let subsurface_u = Point2f::new(ss.v[1], ss.v[2]);
RaySamples {
direct: RaySamplesDirect {
u: direct_u,
uc: direct_uc,
},
indirect: RaySamplesIndirect {
uc: indirect_uc,
rr: indirect_rr,
u: indirect_u,
},
have_subsurface,
subsurface: RaySamplesSubsurface {
uc: subsurface_uc,
u: subsurface_u,
},
}
}
#[cfg(feature = "use_gpu")]
pub unsafe fn set_samples(&self, index: u32, rs: RaySamples) {
if index >= self.capacity {
return;
}
let i = index as usize;
let flags = if rs.have_subsurface { 1.0 } else { 0.0 };
let dir = Float4 {
v: [rs.direct.u.0[0], rs.direct.u.0[1], rs.direct.uc, flags],
};
let ind = Float4 {
v: [
rs.indirect.uc,
rs.indirect.rr,
rs.indirect.u.0[0],
rs.indirect.u.0[1],
],
};
unsafe {
*self.rs_direct_packed.add(i) = dir;
*self.rs_indirect_packed.add(i) = ind;
}
if rs.have_subsurface {
let ss = Float4 {
v: [
rs.subsurface.uc,
rs.subsurface.u.0[0],
rs.subsurface.u.0[1],
0.0,
],
};
unsafe {
*self.rs_subsurface_packed.add(i) = ss;
}
}
}
}