Reshape the matrices so more operations can be vectorized

experimental/segmenter/src/lstm.rs

@@ -3,7 +3,7 @@
 // (online at: https://github.com/unicode-org/icu4x/blob/main/LICENSE ).
 
 use crate::grapheme::GraphemeClusterSegmenter;
-use crate::math_helper::{self, MatrixBorrowedMut, MatrixOwned, MatrixZero};
+use crate::math_helper::{MatrixBorrowedMut, MatrixOwned, MatrixZero};
 use crate::provider::*;
 use alloc::boxed::Box;
 use alloc::string::String;
@@ -161,37 +161,37 @@ impl<'l> LstmSegmenter<'l> {
                 let hunits = h_tm1.dim();
                 let embedd_dim = x_t.dim();
                 c_tm1.as_borrowed().debug_assert_dims([hunits]);
-                w.debug_assert_dims([hunits, 4, embedd_dim]);
-                u.debug_assert_dims([hunits, 4, hunits]);
-                b.debug_assert_dims([hunits, 4]);
+                w.debug_assert_dims([4, hunits, embedd_dim]);
+                u.debug_assert_dims([4, hunits, hunits]);
+                b.debug_assert_dims([4, hunits]);
             }
 
             let mut s_t = b.to_owned();
 
             s_t.as_mut().add_dot_3d_2(x_t, w);
             s_t.as_mut().add_dot_3d_1(h_tm1.as_borrowed(), u);
 
-            #[allow(clippy::unwrap_used)]
-            for i in 0..s_t.dim().0 {
-                let [s0, s1, s2, s3] = s_t
-                    .as_borrowed()
-                    .submatrix::<1>(i)
-                    .unwrap()
-                    .read_4()
-                    .unwrap(); // shape (hunits, 4)
-                let p = math_helper::sigmoid(s0);
-                let f = math_helper::sigmoid(s1);
-                let c = math_helper::tanh(s2);
-                let o = math_helper::sigmoid(s3);
-                let c_old = c_tm1.as_borrowed().as_slice().get(i).unwrap(); // shape (h_units)
-                let c_new = math_helper::fma(p, c, f * c_old);
-                *c_tm1.as_mut_slice().get_mut(i).unwrap() = c_new; // shape (h_units)
-                *h_tm1.as_mut_slice().get_mut(i).unwrap() = o * math_helper::tanh(c_new);
-                // shape (hunits)
-            }
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            s_t.submatrix_mut::<1>(0).unwrap().sigmoid_transform();
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            s_t.submatrix_mut::<1>(1).unwrap().sigmoid_transform();
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            s_t.submatrix_mut::<1>(2).unwrap().tanh_transform();
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            s_t.submatrix_mut::<1>(3).unwrap().sigmoid_transform();
+
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            c_tm1.convolve(
+                s_t.as_borrowed().submatrix(0).unwrap(),
+                s_t.as_borrowed().submatrix(2).unwrap(),
+                s_t.as_borrowed().submatrix(1).unwrap(),
+            );
+
+            #[allow(clippy::unwrap_used)] // first dimension is 4
+            h_tm1.mul_tanh(s_t.as_borrowed().submatrix(3).unwrap(), c_tm1.as_borrowed());
         }
 
-        let hunits = self.fw_u.dim().0;
+        let hunits = self.fw_u.dim().1;
 
         // Forward LSTM
         let mut c_fw = MatrixOwned::<1>::new_zero([hunits]);

experimental/segmenter/src/math_helper.rs

@@ -24,12 +24,6 @@ pub fn sigmoid(x: f32) -> f32 {
     1.0 / (1.0 + (-x).exp())
 }
 
-/// computes x * y + z in one instruction
-#[inline]
-pub fn fma(x: f32, y: f32, z: f32) -> f32 {
-    x.mul_add(y, z)
-}
-
 /// A `D`-dimensional, heap-allocated matrix.
 ///
 /// This matrix implementation supports slicing matrices into tightly-packed
@@ -170,12 +164,6 @@ impl_basic_dim!(
 );
 impl_basic_dim!(MatrixZero<'a, 1>, MatrixZero<'a, 2>, MatrixZero<'a, 3>);
 
-impl<'a> MatrixBorrowed<'a, 1> {
-    pub fn read_4(&self) -> Option<[f32; 4]> {
-        <&[f32; 4]>::try_from(self.data).ok().copied()
-    }
-}
-
 /// A `D`-dimensional, mutably borrowed matrix.
 pub struct MatrixBorrowedMut<'a, const D: usize> {
     pub(crate) data: &'a mut [f32],
@@ -223,6 +211,58 @@ impl<'a, const D: usize> MatrixBorrowedMut<'a, D> {
             *v = v.exp() / sm;
         });
     }
