665 lines
21 KiB
Rust
665 lines
21 KiB
Rust
use crate::core::image::Image;
|
|
use crate::utils::arena::Arena;
|
|
use crate::utils::backend::GpuAllocator;
|
|
use crate::utils::containers::Array2D;
|
|
use shared::core::geometry::{Bounds2f, Point2i, Vector2f, Vector2i};
|
|
use shared::utils::sampling::{
|
|
AliasTable, Bin, DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceSummedAreaTable,
|
|
DeviceWindowedPiecewiseConstant2D, PiecewiseLinear2D,
|
|
};
|
|
use shared::utils::{gpu_array_from_fn, Ptr};
|
|
use shared::Float;
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct PiecewiseConstant1D {
|
|
func: Vec<Float>,
|
|
cdf: Vec<Float>,
|
|
pub min: Float,
|
|
pub max: Float,
|
|
}
|
|
|
|
impl PiecewiseConstant1D {
|
|
pub fn new(f: &[Float]) -> Self {
|
|
Self::new_with_bounds(f.to_vec(), 0.0, 1.0)
|
|
}
|
|
|
|
pub fn from_func<F>(f: F, min: Float, max: Float, n: usize) -> Self
|
|
where
|
|
F: Fn(Float) -> Float,
|
|
{
|
|
let delta = (max - min) / n as Float;
|
|
let values: Vec<Float> = (0..n)
|
|
.map(|i| f(min + (i as Float + 0.5) * delta))
|
|
.collect();
|
|
Self::new_with_bounds(values, min, max)
|
|
}
|
|
|
|
pub fn new_with_bounds(f: Vec<Float>, min: Float, max: Float) -> Self {
|
|
let n = f.len();
|
|
let mut cdf = Vec::with_capacity(n + 1);
|
|
cdf.push(0.0);
|
|
|
|
let delta = (max - min) / n as Float;
|
|
for i in 0..n {
|
|
cdf.push(cdf[i] + f[i] * delta);
|
|
}
|
|
|
|
let func_integral = cdf[n];
|
|
if func_integral > 0.0 {
|
|
for c in &mut cdf {
|
|
*c /= func_integral;
|
|
}
|
|
}
|
|
|
|
Self { func: f, cdf, min, max }
|
|
}
|
|
|
|
// Accessors
|
|
pub fn n(&self) -> usize { self.func.len() }
|
|
pub fn func(&self) -> &[Float] { &self.func }
|
|
pub fn cdf(&self) -> &[Float] { &self.cdf }
|
|
|
|
pub fn integral(&self) -> Float {
|
|
// func_integral is the un-normalized sum. After normalization cdf[n] == 1.0,
|
|
// so we reconstruct from the last CDF entry before normalization.
|
|
// But since we normalized in-place, we need to store it. Let's compute it.
|
|
let n = self.func.len();
|
|
let delta = (self.max - self.min) / n as Float;
|
|
self.func.iter().sum::<Float>() * delta
|
|
}
|
|
|
|
/// Host-side sampling (for scene construction, not rendering).
|
|
/// During rendering, use the device struct via arena-uploaded Ptrs.
|
|
pub fn sample_host(&self, u: Float) -> (Float, Float, usize) {
|
|
let n = self.func.len();
|
|
let offset = self.find_interval_host(u);
|
|
let cdf_offset = self.cdf[offset];
|
|
let cdf_next = self.cdf[offset + 1];
|
|
let du = if cdf_next - cdf_offset > 0.0 {
|
|
(u - cdf_offset) / (cdf_next - cdf_offset)
|
|
} else {
|
|
0.0
|
|
};
|
|
let delta = (self.max - self.min) / n as Float;
|
|
let x = self.min + (offset as Float + du) * delta;
|
|
let func_integral = self.integral();
|
|
let pdf = if func_integral > 0.0 {
|
|
self.func[offset] / func_integral
|
|
} else {
|
|
0.0
|
|
};
|
|
(x, pdf, offset)
|
|
}
|
|
|
|
fn find_interval_host(&self, u: Float) -> usize {
|
|
let n = self.func.len();
|
|
let mut size = n;
|
|
let mut first = 0usize;
|
|
while size > 0 {
|
|
let half = size >> 1;
|
|
let middle = first + half;
|
|
if self.cdf[middle] <= u {
|
|
first = middle + 1;
|
|
size -= half + 1;
|
|
} else {
|
|
size = half;
|
|
}
|
|
}
|
|
first.saturating_sub(1).min(n - 1)
|
|
}
|
|
}
|
|
|
|
#[derive(DeviceRepr)]
|
|
#[device(name = "DevicePiecewiseConstant1D")]
|
|
pub struct PiecewiseConstant1D {
|
|
pub func: Vec<Float>,
|
|
pub cdf: Vec<Float>,
|
|
pub min: Float,
|
|
pub max: Float,
|
|
pub n: u32,
|
|
#[device(expr = "self.integral()")]
|
|
pub func_integral: Float,
|
|
}
|
|
|
|
#[derive(Clone, Debug, DeviceRepr)]
|
|
#[device(name = "DevicePiecewiseConstant2D")]
|
|
pub struct PiecewiseConstant2D {
|
|
pub conditionals: Vec<PiecewiseConstant1D>,
|
|
#[device(flatten)]
|
|
pub marginal: PiecewiseConstant1D,
|
|
pub n_u: u32,
|
|
pub n_v: u32,
|
|
}
|
|
|
|
|
|
impl PiecewiseConstant2D {
|
|
pub fn new(data: &Array2D<Float>) -> Self {
|
|
Self::new_with_bounds(data, Bounds2f::unit())
|
|
}
|
|
|
|
pub fn new_with_bounds(data: &Array2D<Float>, domain: Bounds2f) -> Self {
|
|
Self::from_slice(
|
|
data.as_slice(),
|
|
data.x_size(),
|
|
data.y_size(),
|
|
domain,
|
|
)
|
|
}
|
|
|
|
pub fn from_slice(data: &[Float], n_u: usize, n_v: usize, domain: Bounds2f) -> Self {
|
|
assert_eq!(data.len(), n_u * n_v);
|
|
|
|
let mut conditionals = Vec::with_capacity(n_v);
|
|
let mut marginal_func = Vec::with_capacity(n_v);
|
|
|
|
for v in 0..n_v {
|
|
let row = data[v * n_u..(v + 1) * n_u].to_vec();
|
|
let conditional = PiecewiseConstant1D::new_with_bounds(
|
|
row,
|
|
domain.p_min.x(),
|
|
domain.p_max.x(),
|
|
);
|
|
marginal_func.push(conditional.integral());
|
|
conditionals.push(conditional);
|
|
}
|
|
|
|
let marginal = PiecewiseConstant1D::new_with_bounds(
|
|
marginal_func,
|
|
domain.p_min.y(),
|
|
domain.p_max.y(),
|
|
);
|
|
|
|
Self { conditionals, marginal, n_u, n_v }
|
|
}
|
|
|
|
pub fn from_image(image: &Image) -> Self {
|
|
let res = image.resolution();
|
|
let n_u = res.x() as usize;
|
|
let n_v = res.y() as usize;
|
|
|
|
let mut data = Vec::with_capacity(n_u * n_v);
|
|
for v in 0..n_v {
|
|
for u in 0..n_u {
|
|
data.push(image.get_channels(Point2i::new(u as i32, v as i32)).average());
|
|
}
|
|
}
|
|
|
|
Self::from_slice(&data, n_u, n_v, Bounds2f::unit())
|
|
}
|
|
|
|
pub fn integral(&self) -> Float {
|
|
self.marginal.integral()
|
|
}
|
|
}
|
|
|
|
struct PiecewiseLinear2DStorage<const N: usize> {
|
|
data: Vec<Float>,
|
|
marginal_cdf: Vec<Float>,
|
|
conditional_cdf: Vec<Float>,
|
|
param_values: [Vec<Float>; N],
|
|
}
|
|
|
|
pub struct PiecewiseLinear2DHost<const N: usize> {
|
|
size: Vector2i,
|
|
inv_patch_size: Vector2f,
|
|
param_size: [u32; N],
|
|
param_strides: [u32; N],
|
|
storage: Arc<PiecewiseLinear2DStorage<N>>,
|
|
}
|
|
|
|
|
|
impl<const N: usize> PiecewiseLinear2DHost<N> {
|
|
pub fn new(
|
|
data: &[Float],
|
|
x_size: i32,
|
|
y_size: i32,
|
|
param_res: [usize; N],
|
|
param_values: [&[Float]; N],
|
|
normalize: bool,
|
|
build_cdf: bool,
|
|
) -> Self {
|
|
if build_cdf && !normalize {
|
|
panic!("PiecewiseLinear2D::new: build_cdf implies normalize=true");
|
|
}
|
|
|
|
let size = Vector2i::new(x_size, y_size);
|
|
let inv_patch_size = Vector2f::new(1. / (x_size - 1) as Float, 1. / (y_size - 1) as Float);
|
|
|
|
let mut param_size = [0u32; N];
|
|
let mut param_strides = [0u32; N];
|
|
let param_values = gpu_array_from_fn(|i| param_values[i].to_vec());
|
|
|
|
let mut slices: u32 = 1;
|
|
for i in (0..N).rev() {
|
|
if param_res[i] < 1 {
|
|
panic!("PiecewiseLinear2D::new: parameter resolution must be >= 1!");
|
|
}
|
|
param_size[i] = param_res[i] as u32;
|
|
param_strides[i] = if param_res[i] > 1 { slices } else { 0 };
|
|
slices *= param_size[i];
|
|
}
|
|
|
|
let n_values = (x_size * y_size) as usize;
|
|
let mut new_data = vec![0.0; slices as usize * n_values];
|
|
let mut marginal_cdf = if build_cdf {
|
|
vec![0.0; slices as usize * y_size as usize]
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
let mut conditional_cdf = if build_cdf {
|
|
vec![0.0; slices as usize * n_values]
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
let mut data_offset = 0;
|
|
for slice in 0..slices as usize {
|
|
let slice_offset = slice * n_values;
|
|
let current_data = &data[data_offset..data_offset + n_values];
|
|
let mut sum = 0.;
|
|
|
|
// Construct conditional CDF
|
|
if normalize {
|
|
for y in 0..(y_size - 1) {
|
|
for x in 0..(x_size - 1) {
|
|
let i = (y * x_size + x) as usize;
|
|
let v00 = current_data[i] as f64;
|
|
let v10 = current_data[i + 1] as f64;
|
|
let v01 = current_data[i + x_size as usize] as f64;
|
|
let v11 = current_data[i + 1 + x_size as usize] as f64;
|
|
sum += 0.25 * (v00 + v10 + v01 + v11);
|
|
}
|
|
}
|
|
}
|
|
|
|
let normalization = if normalize && sum > 0.0 {
|
|
1.0 / sum as Float
|
|
} else {
|
|
1.0
|
|
};
|
|
for k in 0..n_values {
|
|
new_data[slice_offset + k] = current_data[k] * normalization;
|
|
}
|
|
|
|
if build_cdf {
|
|
let marginal_slice_offset = slice * y_size as usize;
|
|
// Construct marginal CDF
|
|
for y in 0..y_size as usize {
|
|
let mut cdf_sum = 0.0;
|
|
let i_base = y * x_size as usize;
|
|
conditional_cdf[slice_offset + i_base] = 0.0;
|
|
for x in 0..(x_size - 1) as usize {
|
|
let i = i_base + x;
|
|
cdf_sum +=
|
|
0.5 * (new_data[slice_offset + i] + new_data[slice_offset + i + 1]);
|
|
conditional_cdf[slice_offset + i + 1] = cdf_sum;
|
|
}
|
|
}
|
|
// Construct marginal CDF
|
|
marginal_cdf[marginal_slice_offset] = 0.0;
|
|
let mut marginal_sum = 0.0;
|
|
for y in 0..(y_size - 1) as usize {
|
|
let cdf1 = conditional_cdf[slice_offset + (y + 1) * x_size as usize - 1];
|
|
let cdf2 = conditional_cdf[slice_offset + (y + 2) * x_size as usize - 1];
|
|
marginal_sum += 0.5 * (cdf1 + cdf2);
|
|
marginal_cdf[marginal_slice_offset + y + 1] = marginal_sum;
|
|
}
|
|
}
|
|
data_offset += n_values;
|
|
}
|
|
|
|
let view = PiecewiseLinear2D {
|
|
size,
|
|
inv_patch_size,
|
|
param_size,
|
|
param_strides,
|
|
param_values: param_values.each_ref().map(|x| x.as_ptr().into()),
|
|
data: Ptr::null(),
|
|
marginal_cdf: marginal_cdf.as_ptr().into(),
|
|
conditional_cdf: conditional_cdf.as_ptr().into(),
|
|
};
|
|
|
|
let storage = Arc::new(PiecewiseLinear2DStorage {
|
|
data: new_data,
|
|
marginal_cdf,
|
|
conditional_cdf,
|
|
param_values,
|
|
});
|
|
|
|
let mut final_view = view;
|
|
final_view.data = storage.data.as_ptr().into();
|
|
final_view.marginal_cdf = storage.marginal_cdf.as_ptr().into();
|
|
final_view.conditional_cdf = storage.conditional_cdf.as_ptr().into();
|
|
for i in 0..N {
|
|
final_view.param_values[i] = storage.param_values[i].as_ptr().into();
|
|
}
|
|
|
|
Self {
|
|
view: final_view,
|
|
_storage: storage,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<const N: usize> DeviceRepr for PiecewiseLinear2DHost<N> {
|
|
type Target = PiecewiseLinear2D<N>;
|
|
|
|
fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> PiecewiseLinear2D<N> {
|
|
let s = &self.storage;
|
|
|
|
let (data_ptr, _) = arena.alloc_slice(&s.data);
|
|
let (marginal_ptr, _) = arena.alloc_slice(&s.marginal_cdf);
|
|
let (conditional_ptr, _) = arena.alloc_slice(&s.conditional_cdf);
|
|
|
|
let param_ptrs: [Ptr<Float>; N] = std::array::from_fn(|i| {
|
|
let (ptr, _) = arena.alloc_slice(&s.param_values[i]);
|
|
ptr
|
|
});
|
|
|
|
PiecewiseLinear2D {
|
|
size: self.size,
|
|
inv_patch_size: self.inv_patch_size,
|
|
param_size: self.param_size,
|
|
param_strides: self.param_strides,
|
|
param_values: param_ptrs,
|
|
data: data_ptr,
|
|
marginal_cdf: marginal_ptr,
|
|
conditional_cdf: conditional_ptr,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct AliasTableHost {
|
|
bins: Vec<Bin>,
|
|
}
|
|
|
|
impl AliasTableHost {
|
|
pub fn new(weights: &[Float]) -> Self {
|
|
let n = weights.len();
|
|
if n == 0 {
|
|
return Self { bins: Vec::new() };
|
|
}
|
|
|
|
let sum: f64 = weights.iter().map(|&w| w as f64).sum();
|
|
assert!(sum > 0.0, "Sum of weights must be positive");
|
|
|
|
let mut bins = Vec::with_capacity(n);
|
|
for &w in weights {
|
|
bins.push(Bin {
|
|
p: (w as f64 / sum) as Float,
|
|
q: 0.0,
|
|
alias: 0,
|
|
});
|
|
}
|
|
|
|
struct Outcome { p_hat: f64, index: usize }
|
|
|
|
let mut under = Vec::with_capacity(n);
|
|
let mut over = Vec::with_capacity(n);
|
|
|
|
for (i, bin) in bins.iter().enumerate() {
|
|
let p_hat = (bin.p as f64) * (n as f64);
|
|
if p_hat < 1.0 {
|
|
under.push(Outcome { p_hat, index: i });
|
|
} else {
|
|
over.push(Outcome { p_hat, index: i });
|
|
}
|
|
}
|
|
|
|
while !under.is_empty() && !over.is_empty() {
|
|
let un = under.pop().unwrap();
|
|
let ov = over.pop().unwrap();
|
|
|
|
bins[un.index].q = un.p_hat as Float;
|
|
bins[un.index].alias = ov.index as u32;
|
|
|
|
let p_excess = un.p_hat + ov.p_hat - 1.0;
|
|
if p_excess < 1.0 {
|
|
under.push(Outcome { p_hat: p_excess, index: ov.index });
|
|
} else {
|
|
over.push(Outcome { p_hat: p_excess, index: ov.index });
|
|
}
|
|
}
|
|
|
|
while let Some(ov) = over.pop() {
|
|
bins[ov.index].q = 1.0;
|
|
bins[ov.index].alias = ov.index as u32;
|
|
}
|
|
while let Some(un) = under.pop() {
|
|
bins[un.index].q = 1.0;
|
|
bins[un.index].alias = un.index as u32;
|
|
}
|
|
|
|
Self { bins }
|
|
}
|
|
|
|
pub fn size(&self) -> usize { self.bins.len() }
|
|
pub fn is_empty(&self) -> bool { self.bins.is_empty() }
|
|
}
|
|
|
|
impl DeviceRepr for AliasTableHost {
|
|
type Target = AliasTable;
|
|
|
|
fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> AliasTable {
|
|
if self.bins.is_empty() {
|
|
return AliasTable { bins: Ptr::null(), size: 0 };
|
|
}
|
|
let (bins_ptr, _) = arena.alloc_slice(&self.bins);
|
|
AliasTable {
|
|
bins: bins_ptr,
|
|
size: self.bins.len() as u32,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, DeviceRepr)]
|
|
#[device(name = "DeviceSummedAreaTable")]
|
|
pub struct SummedAreaTable {
|
|
#[device(flatten)]
|
|
sum: Array2D<f64>,
|
|
}
|
|
|
|
impl SummedAreaTable {
|
|
pub fn new(values: &Array2D<Float>) -> Self {
|
|
let width = values.x_size() as i32;
|
|
let height = values.y_size() as i32;
|
|
|
|
let mut sum = Array2D::<f64>::new_dims(width, height);
|
|
sum[(0, 0)] = values[(0, 0)] as f64;
|
|
|
|
for x in 1..width {
|
|
sum[(x, 0)] = values[(x, 0)] as f64 + sum[(x - 1, 0)];
|
|
}
|
|
for y in 1..height {
|
|
sum[(0, y)] = values[(0, y)] as f64 + sum[(0, y - 1)];
|
|
}
|
|
for y in 1..height {
|
|
for x in 1..width {
|
|
sum[(x, y)] = values[(x, y)] as f64
|
|
+ sum[(x - 1, y)]
|
|
+ sum[(x, y - 1)]
|
|
- sum[(x - 1, y - 1)];
|
|
}
|
|
}
|
|
|
|
Self { sum }
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, DeviceRepr)]
|
|
#[device(name = "DeviceWindowedPiecewiseConstant2D")]
|
|
pub struct WindowedPiecewiseConstant2D {
|
|
#[device(flatten)]
|
|
sat: SummedAreaTable,
|
|
#[device(flatten)]
|
|
func: Array2D<Float>,
|
|
}
|
|
|
|
impl WindowedPiecewiseConstant2D {
|
|
pub fn new(func: Array2D<Float>) -> Self {
|
|
let sat = SummedAreaTable::new(&func);
|
|
Self { sat, func }
|
|
}
|
|
}
|
|
|
|
struct PiecewiseLinear2DStorage<const N: usize> {
|
|
data: Vec<Float>,
|
|
marginal_cdf: Vec<Float>,
|
|
conditional_cdf: Vec<Float>,
|
|
param_values: [Vec<Float>; N],
|
|
}
|
|
|
|
pub struct PiecewiseLinear2DHost<const N: usize> {
|
|
size: Vector2i,
|
|
inv_patch_size: Vector2f,
|
|
param_size: [u32; N],
|
|
param_strides: [u32; N],
|
|
storage: Arc<PiecewiseLinear2DStorage<N>>,
|
|
}
|
|
|
|
impl<const N: usize> PiecewiseLinear2DHost<N> {
|
|
pub fn new(
|
|
data: &[Float],
|
|
x_size: i32,
|
|
y_size: i32,
|
|
param_res: [usize; N],
|
|
param_values: [&[Float]; N],
|
|
normalize: bool,
|
|
build_cdf: bool,
|
|
) -> Self {
|
|
if build_cdf && !normalize {
|
|
panic!("PiecewiseLinear2D: build_cdf implies normalize=true");
|
|
}
|
|
|
|
let size = Vector2i::new(x_size, y_size);
|
|
let inv_patch_size = Vector2f::new(
|
|
1.0 / (x_size - 1) as Float,
|
|
1.0 / (y_size - 1) as Float,
|
|
);
|
|
|
|
let mut param_size = [0u32; N];
|
|
let mut param_strides = [0u32; N];
|
|
let owned_param_values: [Vec<Float>; N] = gpu_array_from_fn(|i| param_values[i].to_vec());
|
|
|
|
let mut slices: u32 = 1;
|
|
for i in (0..N).rev() {
|
|
assert!(param_res[i] >= 1, "Parameter resolution must be >= 1");
|
|
param_size[i] = param_res[i] as u32;
|
|
param_strides[i] = if param_res[i] > 1 { slices } else { 0 };
|
|
slices *= param_size[i];
|
|
}
|
|
|
|
let n_values = (x_size * y_size) as usize;
|
|
let mut new_data = vec![0.0; slices as usize * n_values];
|
|
let mut marginal_cdf = if build_cdf {
|
|
vec![0.0; slices as usize * y_size as usize]
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
let mut conditional_cdf = if build_cdf {
|
|
vec![0.0; slices as usize * n_values]
|
|
} else {
|
|
Vec::new()
|
|
};
|
|
|
|
let mut data_offset = 0;
|
|
for slice in 0..slices as usize {
|
|
let slice_offset = slice * n_values;
|
|
let current_data = &data[data_offset..data_offset + n_values];
|
|
let mut sum = 0.0_f64;
|
|
|
|
if normalize {
|
|
for y in 0..(y_size - 1) {
|
|
for x in 0..(x_size - 1) {
|
|
let i = (y * x_size + x) as usize;
|
|
let v00 = current_data[i] as f64;
|
|
let v10 = current_data[i + 1] as f64;
|
|
let v01 = current_data[i + x_size as usize] as f64;
|
|
let v11 = current_data[i + 1 + x_size as usize] as f64;
|
|
sum += 0.25 * (v00 + v10 + v01 + v11);
|
|
}
|
|
}
|
|
}
|
|
|
|
let normalization = if normalize && sum > 0.0 {
|
|
1.0 / sum as Float
|
|
} else {
|
|
1.0
|
|
};
|
|
for k in 0..n_values {
|
|
new_data[slice_offset + k] = current_data[k] * normalization;
|
|
}
|
|
|
|
if build_cdf {
|
|
let marginal_slice_offset = slice * y_size as usize;
|
|
for y in 0..y_size as usize {
|
|
let mut cdf_sum = 0.0;
|
|
let i_base = y * x_size as usize;
|
|
conditional_cdf[slice_offset + i_base] = 0.0;
|
|
for x in 0..(x_size - 1) as usize {
|
|
let i = i_base + x;
|
|
cdf_sum += 0.5
|
|
* (new_data[slice_offset + i] + new_data[slice_offset + i + 1]);
|
|
conditional_cdf[slice_offset + i + 1] = cdf_sum;
|
|
}
|
|
}
|
|
marginal_cdf[marginal_slice_offset] = 0.0;
|
|
let mut marginal_sum = 0.0;
|
|
for y in 0..(y_size - 1) as usize {
|
|
let cdf1 =
|
|
conditional_cdf[slice_offset + (y + 1) * x_size as usize - 1];
|
|
let cdf2 =
|
|
conditional_cdf[slice_offset + (y + 2) * x_size as usize - 1];
|
|
marginal_sum += 0.5 * (cdf1 + cdf2);
|
|
marginal_cdf[marginal_slice_offset + y + 1] = marginal_sum;
|
|
}
|
|
}
|
|
data_offset += n_values;
|
|
}
|
|
|
|
let storage = Arc::new(PiecewiseLinear2DStorage {
|
|
data: new_data,
|
|
marginal_cdf,
|
|
conditional_cdf,
|
|
param_values: owned_param_values,
|
|
});
|
|
|
|
Self {
|
|
size,
|
|
inv_patch_size,
|
|
param_size,
|
|
param_strides,
|
|
storage,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<const N: usize> DeviceRepr for PiecewiseLinear2DHost<N> {
|
|
type Target = PiecewiseLinear2D<N>;
|
|
|
|
fn upload_value<A: GpuAllocator>(&self, arena: &Arena<A>) -> PiecewiseLinear2D<N> {
|
|
let s = &self.storage;
|
|
|
|
let (data_ptr, _) = arena.alloc_slice(&s.data);
|
|
let (marginal_ptr, _) = arena.alloc_slice(&s.marginal_cdf);
|
|
let (conditional_ptr, _) = arena.alloc_slice(&s.conditional_cdf);
|
|
|
|
let param_ptrs: [Ptr<Float>; N] = std::array::from_fn(|i| {
|
|
let (ptr, _) = arena.alloc_slice(&s.param_values[i]);
|
|
ptr
|
|
});
|
|
|
|
PiecewiseLinear2D {
|
|
size: self.size,
|
|
inv_patch_size: self.inv_patch_size,
|
|
param_size: self.param_size,
|
|
param_strides: self.param_strides,
|
|
param_values: param_ptrs,
|
|
data: data_ptr,
|
|
marginal_cdf: marginal_ptr,
|
|
conditional_cdf: conditional_ptr,
|
|
}
|
|
}
|
|
}
|