spirv_backend: use RuntimeArray to remove all this UnsizedArray stuff

crates/spirv_backend/src/adapt.rs

@@ -1,159 +1,101 @@
+use spirv_std::RuntimeArray;
+
 use coremem_types::step::{StepEContext, StepHContext, VolumeSampleNeg, VolumeSamplePos};
-use crate::support::{Array3, Array3Mut, UnsizedArray};
 use coremem_types::mat::Material;
 use coremem_types::real::Real;
 use coremem_types::step::SimMeta;
 use coremem_types::vec::{Vec3, Vec3u};
 
+use crate::support::{DimensionedSlice, SizedArray};
+
 pub(crate) fn step_h<R: Real, M: Material<R>>(
-    id: Vec3u,
+    idx: Vec3u,
     meta: &SimMeta<R>,
-    stimulus_h: &UnsizedArray<Vec3<R>>,
-    material: &UnsizedArray<M>,
-    e: &UnsizedArray<Vec3<R>>,
-    h: &mut UnsizedArray<Vec3<R>>,
-    m: &mut UnsizedArray<Vec3<R>>,
+    stimulus_h: &RuntimeArray<Vec3<R>>,
+    material: &RuntimeArray<M>,
+    e: &RuntimeArray<Vec3<R>>,
+    h: &mut RuntimeArray<Vec3<R>>,
+    m: &mut RuntimeArray<Vec3<R>>,
 ) {
-    if id.x() < meta.dim.x() && id.y() < meta.dim.y() && id.z() < meta.dim.z() {
-        let sim_state = SerializedStepH::new(meta, stimulus_h, material, e, h, m);
-        let update_state = sim_state.index(id);
-        let (new_h, new_m) = update_state.step_h();
-        sim_state.write_output(id, new_h, new_m);
-    }
-}
+    if idx.x() < meta.dim.x() && idx.y() < meta.dim.y() && idx.z() < meta.dim.z() {
+        let dim = meta.dim;
+        let len = dim.product_sum_usize();
 
-pub(crate) fn step_e<R: Real, M: Material<R>>(
-    id: Vec3u,
-    meta: &SimMeta<R>,
-    stimulus_e: &UnsizedArray<Vec3<R>>,
-    material: &UnsizedArray<M>,
-    e: &mut UnsizedArray<Vec3<R>>,
-    h: &UnsizedArray<Vec3<R>>,
-) {
-    if id.x() < meta.dim.x() && id.y() < meta.dim.y() && id.z() < meta.dim.z() {
-        let sim_state = SerializedStepE::new(meta, stimulus_e, material, e, h);
+        let stim_h_array = unsafe { SizedArray::new(stimulus_h, len) };
+        let stim_h_matrix = DimensionedSlice::new(dim, &stim_h_array);
 
-        let update_state = sim_state.index(id);
-        let new_e = update_state.step_e();
-        sim_state.write_output(id, new_e);
-    }
-}
+        let mat_array = unsafe { SizedArray::new(material, len) };
+        let mat_matrix = DimensionedSlice::new(dim, &mat_array);
 
+        let e_array = unsafe { SizedArray::new(e, len) };
+        let e_matrix = DimensionedSlice::new(dim, &e_array);
 
-/// Whatever data we received from the host in their call to step_h
-struct SerializedStepH<'a, R, M> {
-    meta: &'a SimMeta<R>,
-    stimulus_h: &'a UnsizedArray<Vec3<R>>,
-    material: &'a UnsizedArray<M>,
-    e: &'a UnsizedArray<Vec3<R>>,
-    h: &'a mut UnsizedArray<Vec3<R>>,
-    m: &'a mut UnsizedArray<Vec3<R>>,
-}
+        let mut h_array = unsafe { SizedArray::new(h, len) };
+        let mut h_matrix = DimensionedSlice::new(dim, &mut h_array);
 
