pbrt/src/utils/arena.rs

use crate::core::image::Image;
use crate::core::texture::FloatTexture;
use crate::utils::sampling::PiecewiseConstant2D;
use shared::Float;
use shared::core::color::RGBToSpectrumTable;
use shared::core::image::DeviceImage;
use shared::core::shape::Shape;
use shared::core::texture::GPUFloatTexture;
use shared::spectra::{RGBColorSpace, StandardColorSpaces};
use shared::textures::*;
use shared::utils::Ptr;
use shared::utils::sampling::DevicePiecewiseConstant2D;
use std::alloc::Layout;

pub struct Arena {
    buffer: Vec<(*mut u8, Layout)>,
}

impl Arena {
    pub fn new() -> Self {
        Self { buffer: Vec::new() }
    }

    pub fn alloc<T>(&mut self, value: T) -> Ptr<T> {
        let layout = Layout::new::<T>();

        let ptr = unsafe { self.alloc_unified(layout) } as *mut T;

        // Write the value
        unsafe {
            ptr.write(value);
        }

        Ptr::from_raw(ptr)
    }

    pub fn alloc_opt<T>(&mut self, value: Option<T>) -> Ptr<T> {
        match value {
            Some(v) => self.alloc(v),
            None => Ptr::null(),
        }
    }

    pub fn alloc_slice<T: Copy>(&mut self, values: &[T]) -> (Ptr<T>, usize) {
        if values.is_empty() {
            return (Ptr::null(), 0);
        }

        let layout = Layout::array::<T>(values.len()).unwrap();
        let ptr = unsafe { self.alloc_unified(layout) } as *mut T;

        unsafe {
            std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len());
        }

        (Ptr::from_raw(ptr), values.len())
    }

    #[cfg(feature = "cuda")]
    unsafe fn alloc_unified(&mut self, layout: Layout) -> *mut u8 {
        use cuda_runtime_sys::*;

        let mut ptr: *mut std::ffi::c_void = std::ptr::null_mut();
        let size = layout.size().max(layout.align());

        let result = cudaMallocManaged(&mut ptr, size, cudaMemAttachGlobal);

        if result != cudaError::cudaSuccess {
            panic!("cudaMallocManaged failed: {:?}", result);
        }

        self.allocations.push((ptr as *mut u8, layout));
        ptr as *mut u8
    }

    #[cfg(not(feature = "cuda"))]
    unsafe fn alloc_unified(&mut self, layout: Layout) -> *mut u8 {
        // Fallback: regular allocation for CPU-only testing
        let ptr = std::alloc::alloc(layout);
        self.allocations.push((ptr, layout));
        ptr
    }

    pub fn raw_data(&self) -> &[u8] {
        &self.buffer
    }
}

impl Drop for UnifiedArena {
    fn drop(&mut self) {
        for (ptr, layout) in self.allocations.drain(..) {
            unsafe {
                #[cfg(feature = "cuda")]
                {
                    cuda_runtime_sys::cudaFree(ptr as *mut _);
                }
                #[cfg(not(feature = "cuda"))]
                {
                    std::alloc::dealloc(ptr, layout);
                }
            }
        }
    }
}

pub trait Upload {
    type Target: Copy;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target>;
}

impl Upload for Shape {
    type Target = Shape;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        arena.alloc(self.clone())
    }
}

impl Upload for Image {
    type Target = Image;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let pixels_ptr = arena.alloc_slice(&self.storage_as_slice());
        let device_img = Image {
            base: self.base,
            pixels: pixels_ptr,
        };
        arena.alloc(device_img)
    }
}

impl Upload for FloatTexture {
    type Target = GPUFloatTexture;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let gpu_variant = match self {
            FloatTexture::Constant(tex) => GPUFloatTexture::Constant(tex.clone()),
            FloatTexture::Checkerboard(tex) => GPUFloatTexture::Checkerboard(tex.clone()),
            FloatTexture::Dots(tex) => GPUFloatTexture::Dots(tex.clone()),
            FloatTexture::FBm(tex) => GPUFloatTexture::FBm(tex.clone()),
            FloatTexture::Windy(tex) => GPUFloatTexture::Windy(tex.clone()),
            FloatTexture::Wrinkled(tex) => GPUFloatTexture::Wrinkled(tex.clone()),
            FloatTexture::Constant(val) => GPUFloatTexture::Constant(*val),
            FloatTexture::Scaled(tex) => {
                let child_ptr = tex.texture.upload(arena);

                let gpu_scaled = GPUFloatScaledTexture {
                    tex: child_ptr,
                    scale: tex.scale.upload(arena),
                };
                GPUFloatTexture::Scaled(gpu_scaled)
            }

            FloatTexture::Mix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);
                let amount_ptr = tex.amount.upload(arena);

