Skip to main content

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}