-impl<'a, R, M> SerializedStepH<'a, R, M> {
-    fn new(
-        meta: &'a SimMeta<R>,
-        stimulus_h: &'a UnsizedArray<Vec3<R>>,
-        material: &'a UnsizedArray<M>,
-        e: &'a UnsizedArray<Vec3<R>>,
-        h: &'a mut UnsizedArray<Vec3<R>>,
-        m: &'a mut UnsizedArray<Vec3<R>>,
-    ) -> Self {
-        Self {
-            meta, stimulus_h, material, e, h, m
-        }
-    }
-}
+        let mut m_array = unsafe { SizedArray::new(m, len) };
+        let mut m_matrix = DimensionedSlice::new(dim, &mut m_array);
 
-impl<'a, R: Real, M> SerializedStepH<'a, R, M> {
-    /// returns a context which the user can call `step_h` on
-    fn index(&self, idx: Vec3u) -> StepHContext<'a, R, M> {
-        let dim = self.meta.dim;
-        let stim_h_matrix = Array3::new(self.stimulus_h, dim);
-        let mat_matrix = Array3::new(self.material, dim);
-        let e = Array3::new(self.e, dim);
-        let h = Array3::new(self.h, dim);
-        let m = Array3::new(self.m, dim);
+        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 in_e = VolumeSamplePos::from_indexable(e, dim, idx);
-        let in_h = h.get(idx).unwrap();
-        let in_m = m.get(idx).unwrap();
-
-        let mat = mat_matrix.into_ref(idx);
-
-        StepHContext {
-            inv_feature_size: self.meta.inv_feature_size,
-            time_step: self.meta.time_step,
-            stim_h: stim_h_matrix.get(idx).unwrap(),
+        let update_state = StepHContext {
+            inv_feature_size: meta.inv_feature_size,
+            time_step: meta.time_step,
+            stim_h,
             mat,
             in_e,
             in_h,
             in_m,
-        }
-    }
-
-    fn write_output(self, idx: Vec3u, h: Vec3<R>, m: Vec3<R>) {
-        let dim = self.meta.dim;
-        let arr_h = Array3Mut::new(self.h, dim);
-        let arr_m = Array3Mut::new(self.m, dim);
-        let mut out_h = arr_h.into_mut_handle(idx);
-        let mut out_m = arr_m.into_mut_handle(idx);
-        out_h.write(h);
-        out_m.write(m);
+        };
+        let (new_h, new_m) = update_state.step_h();
+        h_matrix[idx] = new_h;
+        m_matrix[idx] = new_m;
     }
 }
 
