diff --git a/crates/coremem/src/sim/spirv/cpu.rs b/crates/coremem/src/sim/spirv/cpu.rs index 0c4cb21..99f7516 100644 --- a/crates/coremem/src/sim/spirv/cpu.rs +++ b/crates/coremem/src/sim/spirv/cpu.rs @@ -1,5 +1,4 @@ -use std::ops::Index; - +use coremem_types::dim::DimensionedSlice; use coremem_types::mat::Material; use coremem_types::real::Real; use coremem_types::step::{SimMeta, StepEContext, StepHContext, VolumeSampleNeg, VolumeSamplePos}; @@ -65,17 +64,27 @@ fn step_e_cell>( h: &[Vec3], m: &[Vec3], ) { - let flat_idx = flat_idx(meta.dim, idx); - let step_e_context = StepEContext { - inv_feature_size: meta.inv_feature_size(), - time_step: meta.time_step(), - stim_e: stim_e[flat_idx].cast(), - mat: &mat[flat_idx], - in_h: sample_neg(h, meta.dim(), idx), - in_e: e[flat_idx], + let dim = meta.dim; + let stim_e_matrix = DimensionedSlice::new(dim, stim_e); + let mat_matrix = DimensionedSlice::new(dim, mat); + let mut e_matrix = DimensionedSlice::new(dim, e); + let h_matrix = DimensionedSlice::new(dim, h); + + let stim_e = stim_e_matrix[idx]; + let mat = &mat_matrix[idx]; + let in_e = e_matrix[idx]; + let in_h = VolumeSampleNeg::from_indexable(&h_matrix, idx); + + let update_state = StepEContext { + inv_feature_size: meta.inv_feature_size, + time_step: meta.time_step, + stim_e: stim_e.cast(), + mat, + in_h, + in_e, }; - let new_e = step_e_context.step_e(); - e[flat_idx] = new_e; + let new_e = update_state.step_e(); + e_matrix[idx] = new_e; } fn step_h_cell>( idx: Vec3u, @@ -86,19 +95,32 @@ fn step_h_cell>( h: &mut [Vec3], m: &mut [Vec3], ) { - let flat_idx = flat_idx(meta.dim, idx); - let step_h_context = StepHContext { - inv_feature_size: meta.inv_feature_size(), - time_step: meta.time_step(), - stim_h: stim_h[flat_idx].cast(), - mat: &mat[flat_idx], - in_e: sample_pos(e, meta.dim(), idx), - in_h: h[flat_idx], - in_m: m[flat_idx], + let dim = meta.dim; + + let stim_h_matrix = DimensionedSlice::new(dim, stim_h); + let mat_matrix = DimensionedSlice::new(dim, mat); + let e_matrix = DimensionedSlice::new(dim, e); + let mut h_matrix = DimensionedSlice::new(dim, h); + let mut m_matrix = DimensionedSlice::new(dim, m); + + let stim_h = stim_h_matrix[idx]; + let mat = &mat_matrix[idx]; + let in_e = VolumeSamplePos::from_indexable(&e_matrix, dim, idx); + let in_h = h_matrix[idx]; + let in_m = m_matrix[idx]; + + let update_state = StepHContext { + inv_feature_size: meta.inv_feature_size, + time_step: meta.time_step, + stim_h: stim_h.cast(), + mat, + in_e, + in_h, + in_m, }; - let (new_h, new_m) = step_h_context.step_h(); - h[flat_idx] = new_h; - m[flat_idx] = new_m; + let (new_h, new_m) = update_state.step_h(); + h_matrix[idx] = new_h; + m_matrix[idx] = new_m; } fn apply_all_cells(dim: Vec3u, mut f: F) { @@ -111,51 +133,3 @@ fn apply_all_cells(dim: Vec3u, mut f: F) { } } -/// provides (x,y,z)-based indexing into a flat memory view -struct DimIndexer<'a, T> { - items: &'a [T], - dim: Vec3u, -} - -impl<'a, T> DimIndexer<'a, T> { - fn new(items: &'a [T], dim: Vec3u) -> Self { - Self { items, dim } - } -} - -impl<'a, T> Index for DimIndexer<'a, T> { - type Output = T; - fn index(&self, idx: Vec3u) -> &Self::Output { - let idx = flat_idx(self.dim, idx); - &self.items[idx] - } -} - -fn flat_idx(dim: Vec3u, idx: Vec3u) -> usize { - let dx = dim.x() as usize; - let dy = dim.y() as usize; - // let dz = dim.z() as usize; - let ix = idx.x() as usize; - let iy = idx.y() as usize; - let iz = idx.z() as usize; - let effective_y = iz * dy + iy; - let effective_x = effective_y * dx + ix; - effective_x -} - -fn sample_pos(arr: &[Vec3], dim: Vec3u, idx: Vec3u) -> VolumeSamplePos { - VolumeSamplePos::from_indexable(&DimIndexer::new(arr, dim), dim, idx) -} - -fn sample_neg(arr: &[Vec3], dim: Vec3u, idx: Vec3u) -> VolumeSampleNeg { - VolumeSampleNeg::from_indexable(&DimIndexer::new(arr, dim), idx) -} - -#[cfg(test)] -mod test { - #[test] - fn flat_idx_() { - // TODO: test all these helper methods! - unimplemented!() - } -} diff --git a/crates/spirv_backend/src/support.rs b/crates/spirv_backend/src/support.rs index 57c140e..3d38fa0 100644 --- a/crates/spirv_backend/src/support.rs +++ b/crates/spirv_backend/src/support.rs @@ -1,8 +1,9 @@ use core::ops::{Index, IndexMut}; -use spirv_std::RuntimeArray; +// TODO: remove this re-export +pub use coremem_types::dim::DimensionedSlice; -use coremem_types::vec::Vec3u; +use spirv_std::RuntimeArray; pub struct SizedArray { items: T, @@ -45,62 +46,3 @@ impl<'a, T> IndexMut for SizedArray<&'a mut RuntimeArray> { } } -pub struct DimensionedSlice { - dim: Vec3u, - items: T, -} - -impl DimensionedSlice { - pub fn new(dim: Vec3u, items: T) -> Self { - Self { dim, items } - } -} - -impl<'a, T: Index + ?Sized> Index for DimensionedSlice<&'a T> { - type Output=T::Output; - - fn index(&self, idx: Vec3u) -> &Self::Output { - let idx = index(idx, self.dim); - &self.items[idx] - } -} - -impl<'a, T: Index + ?Sized> Index for DimensionedSlice<&'a mut T> { - type Output=T::Output; - - fn index(&self, idx: Vec3u) -> &Self::Output { - let idx = index(idx, self.dim); - &self.items[idx] - } -} - -impl<'a, T: IndexMut + ?Sized> IndexMut for DimensionedSlice<&'a mut T> { - fn index_mut(&mut self, idx: Vec3u) -> &mut Self::Output { - let idx = index(idx, self.dim); - &mut self.items[idx] - } -} - -fn index(loc: Vec3u, dim: Vec3u) -> usize { - ((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize -} - - -#[cfg(test)] -mod test { - use super::*; - #[test] - fn test_index() { - let dim = Vec3u::new(2, 3, 7); - assert_eq!(index(Vec3u::new(0, 0, 0), dim), 0); - assert_eq!(index(Vec3u::new(1, 0, 0), dim), 1); - assert_eq!(index(Vec3u::new(0, 1, 0), dim), 2); - assert_eq!(index(Vec3u::new(1, 1, 0), dim), 3); - assert_eq!(index(Vec3u::new(0, 2, 0), dim), 4); - assert_eq!(index(Vec3u::new(0, 0, 1), dim), 6); - assert_eq!(index(Vec3u::new(1, 0, 1), dim), 7); - assert_eq!(index(Vec3u::new(0, 1, 1), dim), 8); - assert_eq!(index(Vec3u::new(1, 2, 1), dim), 11); - assert_eq!(index(Vec3u::new(1, 2, 2), dim), 17); - } -} diff --git a/crates/types/src/dim.rs b/crates/types/src/dim.rs new file mode 100644 index 0000000..0d78e19 --- /dev/null +++ b/crates/types/src/dim.rs @@ -0,0 +1,66 @@ +use core::ops::{Index, IndexMut}; + +use crate::vec::Vec3u; + +/// use this to wrap a flat region of memory into something which can be indexed by coordinates in +/// 3d space. +pub struct DimensionedSlice { + dim: Vec3u, + items: T, +} + +impl DimensionedSlice { + pub fn new(dim: Vec3u, items: T) -> Self { + Self { dim, items } + } +} + +impl<'a, T: Index + ?Sized> Index for DimensionedSlice<&'a T> { + type Output=T::Output; + + fn index(&self, idx: Vec3u) -> &Self::Output { + let idx = index(idx, self.dim); + &self.items[idx] + } +} + +impl<'a, T: Index + ?Sized> Index for DimensionedSlice<&'a mut T> { + type Output=T::Output; + + fn index(&self, idx: Vec3u) -> &Self::Output { + let idx = index(idx, self.dim); + &self.items[idx] + } +} + +impl<'a, T: IndexMut + ?Sized> IndexMut for DimensionedSlice<&'a mut T> { + fn index_mut(&mut self, idx: Vec3u) -> &mut Self::Output { + let idx = index(idx, self.dim); + &mut self.items[idx] + } +} + +fn index(loc: Vec3u, dim: Vec3u) -> usize { + ((loc.z()*dim.y() + loc.y())*dim.x() + loc.x()) as usize +} + + +#[cfg(test)] +mod test { + use super::*; + #[test] + fn test_index() { + let dim = Vec3u::new(2, 3, 7); + assert_eq!(index(Vec3u::new(0, 0, 0), dim), 0); + assert_eq!(index(Vec3u::new(1, 0, 0), dim), 1); + assert_eq!(index(Vec3u::new(0, 1, 0), dim), 2); + assert_eq!(index(Vec3u::new(1, 1, 0), dim), 3); + assert_eq!(index(Vec3u::new(0, 2, 0), dim), 4); + assert_eq!(index(Vec3u::new(0, 0, 1), dim), 6); + assert_eq!(index(Vec3u::new(1, 0, 1), dim), 7); + assert_eq!(index(Vec3u::new(0, 1, 1), dim), 8); + assert_eq!(index(Vec3u::new(1, 2, 1), dim), 11); + assert_eq!(index(Vec3u::new(1, 2, 2), dim), 17); + } +} + diff --git a/crates/types/src/lib.rs b/crates/types/src/lib.rs index 74ee7b6..c51c4ed 100644 --- a/crates/types/src/lib.rs +++ b/crates/types/src/lib.rs @@ -2,6 +2,7 @@ #![feature(core_intrinsics)] pub mod compound; +pub mod dim; pub mod mat; pub mod real; pub mod step;