+
+    pub fn sigmoid_transform(&mut self) {
+        for x in &mut self.data.iter_mut() {
+            *x = sigmoid(*x);
+        }
+    }
+
+    pub fn tanh_transform(&mut self) {
+        for x in &mut self.data.iter_mut() {
+            *x = tanh(*x);
+        }
+    }
+
+    pub fn convolve(
+        &mut self,
+        i: MatrixBorrowed<'_, D>,
+        c: MatrixBorrowed<'_, D>,
+        f: MatrixBorrowed<'_, D>,
+    ) {
+        let i = i.as_slice();
+        let c = c.as_slice();
+        let f = f.as_slice();
+        let len = self.data.len();
+        if len != i.len() || len != c.len() || len != f.len() {
+            debug_assert!(false, "LSTM matrices not the correct dimensions");
+            return;
+        }
+        for idx in 0..len {
+            // Safety: The lengths are all the same (checked above)
+            unsafe {
+                *self.data.get_unchecked_mut(idx) = i.get_unchecked(idx) * c.get_unchecked(idx)
+                    + self.data.get_unchecked(idx) * f.get_unchecked(idx)
+            }
+        }
+    }
+
+    pub fn mul_tanh(&mut self, o: MatrixBorrowed<'_, D>, c: MatrixBorrowed<'_, D>) {
+        let o = o.as_slice();
+        let c = c.as_slice();
+        let len = self.data.len();
+        if len != o.len() || len != c.len() {
+            debug_assert!(false, "LSTM matrices not the correct dimensions");
+            return;
+        }
+        for idx in 0..len {
+            // Safety: The lengths are all the same (checked above)
+            unsafe {
+                *self.data.get_unchecked_mut(idx) =
+                    o.get_unchecked(idx) * tanh(*c.get_unchecked(idx));
+            }
+        }
+    }
 }
 
 impl<'a> MatrixBorrowed<'a, 1> {

experimental/segmenter/src/provider/lstm.rs

@@ -188,14 +188,17 @@ impl<'data> LstmDataFloat32<'data> {
         let dic_len = u16::try_from(dic.len())
             .map_err(|_| DataError::custom("Dictionary does not fit in u16"))?;
 
-        if embedding.dims[0] - 1 != dic_len
-            || fw_w.dims != [fw_u.dims[0], 4, embedding.dims[1]]
-            || fw_u.dims != [fw_u.dims[0], 4, fw_u.dims[0]]
-            || fw_b.dims != [fw_u.dims[0], 4]
-            || bw_w.dims != [fw_u.dims[0], 4, embedding.dims[1]]
-            || bw_u.dims != [fw_u.dims[0], 4, fw_u.dims[0]]
-            || bw_b.dims != [fw_u.dims[0], 4]
-            || time_w.dims != [2, 4, fw_u.dims[0]]
+        let num_classes = embedding.dims[0];
+        let embedd_dim = embedding.dims[1];
+        let hunits = fw_u.dims[2];
+        if num_classes - 1 != dic_len
+            || fw_w.dims != [4, hunits, embedd_dim]
+            || fw_u.dims != [4, hunits, hunits]
+            || fw_b.dims != [4, hunits]
+            || bw_w.dims != [4, hunits, embedd_dim]
+            || bw_u.dims != [4, hunits, hunits]
+            || bw_b.dims != [4, hunits]
+            || time_w.dims != [2, 4, hunits]
             || time_b.dims != [4]
         {
             return Err(DataError::custom("LSTM dimension mismatch"));

provider/datagen/src/transform/segmenter/lstm.rs

@@ -90,17 +90,19 @@ impl RawLstmData {
         // Unwraps okay: dimensions checked above
         let mut fw_w = fw_w.into_shape((embedd_dim, 4, hunits)).unwrap();
         let mut fw_u = fw_u.into_shape((hunits, 4, hunits)).unwrap();
-        let mut fw_b = fw_b.into_shape((4, hunits)).unwrap();
+        let fw_b = fw_b.into_shape((4, hunits)).unwrap();
         let mut bw_w = bw_w.into_shape((embedd_dim, 4, hunits)).unwrap();
         let mut bw_u = bw_u.into_shape((hunits, 4, hunits)).unwrap();
-        let mut bw_b = bw_b.into_shape((4, hunits)).unwrap();
+        let bw_b = bw_b.into_shape((4, hunits)).unwrap();
         let mut time_w = time_w.into_shape((2, hunits, 4)).unwrap();
         fw_w.swap_axes(0, 2);
+        fw_w.swap_axes(0, 1);
         fw_u.swap_axes(0, 2);
-        fw_b.swap_axes(0, 1);
+        fw_u.swap_axes(0, 1);
         bw_w.swap_axes(0, 2);
+        bw_w.swap_axes(0, 1);
         bw_u.swap_axes(0, 2);
-        bw_b.swap_axes(0, 1);
+        bw_u.swap_axes(0, 1);
         time_w.swap_axes(1, 2);
         let fw_w = fw_w.as_standard_layout().into_owned();
         let fw_u = fw_u.as_standard_layout().into_owned();
@@ -110,6 +112,15 @@ impl RawLstmData {
         let bw_b = bw_b.as_standard_layout().into_owned();
         let time_w = time_w.as_standard_layout().into_owned();
 
+        assert_eq!(fw_w.shape(), [4, hunits, embedd_dim]);
+        assert_eq!(fw_u.shape(), [4, hunits, hunits]);
+        assert_eq!(fw_b.shape(), [4, hunits]);
+        assert_eq!(bw_w.shape(), [4, hunits, embedd_dim]);
+        assert_eq!(bw_u.shape(), [4, hunits, hunits]);
+        assert_eq!(bw_b.shape(), [4, hunits]);
+        assert_eq!(time_w.shape(), [2, 4, hunits]);
+        assert_eq!(time_b.shape(), [4]);
+
         let model = if self.model.contains("_codepoints") {
             ModelType::Codepoints
         } else if self.model.contains("_graphclust_") {