Introduce safe SIMD API, port some simple vectorized operations to use it

rten-simd/src/lib.rs

@@ -35,6 +35,8 @@ pub mod isa_detection;
 pub mod span;
 mod vec;
 
+pub mod safe;
+
 pub use vec::{vec_count, Simd, SimdFloat, SimdInt, SimdMask};
 
 #[cfg(feature = "avx512")]

rten-simd/src/safe.rs

@@ -0,0 +1,97 @@
+//! Portable SIMD library.
+//!
+//! _This module contains a new portable SIMD API which is in development. The
+//! focus is to revise the API to reduce the amount of unsafe code needed when
+//! using it._
+//!
+//! ## Usage
+//!
+//! The steps to define a vectorized operation using this module are:
+//!
+//! 1. Create a struct containing the operation's parameters.
+//! 2. Implement the [`SimdOp`] trait for the struct to define how to evaluate
+//!    the operation.
+//! 3. Call [`SimdOp::dispatch`] to select the preferred SIMD instruction set and
+//!    evaluate the operation using it.
+//!
+//! ## Functional utilities
+//!
+//! The [`functional`] module provides utilities for defining vectorized
+//! transforms ([`simd_map`](functional::simd_map)) and reductions
+//! ([`simd_fold`](functional::simd_fold)).
+mod arch;
+mod dispatch;
+pub mod functional;
+mod iter;
+mod vec;
+
+pub use dispatch::{SimdOp, SimdUnaryOp};
+pub use iter::{IntoSimdIterator, Iter};
+pub use vec::{Elem, Isa, Mask, Simd, SimdFloatOps, SimdIntOps, SimdOps};
+
+#[cfg(test)]
+pub(crate) use dispatch::test_simd_op;
+
+/// Test that two [`Simd`] vectors are equal according to a [`PartialEq`]
+/// comparison of their array representations.
+#[cfg(test)]
+macro_rules! assert_simd_eq {
+    ($x:expr, $y:expr) => {
+        assert_eq!($x.to_array(), $y.to_array());
+    };
+}
+
+#[cfg(test)]
+pub(crate) use assert_simd_eq;
+
+#[cfg(test)]
+mod tests {
+    use super::functional::simd_map;
+    use super::{Isa, SimdOp, SimdOps};
+
+    #[test]
+    fn test_simd_f32_op() {
+        struct Square<'a> {
+            xs: &'a mut [f32],
+        }
+
+        impl<'a> SimdOp for Square<'a> {
+            type Output = &'a mut [f32];
+
+            fn eval<I: Isa>(self, isa: I) -> Self::Output {
+                let ops = isa.f32();
+                simd_map(ops, self.xs, |x| ops.mul(x, x))
+            }
+        }
+
+        let mut buf: Vec<_> = (0..32).map(|x| x as f32).collect();
+        let expected: Vec<_> = buf.iter().map(|x| *x * *x).collect();
+
+        let squared = Square { xs: &mut buf }.dispatch();
+
+        assert_eq!(squared, &expected);
+    }
+
+    #[test]
+    fn test_simd_i32_op() {
+        struct Square<'a> {
+            xs: &'a mut [i32],
+        }
+
+        impl<'a> SimdOp for Square<'a> {
+            type Output = &'a mut [i32];
+
+            fn eval<I: Isa>(self, isa: I) -> Self::Output {
+                let ops = isa.i32();
+                simd_map(ops, self.xs, |x| ops.mul(x, x))
+            }
+        }
+
+        let mut buf: Vec<_> = (0..32).collect();
+        let expected: Vec<_> = buf.iter().map(|x| *x * *x).collect();
+
+        let squared = Square { xs: &mut buf }.dispatch();
+
+        assert_eq!(squared, &expected);
+    }
+}

rten-simd/src/safe/arch.rs

@@ -0,0 +1,10 @@
+#[cfg(target_arch = "aarch64")]
+pub mod aarch64;
+
+#[cfg(target_arch = "x86_64")]
+pub mod x86_64;
+
+#[cfg(target_arch = "wasm32")]
+pub mod wasm32;
+
+pub mod generic;