-/// Whatever data we received from the host in their call to step_e
-struct SerializedStepE<'a, R, M> {
-    meta: &'a SimMeta<R>,
-    stimulus_e: &'a UnsizedArray<Vec3<R>>,
-    material: &'a UnsizedArray<M>,
-    e: &'a mut UnsizedArray<Vec3<R>>,
-    h: &'a UnsizedArray<Vec3<R>>,
-}
+pub(crate) fn step_e<R: Real, M: Material<R>>(
+    idx: Vec3u,
+    meta: &SimMeta<R>,
+    stimulus_e: &RuntimeArray<Vec3<R>>,
+    material: &RuntimeArray<M>,
+    e: &mut RuntimeArray<Vec3<R>>,
+    h: &RuntimeArray<Vec3<R>>,
+) {
+    if idx.x() < meta.dim.x() && idx.y() < meta.dim.y() && idx.z() < meta.dim.z() {
+        let dim = meta.dim;
+        let len = dim.product_sum_usize();
 
-impl<'a, R, M> SerializedStepE<'a, R, M> {
-    pub fn new(
-        meta: &'a SimMeta<R>,
-        stimulus_e: &'a UnsizedArray<Vec3<R>>,
-        material: &'a UnsizedArray<M>,
-        e: &'a mut UnsizedArray<Vec3<R>>,
-        h: &'a UnsizedArray<Vec3<R>>,
-    ) -> Self {
-        Self {
-            meta, stimulus_e, material, e, h
-        }
-    }
-}
+        let stim_e_array = unsafe { SizedArray::new(stimulus_e, len) };
+        let stim_e_matrix = DimensionedSlice::new(dim, &stim_e_array);
 
-impl<'a, R: Real, M> SerializedStepE<'a, R, M> {
-    fn index(&self, idx: Vec3u) -> StepEContext<'a, R, M> {
-        let dim = self.meta.dim;
-        let stim_e_matrix = Array3::new(self.stimulus_e, dim);
-        let mat_matrix = Array3::new(self.material, dim);
-        let e = Array3::new(self.e, dim);
-        let h = Array3::new(self.h, dim);
+        let mat_array = unsafe { SizedArray::new(material, len) };
+        let mat_matrix = DimensionedSlice::new(dim, &mat_array);
 
-        let in_e = e.get(idx).unwrap();
-        let in_h = VolumeSampleNeg::from_indexable(h, idx);
+        let mut e_array = unsafe { SizedArray::new(e, len) };
+        let mut e_matrix = DimensionedSlice::new(dim, &mut e_array);
 
-        let mat = mat_matrix.into_ref(idx);
+        let h_array = unsafe { SizedArray::new(h, len) };
+        let h_matrix = DimensionedSlice::new(dim, &h_array);
 
-        StepEContext {
-            inv_feature_size: self.meta.inv_feature_size,
-            time_step: self.meta.time_step,
-            stim_e: stim_e_matrix.get(idx).unwrap(),
+        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,
             mat,
             in_h,
             in_e,
-        }
-    }
-
-    fn write_output(self, idx: Vec3u, e: Vec3<R>) {
-        let dim = self.meta.dim;
-        let arr_e = Array3Mut::new(self.e, dim);
-        let mut out_e = arr_e.into_mut_handle(idx);
-        out_e.write(e);
+        };
+        let new_e = update_state.step_e();
+        e_matrix[idx] = new_e;
     }
 }
 

crates/spirv_backend/src/lib.rs

@@ -8,14 +8,12 @@
 
 extern crate spirv_std;
 
-pub use spirv_std::glam;
+pub use spirv_std::{glam, RuntimeArray};
 #[cfg(not(target_arch = "spirv"))]
 use spirv_std::macros::spirv;
 
 mod adapt;
-pub mod support;
-
-pub use support::UnsizedArray;
+mod support;
 
 use coremem_types::compound::Optional;
 use coremem_types::mat::{Ferroxcube3R1MH, FullyGenericMaterial, IsoConductorOr};
@@ -56,12 +54,12 @@ macro_rules! steps {
             #[spirv(global_invocation_id)] id: glam::UVec3,
             #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] meta: &SimMeta<$flt>,
             // XXX: delete this input?
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] _unused_stimulus_e: &UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] stimulus_h: &UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &UnsizedArray<$mat>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &mut UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] m: &mut UnsizedArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] _unused_stimulus_e: &RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] stimulus_h: &RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &RuntimeArray<$mat>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &mut RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] m: &mut RuntimeArray<Vec3<$flt>>,
         ) {
             adapt::step_h(glam_vec_to_internal(id), meta, stimulus_h, material, e, h, m)
         }
@@ -70,14 +68,14 @@ macro_rules! steps {
         pub fn $step_e(
             #[spirv(global_invocation_id)] id: glam::UVec3,
             #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] meta: &SimMeta<$flt>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] stimulus_e: &UnsizedArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] stimulus_e: &RuntimeArray<Vec3<$flt>>,
             // XXX: delete this input?
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] _unused_stimulus_h: &UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &UnsizedArray<$mat>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &mut UnsizedArray<Vec3<$flt>>,
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &UnsizedArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] _unused_stimulus_h: &RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] material: &RuntimeArray<$mat>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] e: &mut RuntimeArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] h: &RuntimeArray<Vec3<$flt>>,
             // XXX: can/should this m input be deleted?
