sim: add AbstractStimulus::eval_into for bulk evaluation

crates/coremem/src/sim/spirv/mod.rs

@@ -1,4 +1,3 @@
-use ndarray::{self, Array3};
 use serde::{Deserialize, Serialize};
 use log::{info, trace, warn};
 
@@ -7,9 +6,10 @@ use crate::geom::{Coord, Index};
 use crate::real::Real;
 use crate::sim::{AbstractSim, Fields, GenericSim, StaticSim};
 use crate::stim::AbstractStimulus;
-use crate::cross::vec::Vec3;
+use coremem_cross::dim::OffsetDimSlice;
 use coremem_cross::mat::{FullyGenericMaterial, Material};
 use coremem_cross::step::SimMeta;
+use coremem_cross::vec::{Vec3, Vec3u};
 
 mod cpu;
 mod gpu;
@@ -104,8 +104,8 @@ where
         self.backend.step_n(
             self.meta,
             self.mat.as_slice(),
-            stim_e.as_slice().unwrap(),
+            &*stim_e,
-            stim_h.as_slice().unwrap(),
+            &*stim_h,
             self.e.as_mut_slice(),
             self.h.as_mut_slice(),
             self.m.as_mut_slice(),
@@ -225,32 +225,43 @@ where
     }
 
     fn eval_stimulus<S: AbstractStimulus>(&self, stim: &S)
