sigmoid/sigmoid.rs
1// Copyright 2024 the Fearless_SIMD Authors
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! This example demonstrates the typical usage Fearless SIMD.
5//!
6//! The vector size matches the native vector size of the hardware:
7//!
8//! - SSE and NEON get 128 bit chunks
9//! - AVX2 gets 256 bit ones
10//! - AVX-512 gets 512-bit ones
11//!
12//! All from a single function.
13
14use fearless_simd::{Level, dispatch, prelude::*};
15
16/// Applies the sigmoid function to the input and writes to the output
17#[inline(always)] // or #[simd], either works
18fn sigmoid<S: Simd>(simd: S, x: &[f32], out: &mut [f32]) {
19 let n = S::f32s::LEN; // CPU's native vector size
20
21 // fast vectorized loop
22 for (x, y) in x.chunks_exact(n).zip(out.chunks_exact_mut(n)) {
23 let a = S::f32s::from_slice(simd, x);
24 let b = a / (a * a + 1.0).sqrt();
25 b.store_slice(y);
26 }
27
28 // scalar processing of the remainder smaller than a single vector
29 let x_remainder = x.chunks_exact(n).remainder();
30 let y_remainder = out.chunks_exact_mut(n).into_remainder();
31 for (a, b) in x_remainder.iter().zip(y_remainder) {
32 *b = a / (a * a + 1.0).sqrt();
33 }
34}
35
36fn main() {
37 let level = Level::new();
38 let inp = [
39 0.1, -0.2, 0.001, 0.4, 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14.,
40 ];
41 let mut out = [0.; 18];
42 // dispatch! selects the best implementation for the CPU we're running on
43 dispatch!(level, simd => sigmoid(simd, &inp, &mut out));
44
45 println!("{out:?}");
46}