-            #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] _unused_m: &UnsizedArray<Vec3<$flt>>,
+            #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] _unused_m: &RuntimeArray<Vec3<$flt>>,
         ) {
             adapt::step_e(glam_vec_to_internal(id), meta, stimulus_e, material, e, h)
         }

crates/spirv_backend/src/support.rs

@@ -1,200 +1,90 @@
-use core::ops::Index;
+use core::ops::{Index, IndexMut};
+
+use spirv_std::RuntimeArray;
 
-use coremem_types::compound::Optional;
 use coremem_types::vec::Vec3u;
 
-/// This struct allows doing things like *(ptr + offset) = value.
-/// Such code wouldn't ordinarily compile with the spirv target because of
-/// unsupported pointer math. RuntimeArray exists to overcome this, however
-/// it emits invalid code for non-primitive types. Hence, this hack.
-// XXX: maximum bytes an array may occupy is 0x7fff_ffff
-// We don't know the element size, so assume it to be <= 44 bytes
-// pub const MAX_UNSIZED_ARRAY: usize = 0x2e7_ffff;
-pub const MAX_UNSIZED_ARRAY: usize = 0x1ff_ffff;
-pub struct UnsizedArray<T>(pub [T; MAX_UNSIZED_ARRAY]);
+pub struct SizedArray<T> {
+    items: T,
+    len: usize,
+}
 
