pbrt/src/wavefront/aggregate.rs

use crate::globals::get_options;
use log::debug;
use std::sync::atomic::{AtomicU32, Ordering};

pub static DIAG_SHADOW_UNOCCLUDED: AtomicU32 = AtomicU32::new(0);
use rayon::prelude::*;
use shared::core::geometry::{Bounds3f, Ray, VectorLike};
use shared::core::interaction::{InteractionTrait, SurfaceInteraction};
use shared::core::material::{Material, MaterialTrait};
use shared::core::primitive::{Primitive, PrimitiveTrait};
use shared::core::texture::BasicTextureEvaluator;
use shared::core::texture::TextureEvaluator;
use shared::wavefront::workitems::*;
use shared::wavefront::WavefrontAggregate;
use shared::{Float, Ptr, GVec, gvec_from_slice};

pub struct CpuAggregate {
    pub aggregate: Primitive,
    pub materials: GVec<Material>,
}

impl CpuAggregate {
    pub fn new(aggregate: Primitive, materials: &[Material]) -> Self {
        Self { aggregate, materials: gvec_from_slice(materials) }
    }
}

// fn enqueue_after_intersection(
//     r: RayWorkItem, ray_medium: Medium, t_max: Float, intr: SurfaceInteraction,
//     mut medium_sample_queue: &MediumQueue, mut next_ray_queue: &RayQueue,
//     mut hit_area_light_queue: &HitAreaLightQueue, mut basic_eval_mtl_q: &MaterialEvalQueue,
//     mut universal_eval_mlt_q: MaterialEvalQueue) {
// }

impl WavefrontAggregate for CpuAggregate {
    fn bounds(&self) -> Bounds3f {
        self.aggregate.bounds()
    }

    fn intersect_closest(
        &self,
        max_rays: usize,
        ray_q: &RayQueue,
        escaped_ray_q: &EscapedRayQueue,
        hit_area_light_q: &HitAreaLightQueue,
        basic_eval_mtl_q: &MaterialEvalQueue,
        universal_eval_mtl_q: &MaterialEvalQueue,
        next_ray_q: &RayQueue,
        _pixel_sample_state: &PixelSampleState,
    ) {
        let n_rays = ray_q.size().min(max_rays as u32);

        // Intersect ray with the scene and enqueue resulting work
        (0..n_rays as usize).into_par_iter().for_each(|i| {
            let r = unsafe { ray_q.get(i) };

            let Some(si) = self.aggregate.intersect(&r.ray, None) else {
                // EnqueueMiss
                escaped_ray_q.push(EscapedRayWorkItem {
                    ray_o: r.ray.o,
                    ray_d: r.ray.d,
                    lambda: r.lambda,
                    pixel_index: r.pixel_index,
                    beta: r.beta,
                    r_u: r.r_u,
                    r_l: r.r_l,
                    depth: r.depth,
                    specular_bounce: r.specular_bounce,
                    prev_intr_ctx: r.prev_intr_ctx,
                });
                return;
            };

            let intr = &si.intr;

            // Medium transition
            if intr.material.is_none() {
                let mut next = r;
                next.ray = intr.spawn_ray(r.ray.d);
                next_ray_q.push(next);
                return;
            }

            // Area light hit
            if !intr.area_light.is_none() {
                hit_area_light_q.push(HitAreaLightWorkItem {
                    area_light: intr.area_light,
                    p: intr.p(),
                    n: intr.n(),
                    uv: intr.common.uv,
                    wo: intr.wo(),
                    lambda: r.lambda,
                    pixel_index: r.pixel_index,
                    beta: r.beta,
                    r_u: r.r_u,
                    r_l: r.r_l,
                    depth: r.depth,
                    specular_bounce: r.specular_bounce,
                    prev_intr_ctx: r.prev_intr_ctx,
                });
            }

            // Material eval queue dispatch
            let material = intr.material.get(&self.materials);
            let eval_q = if material.can_evaluate_textures(&BasicTextureEvaluator) {
                basic_eval_mtl_q
            } else {
                universal_eval_mtl_q
            };

            if material.is_conductor() {
                debug!("shading frame: {:?}", intr.shading.dpdu);
                debug!(
                    "dot product: {:?}",
                    intr.shading.dpdu.normalize().dot(intr.n().into())
                );
            }

            let item = MaterialEvalWorkItem {
                p: intr.pi(),
                n: intr.n(),
                ns: intr.shading.n,
                dpdu: intr.dpdu,
                dpdv: intr.dpdv,
                uv: intr.common.uv,
                wo: intr.wo(),
                time: r.ray.time,
                face_index: intr.face_index,
                material: intr.material,
                area_light: intr.area_light,
                medium_interface: intr.common.medium_interface,
                pixel_index: r.pixel_index,
                lambda: r.lambda,
                beta: r.beta,
                r_u: r.r_u,
                any_non_specular_bounces: r.any_non_specular_bounces,
                depth: r.depth,
                eta_scale: r.eta_scale,
                dpdus: intr.shading.dpdu,
                dpdvs: intr.shading.dpdv,
                dndus: intr.shading.dndu,
                dndvs: intr.shading.dndv,
            };
            if let Some(slot) = eval_q.push(item) {
                if slot < 10 {
                    eprintln!(
                        "ENQUEUE[{slot}] pixel={:?} depth={} \
                         p={:?} n={:?} ns={:?} \
                         dpdu={:?} dpdv={:?} \
                         dpdus={:?} dpdvs={:?} \
                         uv={:?} material={:?} area_light={:?} face_index={}",
                        item.pixel_index, item.depth,
                        item.p, item.n, item.ns,
                        item.dpdu, item.dpdv,
                        item.dpdus, item.dpdvs,
                        item.uv, item.material, item.area_light, item.face_index,
                    );
                }
            }
        });
    }

    fn intersect_shadow(
        &self,
        max_rays: usize,
        shadow_ray_q: &ShadowRayQueue,
        pixel_sample_state: &PixelSampleState,
    ) {
        let n_rays = shadow_ray_q.size().min(max_rays as u32);

        (0..n_rays as usize).into_par_iter().for_each(|i| {
            let work = unsafe { shadow_ray_q.get(i) };

            let ray = Ray::new(work.ray.o, work.ray.d, Some(work.ray.time), Ptr::null());
            if !self.aggregate.intersect_p(&ray, Some(work.t_max)) {
                let pi = work.pixel_index as usize;
                let ld = work.l_d / (work.r_u + work.r_l).average();
                let n = DIAG_SHADOW_UNOCCLUDED.fetch_add(1, Ordering::Relaxed);
                if n < 10 {
                    eprintln!(
                        "SHADOW_UNOCCLUDED[{n}] pixel={} l_d={:?} r_u={:?} r_l={:?} \
                         denom={:.6} ld={:?}",
                        pi, work.l_d, work.r_u, work.r_l,
                        (work.r_u + work.r_l).average(), ld
                    );
                }
                let mut l = pixel_sample_state.l.get(pi);
                l += ld;
                pixel_sample_state.l.set(pi, l);
            }
        });
    }

    fn intersect_shadow_tr(
        &self,
        max_rays: usize,
        shadow_ray_q: &ShadowRayQueue,
        pixel_sample_state: &PixelSampleState,
    ) {
        self.intersect_shadow(max_rays, shadow_ray_q, pixel_sample_state);
    }
}