-        -> (Array3<Vec3<R>>, Array3<Vec3<R>>)
+        -> (Vec<Vec3<R>>, Vec<Vec3<R>>)
     {
         trace!("eval_stimulus begin");
         let (e, h) = self.diag.instrument_stimuli(|| {
             let dim = self.size();
+            let dim_len = dim.product_sum_usize();
             let feature_size = self.feature_size();
             let t_sec = self.time();
             let timestep = self.meta.time_step;
 
-            // TODO(perf): do this in one loop!
+            // we'll evaluate in parallel each row (const y/z) of the stimulus.
-            let e = ndarray::Zip::from(ndarray::indices(
+            let mut backing = Vec::new();
-                [dim.z() as usize, dim.y() as usize, dim.x() as usize]
+            backing.resize_with(dim_len, Default::default);
-            )).par_map_collect(|(z, y, x)| {
+
-                let pos_idx = Index::new(x as _, y as _, z as _);
+            rayon::scope(|s| {
-                let pos_meters = pos_idx.to_meters(feature_size);
+                let mut undispatched_backing = &mut backing[..];
-                let densities = stim.at(t_sec, pos_meters);
+                for z in 0..dim.z() {
-                densities.e.cast::<R>() * timestep
+                    for y in 0..dim.y() {
-            });
+                        let this_slice;
-            let h = ndarray::Zip::from(ndarray::indices(
+                        (this_slice, undispatched_backing) = undispatched_backing.split_at_mut(dim.x() as usize);
-                [dim.z() as usize, dim.y() as usize, dim.x() as usize]
+                        let view = OffsetDimSlice::new(Vec3u::new(0, y, z), Vec3u::new(this_slice.len() as u32, 1, 1), this_slice);
-            )).par_map_collect(|(z, y, x)| {
+                        s.spawn(move |_| stim.eval_into(t_sec, feature_size, view));
-                let pos_idx = Index::new(x as _, y as _, z as _);
+                    }
-                let pos_meters = pos_idx.to_meters(feature_size);
+                }
-                let densities = stim.at(t_sec, pos_meters);
-                densities.h.cast::<R>() * timestep
             });
+
+            // unpack the E and H portions of the stimulus
+            let mut e = Vec::new();
+            e.reserve(dim_len);
+            let mut h = Vec::new();
+            h.reserve(dim_len);
+
+            for field in backing {
+                e.push(field.e.cast::<R>() * timestep);
+                h.push(field.h.cast::<R>() * timestep);
+            }
+
             (e, h)
         });
         trace!("eval_stimulus end");

crates/coremem/src/stim.rs

@@ -59,14 +59,41 @@ impl<S: AbstractStimulus> Visitor<&S> for &mut StimulusEvaluator {
     }
 }
 
+struct StimEvalInto<'a> {
+    into: OffsetDimSlice<&'a mut [Fields]>,
+    t_sec: f32,
+    feat_size: f32,
+}
+
+impl<'a, S: AbstractStimulus> Visitor<&S> for StimEvalInto<'a> {
+    fn visit(&mut self, next: &S) {
+        next.eval_into(self.t_sec, self.feat_size, self.into.as_mut());
+    }
+}
+
+// impl<L: Sync> AbstractStimulus for L
+// where
+//     for<'a, 'b> &'a L: Visit<&'b mut StimulusEvaluator>,
+// {
+//     fn at(&self, t_sec: f32, pos: Meters) -> Fields {
+//         let mut ev = StimulusEvaluator { t_sec, pos, fields: Fields::default()};
+//         self.visit(&mut ev);
+//         ev.fields
+//     }
+// }
+
 impl<L: Sync> AbstractStimulus for L
 where
-    for<'a, 'b> &'a L: Visit<&'b mut StimulusEvaluator>,
+    for<'a, 'b> &'a L: Visit<StimEvalInto<'b>>,
 {
-    fn at(&self, t_sec: f32, pos: Meters) -> Fields {
+    fn at(&self, _t_sec: f32, _pos: Meters) -> Fields {
-        let mut ev = StimulusEvaluator { t_sec, pos, fields: Fields::default()};
+        // TODO: if we replace `_pos` with `feat_size`, `Index`, we can replace this with a generic
-        self.visit(&mut ev);
+        // impl that calls eval_into one a one-len slice.
-        ev.fields
+        unimplemented!();
+    }
+    fn eval_into(&self, t_sec: f32, feat_size: f32, into: OffsetDimSlice<&mut [Fields]>) {
+        let ev = StimEvalInto { t_sec, feat_size, into };
+        self.visit(ev);
     }
 }
 
@@ -82,6 +109,12 @@ where
 pub struct DynStimuli(Vec<Box<dyn AbstractStimulus>>);
 
 impl DynStimuli {
+    pub fn new() -> Self {
+        Self::default()
+    }
+    pub fn from_vec(stim: Vec<Box<dyn AbstractStimulus>>) -> Self {
+        Self(stim)
+    }
     pub fn push(&mut self, a: Box<dyn AbstractStimulus>) {
         self.0.push(a)
     }
@@ -91,6 +124,11 @@ impl AbstractStimulus for DynStimuli {
         self.0.iter().map(|i| i.at(t_sec, pos))
             .fold(Fields::default(), core::ops::Add::add)
     }
+    fn eval_into(&self, t_sec: f32, feat_size: f32, mut into: OffsetDimSlice<&mut [Fields]>) {
+        for i in &self.0 {
+            i.eval_into(t_sec, feat_size, into.as_mut());
+        }
+    }
 }
 // 
 // impl AbstractStimulus for Box<dyn AbstractStimulus> {
@@ -405,6 +443,11 @@ impl<T: AbstractStimulus> AbstractStimulus for Gated<T> {
             Default::default()
         }
     }
+    fn eval_into(&self, t_sec: f32, feat_size: f32, into: OffsetDimSlice<&mut [Fields]>) {
+        if (self.start..self.end).contains(&t_sec) {
+            self.inner.eval_into(t_sec, feat_size, into);
+        }
+    }
 }
 
 #[derive(Clone)]
@@ -436,6 +479,9 @@ impl<T: AbstractStimulus> AbstractStimulus for Shifted<T> {
     fn at(&self, t_sec: f32, pos: Meters) -> Fields {
         self.inner.at(t_sec - self.start_at, pos)
     }
+    fn eval_into(&self, t_sec: f32, feat_size: f32, into: OffsetDimSlice<&mut [Fields]>) {
+        self.inner.eval_into(t_sec - self.start_at, feat_size, into)
+    }
 }
 
 #[cfg(test)]