-impl<T: Copy> UnsizedArray<T> {
-    /// user is responsible for bounds-checking the index before calling this.
-    /// unsafe because this can trigger OOB reads if used incorrectly.
-    pub unsafe fn index(&self, index: usize) -> T {
-        // *self.0.index(index)
-        self.0[index]
-        // *self.0.get_unchecked(index)
-        // &self.0
-        //asm! {
-        //    "%result = OpAccessChain _ {arr} {index}",
-        //    "OpReturnValue %result",
-        //    "%unused = OpLabel",
-        //    arr = in(reg) self,
-        //    index = in(reg) index,
-        //}
+impl<T> SizedArray<T> {
+    pub unsafe fn new(items: T, len: usize) -> Self {
+        Self { items, len }
     }
 }
 
-impl<T> UnsizedArray<T> {
-    /// user is responsible for bounds-checking the index before calling this.
-    /// unsafe because this can trigger OOB writes if used incorrectly.
-    pub unsafe fn write(&mut self, index: usize, value: T) {
-        self.0[index] = value;
-    }
-}
+impl<'a, T> Index<usize> for SizedArray<&'a RuntimeArray<T>> {
+    type Output=T;
 
-impl<T> UnsizedArray<T> {
-    /// user is responsible for bounds-checking the index before calling this.
-    /// unsafe because this can trigger OOB access if used incorrectly.
-    pub unsafe fn index_ref(&self, index: usize) -> &T {
-        &self.0[index]
-    }
-
-    pub unsafe fn get_handle<'a>(&'a self, index: usize) -> ArrayHandle<'a, T> {
-        ArrayHandle {
-            base: self,
-            index,
-        }
-    }
-
-    pub unsafe fn get_handle_mut<'a>(&'a mut self, index: usize) -> ArrayHandleMut<'a, T> {
-        ArrayHandleMut {
-            base: self,
-            index,
-        }
-    }
-}
-
-pub struct ArrayHandle<'a, T> {
-    base: &'a UnsizedArray<T>,
-    index: usize,
-}
-
-impl<'a, T: Copy> ArrayHandle<'a, T> {
-    pub fn get(&self) -> T {
+    fn index(&self, idx: usize) -> &Self::Output {
+        assert!(idx < self.len);
         unsafe {
-            self.base.index(self.index)
+            self.items.index(idx)
         }
     }
-    pub fn get_into(&self, out: &mut T) {
-        core::clone::Clone::clone_from(out, self.get_ref());
-    }
 }
-impl<'a, T> ArrayHandle<'a, T> {
-    pub fn get_ref(&self) -> &'a T {
+
+impl<'a, T> Index<usize> for SizedArray<&'a mut RuntimeArray<T>> {
+    type Output=T;
+
+    fn index(&self, idx: usize) -> &Self::Output {
+        assert!(idx < self.len);
         unsafe {
-            self.base.index_ref(self.index)
+            self.items.index(idx)
         }
     }
 }
-
-
-pub struct ArrayHandleMut<'a, T> {
-    base: &'a mut UnsizedArray<T>,
-    index: usize,
-}
-
-impl<'a, T: Copy> ArrayHandleMut<'a, T> {
-    pub fn write(&mut self, value: T) {
+impl<'a, T> IndexMut<usize> for SizedArray<&'a mut RuntimeArray<T>> {
+    fn index_mut(&mut self, idx: usize) -> &mut Self::Output {
+        assert!(idx < self.len);
         unsafe {
-            self.base.write(self.index, value);
-        }
-    }
-    pub fn get(&self) -> T {
-        unsafe {
-            self.base.index(self.index)
+            self.items.index_mut(idx)
         }
     }
 }
 
-/// 3d dynamically-sized array backed by a borrowed buffer.
-#[derive(Clone, Copy)]
-pub struct Array3<'a, T> {
-    data: &'a UnsizedArray<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
 }
 
-fn checked_index(idx: Vec3u, dim: Vec3u) -> Optional<usize> {
-    if idx.x() < dim.x() && idx.y() < dim.y() && idx.z() < dim.z() {
-        let flat_idx = index(idx, dim);
-        Optional::some(flat_idx)
-    } else {
-        Optional::none()
-    }
-}
-
-impl<'a, T> Array3<'a, T> {
-    pub fn new(data: &'a UnsizedArray<T>, dim: Vec3u) -> Self {
-        Self {
-            data,
-            dim,
-        }
-    }
-
-    pub fn into_ref(self, idx: Vec3u) -> &'a T {
-        let idx = checked_index(idx, self.dim).unwrap();
-        unsafe {
-            self.data.index_ref(idx)
-        }
-    }
-}
-
-impl<'a, T: Copy + Default> Array3<'a, T> {
-    pub fn get(&self, idx: Vec3u) -> Optional<T> {
-        let idx = checked_index(idx, self.dim);
-        if idx.is_some() {
-            Optional::some(unsafe {
-                self.data.index(idx.unwrap())
-            })
-        } else {
-            Optional::none()
-        }
-    }
-}
-
-impl<'a, T> Index<Vec3u> for Array3<'a, T> {
-    type Output = T;
-    fn index(&self, idx: Vec3u) -> &Self::Output {
-        let idx = checked_index(idx, self.dim).unwrap();
-        unsafe {
-            self.data.index_ref(idx)
-        }
-    }
-}
-
-/// 3d dynamically-sized array backed by a mutably borrowed buffer.
-pub struct Array3Mut<'a, T> {
-    data: &'a mut UnsizedArray<T>,
-    dim: Vec3u,
-}
-
-impl<'a, T> Array3Mut<'a, T> {
-    pub fn new(data: &'a mut UnsizedArray<T>, dim: Vec3u) -> Self {
-        Self {
-            data,
-            dim,
-        }
-    }
-
-    pub fn index(&self, idx: Vec3u) -> Optional<usize> {
-        if idx.x() < self.dim.x() && idx.y() < self.dim.y() && idx.z() < self.dim.z() {
-            let flat_idx = index(idx, self.dim);
-            Optional::some(flat_idx)
-        } else {
-            Optional::none()
-        }
-    }
-}
-
-impl<'a, T: Copy> Array3Mut<'a, T> {
-    pub fn into_mut_handle(self, idx: Vec3u) -> ArrayHandleMut<'a, T> {
-        let idx = self.index(idx).unwrap();
-        unsafe {
-            self.data.get_handle_mut(idx)
-        }
-    }
-}
-
 
 #[cfg(test)]
 mod test {