                let gpu_mix = GPUFloatMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    amount: amount_ptr,
                };
                GPUFloatTexture::Mix(gpu_mix)
            }

            FloatTexture::DirectionMix(tex) => {
                let tex1_ptr = tex.tex1.upload(arena);
                let tex2_ptr = tex.tex2.upload(arena);
                let gpu_dmix = shared::textures::GPUFloatDirectionMixTexture {
                    tex1: tex1_ptr,
                    tex2: tex2_ptr,
                    dir: tex.dir,
                };
                GPUFloatTexture::DirectionMix(gpu_dmix)
            }
            FloatTexture::Image(tex) => {
                let image_ptr = tex.image.upload(arena);

                let gpu_image_tex = GPUFloatImageTexture {
                    mapping: tex.mapping,
                    tex_obj: image_ptr.offset as u64,
                    scale: tex.scale,
                    invert: tex.invert,
                    mapping: tex.mapping,
                };
                GPUFloatTexture::Image(gpu_image_tex)
            }

            FloatTexture::Ptex(tex) => {
                todo!("Implement Ptex buffer upload")
            }

            FloatTexture::Bilerp(tex) => GPUFloatTexture::Bilerp(tex.clone()),
        };

        arena.alloc(gpu_variant)
    }
}

impl Upload for RGBToSpectrumTable {
    type Target = RGBToSpectrumTable;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let z_ptr = arena.alloc_slice(&self.z_nodes);
        let c_ptr = arena.alloc_slice(&self.coeffs);

        let shared_table = RGBToSpectrumTable {
            z_nodes: z_ptr,
            coeffs: c_ptr,
        };

        arena.alloc(shared_table)
    }
}

impl Upload for RGBColorSpace {
    type Target = RGBColorSpace;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let table_ptr = self.rgb_to_spectrum_table.upload(arena);

        let shared_space = RGBColorSpace {
            r: self.r,
            g: self.g,
            b: self.b,
            w: self.w,
            illuminant: self.illuminant.clone(),
            rgb_to_spectrum_table: table_ptr,
            xyz_from_rgb: self.xyz_from_rgb,
            rgb_from_xyz: self.rgb_from_xyz,
        };

        arena.alloc(shared_space)
    }
}

impl Upload for StandardColorSpaces {
    type Target = StandardColorSpaces;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let srgb_ptr = self.srgb.upload(arena);
        let dci_ptr = self.dci_p3.upload(arena);
        let rec_ptr = self.rec2020.upload(arena);
        let aces_ptr = self.aces2065_1.upload(arena);

        let registry = StandardColorSpaces {
            srgb: srgb_ptr,
            dci_p3: dci_ptr,
            rec2020: rec_ptr,
            aces2065_1: aces_ptr,
        };

        arena.alloc(registry)
    }
}

impl Upload for PiecewiseConstant2D {
    type Target = DevicePiecewiseConstant2D;

    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        let marginal_shared = self.p_marginal.to_shared(arena);

        let conditionals_shared: Vec<DevicePiecewiseConstant1D> = self
            .p_conditionals
            .iter()
            .map(|c| c.to_shared(arena))
            .collect();

        let conditionals_ptr = arena.alloc_slice(&conditionals_shared);

        let shared_2d = DevicePiecewiseConstant2D {
            domain: self.domain,
            p_marginal: marginal_shared,
            n_conditionals: self.p_conditionals.len(),
            p_conditional_v: conditionals_ptr,
        };

        arena.alloc(shared_2d)
    }
}

impl<T: Upload> Upload for Option<T> {
    type Target = T::Target;
    fn upload(&self, arena: &mut Arena) -> Ptr<Self::Target> {
        match self {
            Some(val) => val.upload(arena),
            None => Ptr::null(),
        }
    }
}