Skip to main content

sigmoid_generic/
sigmoid_generic.rs

1// Copyright 2026 the Fearless_SIMD Authors
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! Apply the sigmoid function `x / sqrt(x * x + 1)` using either `f32` or `f64`.
5//!
6//! `SimdFloatElement::Native<S>` selects the native-width vector for the given
7//! scalar type and the current SIMD backend.
8
9use fearless_simd::{Level, dispatch, prelude::*};
10
11fn sigmoid<T: SimdFloatElement>(level: Level, input: &[T], output: &mut [T]) {
12    assert_eq!(
13        input.len(),
14        output.len(),
15        "input and output lengths must match"
16    );
17    dispatch!(level, simd => sigmoid_simd(simd, input, output));
18}
19
20#[inline(always)]
21fn sigmoid_simd<S: Simd, T: SimdFloatElement>(simd: S, input: &[T], output: &mut [T]) {
22    let n = T::Native::<S>::LEN;
23    let one = T::Native::<S>::splat(simd, T::from(1_u8));
24    let mut inputs = input.chunks_exact(n);
25    let mut outputs = output.chunks_exact_mut(n);
26
27    for (input, output) in inputs.by_ref().zip(outputs.by_ref()) {
28        let a = T::Native::<S>::from_slice(simd, input);
29        let b = a / (a * a + one).sqrt();
30        b.store_slice(output);
31    }
32
33    let input = inputs.remainder();
34    let output = outputs.into_remainder();
35    if !input.is_empty() {
36        // Padding keeps sqrt in SIMD, without requiring a scalar sqrt trait.
37        let a = T::Native::<S>::from_fn(simd, |i| input.get(i).copied().unwrap_or_default());
38        let b = a / (a * a + one).sqrt();
39        output.copy_from_slice(&b.as_slice()[..input.len()]);
40    }
41}
42
43fn main() {
44    let level = Level::new();
45    let input_f32 = [
46        0.1_f32, -0.2, 0.001, 0.4, 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11., 12., 13.,
47    ];
48    let input_f64 = input_f32.map(f64::from);
49    let mut output_f32 = [0.0; 17];
50    let mut output_f64 = [0.0; 17];
51
52    sigmoid(level, &input_f32, &mut output_f32);
53    sigmoid(level, &input_f64, &mut output_f64);
54
55    println!("f32: {output_f32:?}");
56    println!("f64: {output_f64:?}");
57}