move the dimensioned indexing out of spirv_backend and into coremem_types

this allows us to use it from the CPU implementation.
This commit is contained in:
2022-07-26 18:03:21 -07:00
parent d93d14d260
commit dbd666d272
4 changed files with 116 additions and 133 deletions

View File

@@ -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<R: Real, M: Material<R>>(
h: &[Vec3<R>],
m: &[Vec3<R>],
) {
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<R: Real, M: Material<R>>(
idx: Vec3u,
@@ -86,19 +95,32 @@ fn step_h_cell<R: Real, M: Material<R>>(
h: &mut [Vec3<R>],
m: &mut [Vec3<R>],
) {
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<F: FnMut(Vec3u)>(dim: Vec3u, mut f: F) {
@@ -111,51 +133,3 @@ fn apply_all_cells<F: FnMut(Vec3u)>(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<Vec3u> 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<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSamplePos<R> {
VolumeSamplePos::from_indexable(&DimIndexer::new(arr, dim), dim, idx)
}
fn sample_neg<R: Copy + Default>(arr: &[Vec3<R>], dim: Vec3u, idx: Vec3u) -> VolumeSampleNeg<R> {
VolumeSampleNeg::from_indexable(&DimIndexer::new(arr, dim), idx)
}
#[cfg(test)]
mod test {
#[test]
fn flat_idx_() {
// TODO: test all these helper methods!
unimplemented!()
}
}

View File

@@ -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<T> {
items: T,
@@ -45,62 +46,3 @@ impl<'a, T> IndexMut<usize> for SizedArray<&'a mut RuntimeArray<T>> {
}
}
pub struct DimensionedSlice<T> {
dim: Vec3u,
items: T,
}
impl<T> DimensionedSlice<T> {
pub fn new(dim: Vec3u, items: T) -> Self {
Self { dim, items }
}
}
impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> 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<usize> + ?Sized> Index<Vec3u> 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<usize> + ?Sized> IndexMut<Vec3u> 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);
}
}

66
crates/types/src/dim.rs Normal file
View File

@@ -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<T> {
dim: Vec3u,
items: T,
}
impl<T> DimensionedSlice<T> {
pub fn new(dim: Vec3u, items: T) -> Self {
Self { dim, items }
}
}
impl<'a, T: Index<usize> + ?Sized> Index<Vec3u> 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<usize> + ?Sized> Index<Vec3u> 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<usize> + ?Sized> IndexMut<Vec3u> 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);
}
}

View File

@@ -2,6 +2,7 @@
#![feature(core_intrinsics)]
pub mod compound;
pub mod dim;
pub mod mat;
pub mod real;
pub mod step;