#[cfg(target_arch = "x86_64")]
use super::horizontal::horizontal_sum_256;
#[cfg(target_arch = "x86_64")]
use super::is_avx2_fma_available;
const CEPHES_LN2_HI: f32 = 6.931_457_5e-1;
const CEPHES_LN2_LO: f32 = 1.428_606_8e-6;
const CEPHES_INV_LN2: f32 = std::f32::consts::LOG2_E;
#[inline(always)]
pub fn simd_exp_inplace(x: &mut [f32]) {
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_exp_inplace(x) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_exp_inplace(x) }
} else {
scalar_exp_inplace(x)
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_exp_inplace(x) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_exp_inplace(x)
}
}
#[inline(always)]
pub fn simd_exp_sum_inplace(x: &mut [f32]) -> f32 {
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_exp_sum_inplace(x) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_exp_sum_inplace(x) }
} else {
scalar_exp_sum_inplace(x)
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_exp_sum_inplace(x) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_exp_sum_inplace(x)
}
}
#[inline(always)]
pub fn simd_reciprocal_inplace(x: &mut [f32]) {
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_reciprocal_inplace(x) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_reciprocal_inplace(x) }
} else {
scalar_reciprocal_inplace(x)
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_reciprocal_inplace(x) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_reciprocal_inplace(x)
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_reciprocal_inplace(x: &mut [f32]) {
for val in x.iter_mut() {
*val = 1.0 / *val;
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_exp_sum_inplace(x: &mut [f32]) -> f32 {
let mut sum = 0.0f32;
for val in x.iter_mut() {
let e = cephes_exp_scalar(*val);
*val = e;
sum += e;
}
sum
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_reciprocal_inplace(x: &mut [f32]) {
use std::arch::aarch64::*;
unsafe {
let len = x.len();
let chunks = len / 4;
let ones = vdupq_n_f32(1.0);
for i in 0..chunks {
let v = vld1q_f32(x.as_ptr().add(i * 4));
let r = vdivq_f32(ones, v);
vst1q_f32(x.as_mut_ptr().add(i * 4), r);
}
for i in (chunks * 4)..len {
*x.get_unchecked_mut(i) = 1.0 / *x.get_unchecked(i);
}
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_reciprocal_inplace(x: &mut [f32]) {
use std::arch::x86_64::*;
unsafe {
let len = x.len();
let chunks = len / 8;
let ones = _mm256_set1_ps(1.0);
for i in 0..chunks {
let v = _mm256_loadu_ps(x.as_ptr().add(i * 8));
let r = _mm256_div_ps(ones, v);
_mm256_storeu_ps(x.as_mut_ptr().add(i * 8), r);
}
for i in (chunks * 8)..len {
*x.get_unchecked_mut(i) = 1.0 / *x.get_unchecked(i);
}
}
}
#[inline(always)]
pub fn cephes_exp_scalar(x: f32) -> f32 {
let n = (x * CEPHES_INV_LN2).round() as i32;
if n < -126 {
return 0.0;
}
if n > 127 {
return f32::INFINITY;
}
let g = x - n as f32 * CEPHES_LN2_HI - n as f32 * CEPHES_LN2_LO;
let q = 1.0
+ g * (1.0
+ g * 0.5
* (1.0
+ g * (1.0 / 3.0)
* (1.0 + g * 0.25 * (1.0 + g * 0.2 * (1.0 + g * (1.0 / 6.0))))));
let bits = ((n + 127) as u32) << 23;
let scale = f32::from_bits(bits);
scale * q
}
#[inline(always)]
pub fn fast_exp(x: f32) -> f32 {
cephes_exp_scalar(x)
}
#[inline(always)]
pub fn fast_sigmoid(x: f32) -> f32 {
if x > 40.0 {
return 1.0;
}
if x < -40.0 {
return 0.0;
}
1.0 / (1.0 + cephes_exp_scalar(-x))
}
#[inline(always)]
pub fn fast_tanh(x: f32) -> f32 {
let ax = x.abs();
if ax > 3.0 {
return x.signum();
}
let x2 = x * x;
x * (27.0 + x2) / (27.0 + 9.0 * x2)
}
#[inline(always)]
pub fn simd_tanh_inplace(x: &mut [f32]) {
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_tanh_inplace(x) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_tanh_inplace(x) }
} else {
scalar_tanh_inplace(x)
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_tanh_inplace(x) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_tanh_inplace(x)
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_tanh_inplace(x: &mut [f32]) {
for val in x.iter_mut() {
*val = fast_tanh(*val);
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_tanh_inplace(x: &mut [f32]) {
use std::arch::aarch64::*;
unsafe {
let len = x.len();
let chunks = len / 4;
let c27 = vdupq_n_f32(27.0);
let c9 = vdupq_n_f32(9.0);
let threshold = vdupq_n_f32(3.0);
let ones = vdupq_n_f32(1.0);
let sign_mask = vdupq_n_f32(-0.0);
for i in 0..chunks {
let v = vld1q_f32(x.as_ptr().add(i * 4));
let x2 = vmulq_f32(v, v);
let num = vfmaq_f32(vmulq_f32(v, c27), v, x2); let den = vfmaq_f32(c27, c9, x2); let pade = vdivq_f32(num, den);
let mask = vcagtq_f32(v, threshold);
let sign = vreinterpretq_f32_u32(vorrq_u32(
vandq_u32(vreinterpretq_u32_f32(v), vreinterpretq_u32_f32(sign_mask)),
vreinterpretq_u32_f32(ones),
));
let result = vbslq_f32(mask, sign, pade);
vst1q_f32(x.as_mut_ptr().add(i * 4), result);
}
for i in (chunks * 4)..len {
*x.get_unchecked_mut(i) = fast_tanh(*x.get_unchecked(i));
}
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_tanh_inplace(x: &mut [f32]) {
use std::arch::x86_64::*;
unsafe {
let len = x.len();
let chunks = len / 8;
let c27 = _mm256_set1_ps(27.0);
let c9 = _mm256_set1_ps(9.0);
let threshold = _mm256_set1_ps(3.0);
let ones = _mm256_set1_ps(1.0);
let sign_mask = _mm256_set1_ps(-0.0f32);
for i in 0..chunks {
let v = _mm256_loadu_ps(x.as_ptr().add(i * 8));
let x2 = _mm256_mul_ps(v, v);
let num = _mm256_fmadd_ps(v, x2, _mm256_mul_ps(v, c27)); let den = _mm256_fmadd_ps(c9, x2, c27); let pade = _mm256_div_ps(num, den);
let ax = _mm256_andnot_ps(sign_mask, v);
let mask = _mm256_cmp_ps(ax, threshold, _CMP_GT_OQ);
let sign = _mm256_or_ps(ones, _mm256_and_ps(v, sign_mask));
let result = _mm256_blendv_ps(pade, sign, mask);
_mm256_storeu_ps(x.as_mut_ptr().add(i * 8), result);
}
for i in (chunks * 8)..len {
*x.get_unchecked_mut(i) = fast_tanh(*x.get_unchecked(i));
}
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_tanh_inplace(x: &mut [f32]) {
use std::arch::wasm32::*;
unsafe {
let len = x.len();
let chunks = len / 4;
let c27 = f32x4_splat(27.0);
let c9 = f32x4_splat(9.0);
let threshold = f32x4_splat(3.0);
let ones = f32x4_splat(1.0);
let neg_zero = f32x4_splat(-0.0);
for i in 0..chunks {
let v = v128_load(x.as_ptr().add(i * 4) as *const v128);
let x2 = f32x4_mul(v, v);
let num = f32x4_add(f32x4_mul(v, c27), f32x4_mul(v, x2));
let den = f32x4_add(c27, f32x4_mul(c9, x2));
let pade = f32x4_div(num, den);
let ax = v128_andnot(neg_zero, v);
let mask = f32x4_gt(ax, threshold);
let sign_v = v128_or(ones, v128_and(v, neg_zero));
let result = v128_bitselect(sign_v, pade, mask);
v128_store(x.as_mut_ptr().add(i * 4) as *mut v128, result);
}
for i in (chunks * 4)..len {
*x.get_unchecked_mut(i) = fast_tanh(*x.get_unchecked(i));
}
}
}
#[inline(always)]
pub fn simd_sigmoid_tanh_clamp_inplace(out: &mut [f32], a: &[f32], q: &[f32], clamp: f32) {
debug_assert!(clamp > 0.0, "clamp must be positive: got {clamp}");
debug_assert_eq!(a.len(), q.len(), "a/q length mismatch");
debug_assert!(out.len() >= a.len(), "out too short");
let len = a.len().min(out.len()).min(q.len());
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_sigmoid_tanh_clamp(&mut out[..len], &a[..len], &q[..len], clamp) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_sigmoid_tanh_clamp(&mut out[..len], &a[..len], &q[..len], clamp) }
} else {
scalar_sigmoid_tanh_clamp(&mut out[..len], &a[..len], &q[..len], clamp)
}
return;
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_sigmoid_tanh_clamp(&mut out[..len], &a[..len], &q[..len], clamp) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_sigmoid_tanh_clamp(&mut out[..len], &a[..len], &q[..len], clamp)
}
}
#[inline]
pub fn simd_sigmoid_inplace(x: &mut [f32]) {
#[cfg(target_arch = "aarch64")]
{
unsafe { neon_sigmoid_inplace(x) }
}
#[cfg(target_arch = "x86_64")]
{
if is_avx2_fma_available() {
unsafe { avx2_sigmoid_inplace(x) }
} else {
scalar_sigmoid_inplace(x)
}
return;
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
{
unsafe { wasm32_sigmoid_inplace(x) }
}
#[cfg(not(any(
target_arch = "aarch64",
target_arch = "x86_64",
all(target_arch = "wasm32", target_feature = "simd128")
)))]
{
scalar_sigmoid_inplace(x)
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_sigmoid_inplace(x: &mut [f32]) {
for v in x.iter_mut() {
*v = fast_sigmoid(*v);
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_sigmoid_tanh_clamp(out: &mut [f32], a: &[f32], q: &[f32], clamp: f32) {
for i in 0..out.len() {
let s = fast_sigmoid(a[i] + q[i]);
let v = 2.0 * s - 1.0;
out[i] = v.clamp(-clamp, clamp);
}
}
#[inline(always)]
#[allow(dead_code)]
pub(super) fn scalar_exp_inplace(x: &mut [f32]) {
for val in x.iter_mut() {
*val = cephes_exp_scalar(*val);
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_exp_inplace(x: &mut [f32]) {
use core::arch::x86_64::{
_mm256_add_epi32, _mm256_add_ps, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_loadu_ps,
_mm256_max_epi32, _mm256_min_epi32, _mm256_mul_ps, _mm256_round_ps, _mm256_set1_epi32,
_mm256_set1_ps, _mm256_slli_epi32, _mm256_storeu_ps, _mm256_sub_ps,
};
unsafe {
const ROUND_NEAREST: i32 = 0x00;
let v_inv_ln2 = _mm256_set1_ps(CEPHES_INV_LN2);
let v_ln2_hi = _mm256_set1_ps(CEPHES_LN2_HI);
let v_ln2_lo = _mm256_set1_ps(CEPHES_LN2_LO);
let v_one = _mm256_set1_ps(1.0);
let v_half = _mm256_set1_ps(0.5);
let v_third = _mm256_set1_ps(1.0 / 3.0);
let v_quarter = _mm256_set1_ps(0.25);
let v_fifth = _mm256_set1_ps(0.2);
let v_sixth = _mm256_set1_ps(1.0 / 6.0);
let mut i = 0;
let chunks = x.len() / 8;
for _ in 0..chunks {
let vx = _mm256_loadu_ps(x.as_ptr().add(i));
let vn_f = _mm256_round_ps(_mm256_mul_ps(vx, v_inv_ln2), ROUND_NEAREST);
let vn_i = _mm256_cvtps_epi32(vn_f);
let vg = _mm256_sub_ps(
_mm256_sub_ps(vx, _mm256_mul_ps(vn_f, v_ln2_hi)),
_mm256_mul_ps(vn_f, v_ln2_lo),
);
let p6 = _mm256_add_ps(v_one, _mm256_mul_ps(vg, v_sixth));
let p5 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_fifth), p6));
let p4 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_quarter), p5));
let p3 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_third), p4));
let p2 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_half), p3));
let q = _mm256_add_ps(v_one, _mm256_mul_ps(vg, p2));
let vn_clamped = _mm256_max_epi32(
_mm256_min_epi32(vn_i, _mm256_set1_epi32(127)),
_mm256_set1_epi32(-126),
);
let vn_shifted_i = _mm256_add_epi32(vn_clamped, _mm256_set1_epi32(127));
let v_scale_bits = _mm256_slli_epi32::<23>(vn_shifted_i);
let v_scale = _mm256_castsi256_ps(v_scale_bits);
let result = _mm256_mul_ps(v_scale, q);
_mm256_storeu_ps(x.as_mut_ptr().add(i), result);
i += 8;
}
while i < x.len() {
*x.get_unchecked_mut(i) = cephes_exp_scalar(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_sigmoid_tanh_clamp(out: &mut [f32], a: &[f32], q: &[f32], clamp: f32) {
use core::arch::x86_64::{
_mm256_add_epi32, _mm256_add_ps, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_div_ps,
_mm256_loadu_ps, _mm256_max_epi32, _mm256_max_ps, _mm256_min_epi32, _mm256_min_ps,
_mm256_mul_ps, _mm256_round_ps, _mm256_set1_epi32, _mm256_set1_ps, _mm256_slli_epi32,
_mm256_storeu_ps, _mm256_sub_ps, _mm256_xor_ps,
};
unsafe {
const ROUND_NEAREST: i32 = 0x00;
let v_inv_ln2 = _mm256_set1_ps(CEPHES_INV_LN2);
let v_ln2_hi = _mm256_set1_ps(CEPHES_LN2_HI);
let v_ln2_lo = _mm256_set1_ps(CEPHES_LN2_LO);
let v_one = _mm256_set1_ps(1.0);
let v_half = _mm256_set1_ps(0.5);
let v_third = _mm256_set1_ps(1.0 / 3.0);
let v_quarter = _mm256_set1_ps(0.25);
let v_fifth = _mm256_set1_ps(0.2);
let v_sixth = _mm256_set1_ps(1.0 / 6.0);
let v_two = _mm256_set1_ps(2.0);
let v_sign_flip = _mm256_set1_ps(f32::from_bits(0x8000_0000));
let v_clamp = _mm256_set1_ps(clamp);
let v_neg_clamp = _mm256_xor_ps(v_clamp, v_sign_flip);
let mut i = 0;
let chunks = out.len() / 8;
for _ in 0..chunks {
let va = _mm256_loadu_ps(a.as_ptr().add(i));
let vq = _mm256_loadu_ps(q.as_ptr().add(i));
let vy = _mm256_add_ps(va, vq);
let vx = _mm256_xor_ps(vy, v_sign_flip);
let vn_f = _mm256_round_ps(_mm256_mul_ps(vx, v_inv_ln2), ROUND_NEAREST);
let vn_i = _mm256_cvtps_epi32(vn_f);
let vg = _mm256_sub_ps(
_mm256_sub_ps(vx, _mm256_mul_ps(vn_f, v_ln2_hi)),
_mm256_mul_ps(vn_f, v_ln2_lo),
);
let gc_sixth = _mm256_mul_ps(vg, v_sixth);
let p6 = _mm256_add_ps(v_one, gc_sixth);
let gc_fifth = _mm256_mul_ps(vg, v_fifth);
let p5 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_fifth, p6));
let gc_quarter = _mm256_mul_ps(vg, v_quarter);
let p4 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_quarter, p5));
let gc_third = _mm256_mul_ps(vg, v_third);
let p3 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_third, p4));
let gc_half = _mm256_mul_ps(vg, v_half);
let p2 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_half, p3));
let qpoly = _mm256_add_ps(v_one, _mm256_mul_ps(vg, p2));
let vn_clamped = _mm256_max_epi32(
_mm256_min_epi32(vn_i, _mm256_set1_epi32(127)),
_mm256_set1_epi32(-126),
);
let vn_shifted_i = _mm256_add_epi32(vn_clamped, _mm256_set1_epi32(127));
let v_scale_bits = _mm256_slli_epi32::<23>(vn_shifted_i);
let v_scale = _mm256_castsi256_ps(v_scale_bits);
let exp_neg_y = _mm256_mul_ps(v_scale, qpoly);
let denom = _mm256_add_ps(v_one, exp_neg_y);
let sigma = _mm256_div_ps(v_one, denom);
let tanh_like = _mm256_sub_ps(_mm256_mul_ps(v_two, sigma), v_one);
let clamped = _mm256_max_ps(_mm256_min_ps(tanh_like, v_clamp), v_neg_clamp);
_mm256_storeu_ps(out.as_mut_ptr().add(i), clamped);
i += 8;
}
while i < out.len() {
let s = fast_sigmoid(*a.get_unchecked(i) + *q.get_unchecked(i));
let v = 2.0 * s - 1.0;
*out.get_unchecked_mut(i) = v.clamp(-clamp, clamp);
i += 1;
}
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_sigmoid_inplace(x: &mut [f32]) {
use core::arch::x86_64::{
_mm256_add_epi32, _mm256_add_ps, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_div_ps,
_mm256_loadu_ps, _mm256_max_epi32, _mm256_min_epi32, _mm256_mul_ps, _mm256_round_ps,
_mm256_set1_epi32, _mm256_set1_ps, _mm256_slli_epi32, _mm256_storeu_ps, _mm256_sub_ps,
_mm256_xor_ps,
};
unsafe {
const ROUND_NEAREST: i32 = 0x00;
let v_inv_ln2 = _mm256_set1_ps(CEPHES_INV_LN2);
let v_ln2_hi = _mm256_set1_ps(CEPHES_LN2_HI);
let v_ln2_lo = _mm256_set1_ps(CEPHES_LN2_LO);
let v_one = _mm256_set1_ps(1.0);
let v_half = _mm256_set1_ps(0.5);
let v_third = _mm256_set1_ps(1.0 / 3.0);
let v_quarter = _mm256_set1_ps(0.25);
let v_fifth = _mm256_set1_ps(0.2);
let v_sixth = _mm256_set1_ps(1.0 / 6.0);
let v_sign_flip = _mm256_set1_ps(f32::from_bits(0x8000_0000));
let mut i = 0;
let chunks = x.len() / 8;
for _ in 0..chunks {
let vx = _mm256_xor_ps(_mm256_loadu_ps(x.as_ptr().add(i)), v_sign_flip);
let vn_f = _mm256_round_ps(_mm256_mul_ps(vx, v_inv_ln2), ROUND_NEAREST);
let vn_i = _mm256_cvtps_epi32(vn_f);
let vg = _mm256_sub_ps(
_mm256_sub_ps(vx, _mm256_mul_ps(vn_f, v_ln2_hi)),
_mm256_mul_ps(vn_f, v_ln2_lo),
);
let gc_sixth = _mm256_mul_ps(vg, v_sixth);
let p6 = _mm256_add_ps(v_one, gc_sixth);
let gc_fifth = _mm256_mul_ps(vg, v_fifth);
let p5 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_fifth, p6));
let gc_quarter = _mm256_mul_ps(vg, v_quarter);
let p4 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_quarter, p5));
let gc_third = _mm256_mul_ps(vg, v_third);
let p3 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_third, p4));
let gc_half = _mm256_mul_ps(vg, v_half);
let p2 = _mm256_add_ps(v_one, _mm256_mul_ps(gc_half, p3));
let qpoly = _mm256_add_ps(v_one, _mm256_mul_ps(vg, p2));
let vn_clamped = _mm256_max_epi32(
_mm256_min_epi32(vn_i, _mm256_set1_epi32(127)),
_mm256_set1_epi32(-126),
);
let vn_shifted_i = _mm256_add_epi32(vn_clamped, _mm256_set1_epi32(127));
let v_scale_bits = _mm256_slli_epi32::<23>(vn_shifted_i);
let v_scale = _mm256_castsi256_ps(v_scale_bits);
let exp_neg_x = _mm256_mul_ps(v_scale, qpoly);
let denom = _mm256_add_ps(v_one, exp_neg_x);
let sigma = _mm256_div_ps(v_one, denom);
_mm256_storeu_ps(x.as_mut_ptr().add(i), sigma);
i += 8;
}
while i < x.len() {
*x.get_unchecked_mut(i) = fast_sigmoid(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
#[inline]
unsafe fn avx2_exp_sum_inplace(x: &mut [f32]) -> f32 {
use core::arch::x86_64::{
_mm256_add_epi32, _mm256_add_ps, _mm256_castsi256_ps, _mm256_cvtps_epi32, _mm256_loadu_ps,
_mm256_mul_ps, _mm256_round_ps, _mm256_set1_epi32, _mm256_set1_ps, _mm256_setzero_ps,
_mm256_slli_epi32, _mm256_storeu_ps, _mm256_sub_ps,
};
unsafe {
const ROUND_NEAREST: i32 = 0x00;
let v_inv_ln2 = _mm256_set1_ps(CEPHES_INV_LN2);
let v_ln2_hi = _mm256_set1_ps(CEPHES_LN2_HI);
let v_ln2_lo = _mm256_set1_ps(CEPHES_LN2_LO);
let v_one = _mm256_set1_ps(1.0);
let v_half = _mm256_set1_ps(0.5);
let v_third = _mm256_set1_ps(1.0 / 3.0);
let v_quarter = _mm256_set1_ps(0.25);
let v_fifth = _mm256_set1_ps(0.2);
let v_sixth = _mm256_set1_ps(1.0 / 6.0);
let mut acc0 = _mm256_setzero_ps();
let mut acc1 = _mm256_setzero_ps();
let mut acc2 = _mm256_setzero_ps();
let mut acc3 = _mm256_setzero_ps();
let mut i = 0;
let len = x.len();
let chunks4 = len / 32;
macro_rules! step {
($acc:expr, $off:expr) => {{
let vx = _mm256_loadu_ps(x.as_ptr().add(i + $off));
let vn_f = _mm256_round_ps(_mm256_mul_ps(vx, v_inv_ln2), ROUND_NEAREST);
let vn_i = _mm256_cvtps_epi32(vn_f);
let vg = _mm256_sub_ps(
_mm256_sub_ps(vx, _mm256_mul_ps(vn_f, v_ln2_hi)),
_mm256_mul_ps(vn_f, v_ln2_lo),
);
let p6 = _mm256_add_ps(v_one, _mm256_mul_ps(vg, v_sixth));
let p5 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_fifth), p6));
let p4 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_quarter), p5));
let p3 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_third), p4));
let p2 = _mm256_add_ps(v_one, _mm256_mul_ps(_mm256_mul_ps(vg, v_half), p3));
let q = _mm256_add_ps(v_one, _mm256_mul_ps(vg, p2));
let vn_shifted_i = _mm256_add_epi32(vn_i, _mm256_set1_epi32(127));
let v_scale_bits = _mm256_slli_epi32::<23>(vn_shifted_i);
let v_scale = _mm256_castsi256_ps(v_scale_bits);
let r = _mm256_mul_ps(v_scale, q);
_mm256_storeu_ps(x.as_mut_ptr().add(i + $off), r);
$acc = _mm256_add_ps($acc, r);
}};
}
for _ in 0..chunks4 {
step!(acc0, 0);
step!(acc1, 8);
step!(acc2, 16);
step!(acc3, 24);
i += 32;
}
let mut sum = horizontal_sum_256(_mm256_add_ps(
_mm256_add_ps(acc0, acc1),
_mm256_add_ps(acc2, acc3),
));
let mut acc_rem = _mm256_setzero_ps();
let remaining = (len - i) / 8;
for _ in 0..remaining {
step!(acc_rem, 0);
i += 8;
}
sum += horizontal_sum_256(acc_rem);
while i < len {
let e = cephes_exp_scalar(*x.get_unchecked(i));
*x.get_unchecked_mut(i) = e;
sum += e;
i += 1;
}
sum
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_exp_inplace(x: &mut [f32]) {
use core::arch::aarch64::{
vaddq_f32, vaddq_s32, vcvtq_s32_f32, vdupq_n_f32, vdupq_n_s32, vld1q_f32, vmaxq_s32,
vminq_s32, vmulq_f32, vreinterpretq_f32_s32, vrndq_f32, vshlq_n_s32, vst1q_f32, vsubq_f32,
};
unsafe {
let v_inv_ln2 = vdupq_n_f32(CEPHES_INV_LN2);
let v_ln2_hi = vdupq_n_f32(CEPHES_LN2_HI);
let v_ln2_lo = vdupq_n_f32(CEPHES_LN2_LO);
let v_one = vdupq_n_f32(1.0);
let v_half = vdupq_n_f32(0.5);
let v_third = vdupq_n_f32(1.0 / 3.0);
let v_quarter = vdupq_n_f32(0.25);
let v_fifth = vdupq_n_f32(0.2);
let v_sixth = vdupq_n_f32(1.0 / 6.0);
let mut i = 0;
let chunks = x.len() / 4;
for _ in 0..chunks {
let vx = vld1q_f32(x.as_ptr().add(i));
let vn_f = vrndq_f32(vmulq_f32(vx, v_inv_ln2));
let vn_i = vcvtq_s32_f32(vn_f);
let vg = vsubq_f32(
vsubq_f32(vx, vmulq_f32(vn_f, v_ln2_hi)),
vmulq_f32(vn_f, v_ln2_lo),
);
let p6 = vaddq_f32(v_one, vmulq_f32(vg, v_sixth)); let p5 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_fifth), p6)); let p4 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_quarter), p5)); let p3 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_third), p4)); let p2 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_half), p3)); let q = vaddq_f32(v_one, vmulq_f32(vg, p2));
let v127 = vdupq_n_s32(127);
let vneg126 = vdupq_n_s32(-126);
let vn_clamped = vmaxq_s32(vminq_s32(vn_i, v127), vneg126);
let v_bias = vdupq_n_s32(127);
let v_shifted = vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(vn_clamped, v_bias)));
let vresult = vmulq_f32(v_shifted, q);
vst1q_f32(x.as_mut_ptr().add(i), vresult);
i += 4;
}
while i < x.len() {
*x.get_unchecked_mut(i) = cephes_exp_scalar(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_sigmoid_tanh_clamp(out: &mut [f32], a: &[f32], q: &[f32], clamp: f32) {
use core::arch::aarch64::{
vaddq_f32, vaddq_s32, vcvtq_s32_f32, vdivq_f32, vdupq_n_f32, vdupq_n_s32, vld1q_f32,
vmaxq_f32, vmaxq_s32, vminq_f32, vminq_s32, vmulq_f32, vnegq_f32, vreinterpretq_f32_s32,
vrndq_f32, vshlq_n_s32, vst1q_f32, vsubq_f32,
};
unsafe {
let v_inv_ln2 = vdupq_n_f32(CEPHES_INV_LN2);
let v_ln2_hi = vdupq_n_f32(CEPHES_LN2_HI);
let v_ln2_lo = vdupq_n_f32(CEPHES_LN2_LO);
let v_one = vdupq_n_f32(1.0);
let v_half = vdupq_n_f32(0.5);
let v_third = vdupq_n_f32(1.0 / 3.0);
let v_quarter = vdupq_n_f32(0.25);
let v_fifth = vdupq_n_f32(0.2);
let v_sixth = vdupq_n_f32(1.0 / 6.0);
let v_two = vdupq_n_f32(2.0);
let v_clamp = vdupq_n_f32(clamp);
let v_neg_clamp = vnegq_f32(v_clamp);
let mut i = 0;
let chunks = out.len() / 4;
for _ in 0..chunks {
let va = vld1q_f32(a.as_ptr().add(i));
let vq = vld1q_f32(q.as_ptr().add(i));
let vy = vaddq_f32(va, vq);
let vx = vnegq_f32(vy);
let vn_f = vrndq_f32(vmulq_f32(vx, v_inv_ln2));
let vn_i = vcvtq_s32_f32(vn_f);
let vg = vsubq_f32(
vsubq_f32(vx, vmulq_f32(vn_f, v_ln2_hi)),
vmulq_f32(vn_f, v_ln2_lo),
);
let gc_sixth = vmulq_f32(vg, v_sixth);
let p6 = vaddq_f32(v_one, gc_sixth); let gc_fifth = vmulq_f32(vg, v_fifth);
let p5 = vaddq_f32(v_one, vmulq_f32(gc_fifth, p6)); let gc_quarter = vmulq_f32(vg, v_quarter);
let p4 = vaddq_f32(v_one, vmulq_f32(gc_quarter, p5)); let gc_third = vmulq_f32(vg, v_third);
let p3 = vaddq_f32(v_one, vmulq_f32(gc_third, p4)); let gc_half = vmulq_f32(vg, v_half);
let p2 = vaddq_f32(v_one, vmulq_f32(gc_half, p3)); let qpoly = vaddq_f32(v_one, vmulq_f32(vg, p2));
let v127 = vdupq_n_s32(127);
let vneg126 = vdupq_n_s32(-126);
let vn_clamped = vmaxq_s32(vminq_s32(vn_i, v127), vneg126);
let v_bias = vdupq_n_s32(127);
let v_shifted = vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(vn_clamped, v_bias)));
let exp_neg_y = vmulq_f32(v_shifted, qpoly);
let denom = vaddq_f32(v_one, exp_neg_y);
let sigma = vdivq_f32(v_one, denom);
let tanh_like = vsubq_f32(vmulq_f32(v_two, sigma), v_one);
let clamped = vmaxq_f32(vminq_f32(tanh_like, v_clamp), v_neg_clamp);
vst1q_f32(out.as_mut_ptr().add(i), clamped);
i += 4;
}
while i < out.len() {
let s = fast_sigmoid(*a.get_unchecked(i) + *q.get_unchecked(i));
let v = 2.0 * s - 1.0;
*out.get_unchecked_mut(i) = v.clamp(-clamp, clamp);
i += 1;
}
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_sigmoid_inplace(x: &mut [f32]) {
use core::arch::aarch64::{
vaddq_f32, vaddq_s32, vcvtq_s32_f32, vdivq_f32, vdupq_n_f32, vdupq_n_s32, vld1q_f32,
vmaxq_s32, vminq_s32, vmulq_f32, vnegq_f32, vreinterpretq_f32_s32, vrndq_f32, vshlq_n_s32,
vst1q_f32, vsubq_f32,
};
unsafe {
let v_inv_ln2 = vdupq_n_f32(CEPHES_INV_LN2);
let v_ln2_hi = vdupq_n_f32(CEPHES_LN2_HI);
let v_ln2_lo = vdupq_n_f32(CEPHES_LN2_LO);
let v_one = vdupq_n_f32(1.0);
let v_half = vdupq_n_f32(0.5);
let v_third = vdupq_n_f32(1.0 / 3.0);
let v_quarter = vdupq_n_f32(0.25);
let v_fifth = vdupq_n_f32(0.2);
let v_sixth = vdupq_n_f32(1.0 / 6.0);
let v127 = vdupq_n_s32(127);
let vneg126 = vdupq_n_s32(-126);
let v_bias = vdupq_n_s32(127);
let mut i = 0;
let chunks = x.len() / 4;
for _ in 0..chunks {
let vx = vnegq_f32(vld1q_f32(x.as_ptr().add(i)));
let vn_f = vrndq_f32(vmulq_f32(vx, v_inv_ln2));
let vn_i = vcvtq_s32_f32(vn_f);
let vg = vsubq_f32(
vsubq_f32(vx, vmulq_f32(vn_f, v_ln2_hi)),
vmulq_f32(vn_f, v_ln2_lo),
);
let gc_sixth = vmulq_f32(vg, v_sixth);
let p6 = vaddq_f32(v_one, gc_sixth);
let gc_fifth = vmulq_f32(vg, v_fifth);
let p5 = vaddq_f32(v_one, vmulq_f32(gc_fifth, p6));
let gc_quarter = vmulq_f32(vg, v_quarter);
let p4 = vaddq_f32(v_one, vmulq_f32(gc_quarter, p5));
let gc_third = vmulq_f32(vg, v_third);
let p3 = vaddq_f32(v_one, vmulq_f32(gc_third, p4));
let gc_half = vmulq_f32(vg, v_half);
let p2 = vaddq_f32(v_one, vmulq_f32(gc_half, p3));
let qpoly = vaddq_f32(v_one, vmulq_f32(vg, p2));
let vn_clamped = vmaxq_s32(vminq_s32(vn_i, v127), vneg126);
let v_shifted = vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(vn_clamped, v_bias)));
let exp_neg_x = vmulq_f32(v_shifted, qpoly);
let denom = vaddq_f32(v_one, exp_neg_x);
let sigma = vdivq_f32(v_one, denom);
vst1q_f32(x.as_mut_ptr().add(i), sigma);
i += 4;
}
while i < x.len() {
*x.get_unchecked_mut(i) = fast_sigmoid(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn neon_exp_sum_inplace(x: &mut [f32]) -> f32 {
use core::arch::aarch64::{
vaddq_f32, vaddq_s32, vaddvq_f32, vcvtq_s32_f32, vdupq_n_f32, vdupq_n_s32, vld1q_f32,
vmaxq_s32, vminq_s32, vmulq_f32, vreinterpretq_f32_s32, vrndq_f32, vshlq_n_s32, vst1q_f32,
vsubq_f32,
};
unsafe {
let v_inv_ln2 = vdupq_n_f32(CEPHES_INV_LN2);
let v_ln2_hi = vdupq_n_f32(CEPHES_LN2_HI);
let v_ln2_lo = vdupq_n_f32(CEPHES_LN2_LO);
let v_one = vdupq_n_f32(1.0);
let v_half = vdupq_n_f32(0.5);
let v_third = vdupq_n_f32(1.0 / 3.0);
let v_quarter = vdupq_n_f32(0.25);
let v_fifth = vdupq_n_f32(0.2);
let v_sixth = vdupq_n_f32(1.0 / 6.0);
let v127 = vdupq_n_s32(127);
let vneg126 = vdupq_n_s32(-126);
let v_bias = vdupq_n_s32(127);
let mut acc0 = vdupq_n_f32(0.0);
let mut acc1 = vdupq_n_f32(0.0);
let mut acc2 = vdupq_n_f32(0.0);
let mut acc3 = vdupq_n_f32(0.0);
let mut i = 0;
let len = x.len();
let chunks4 = len / 16;
for _ in 0..chunks4 {
macro_rules! step {
($acc:expr, $off:expr) => {{
let vx = vld1q_f32(x.as_ptr().add(i + $off));
let vn_f = vrndq_f32(vmulq_f32(vx, v_inv_ln2));
let vn_i = vcvtq_s32_f32(vn_f);
let vg = vsubq_f32(
vsubq_f32(vx, vmulq_f32(vn_f, v_ln2_hi)),
vmulq_f32(vn_f, v_ln2_lo),
);
let p6 = vaddq_f32(v_one, vmulq_f32(vg, v_sixth));
let p5 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_fifth), p6));
let p4 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_quarter), p5));
let p3 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_third), p4));
let p2 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_half), p3));
let q = vaddq_f32(v_one, vmulq_f32(vg, p2));
let vn_clamped = vmaxq_s32(vminq_s32(vn_i, v127), vneg126);
let v_shifted =
vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(vn_clamped, v_bias)));
let r = vmulq_f32(v_shifted, q);
vst1q_f32(x.as_mut_ptr().add(i + $off), r);
$acc = vaddq_f32($acc, r);
}};
}
step!(acc0, 0);
step!(acc1, 4);
step!(acc2, 8);
step!(acc3, 12);
i += 16;
}
let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3)));
let mut acc_rem = vdupq_n_f32(0.0);
let remaining = (len - i) / 4;
for _ in 0..remaining {
let vx = vld1q_f32(x.as_ptr().add(i));
let vn_f = vrndq_f32(vmulq_f32(vx, v_inv_ln2));
let vn_i = vcvtq_s32_f32(vn_f);
let vg = vsubq_f32(
vsubq_f32(vx, vmulq_f32(vn_f, v_ln2_hi)),
vmulq_f32(vn_f, v_ln2_lo),
);
let p6 = vaddq_f32(v_one, vmulq_f32(vg, v_sixth));
let p5 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_fifth), p6));
let p4 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_quarter), p5));
let p3 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_third), p4));
let p2 = vaddq_f32(v_one, vmulq_f32(vmulq_f32(vg, v_half), p3));
let q = vaddq_f32(v_one, vmulq_f32(vg, p2));
let vn_clamped = vmaxq_s32(vminq_s32(vn_i, v127), vneg126);
let v_shifted = vreinterpretq_f32_s32(vshlq_n_s32::<23>(vaddq_s32(vn_clamped, v_bias)));
let r = vmulq_f32(v_shifted, q);
vst1q_f32(x.as_mut_ptr().add(i), r);
acc_rem = vaddq_f32(acc_rem, r);
i += 4;
}
sum += vaddvq_f32(acc_rem);
while i < len {
let e = cephes_exp_scalar(*x.get_unchecked(i));
*x.get_unchecked_mut(i) = e;
sum += e;
i += 1;
}
sum
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_reciprocal_inplace(x: &mut [f32]) {
use core::arch::wasm32::{f32x4_div, f32x4_splat, v128_load, v128_store};
unsafe {
let len = x.len();
let chunks = len / 4;
let ones = f32x4_splat(1.0);
for i in 0..chunks {
let v = v128_load(x.as_ptr().add(i * 4).cast());
let r = f32x4_div(ones, v);
v128_store(x.as_mut_ptr().add(i * 4).cast(), r);
}
for i in (chunks * 4)..len {
*x.get_unchecked_mut(i) = 1.0 / *x.get_unchecked(i);
}
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_exp_inplace(x: &mut [f32]) {
use core::arch::wasm32::{
f32x4_add, f32x4_mul, f32x4_nearest, f32x4_splat, f32x4_sub, i32x4_add, i32x4_max,
i32x4_min, i32x4_shl, i32x4_splat, i32x4_trunc_sat_f32x4, v128_load, v128_store,
};
unsafe {
let v_inv_ln2 = f32x4_splat(CEPHES_INV_LN2);
let v_ln2_hi = f32x4_splat(CEPHES_LN2_HI);
let v_ln2_lo = f32x4_splat(CEPHES_LN2_LO);
let v_one = f32x4_splat(1.0);
let v_half = f32x4_splat(0.5);
let v_third = f32x4_splat(1.0 / 3.0);
let v_quarter = f32x4_splat(0.25);
let v_fifth = f32x4_splat(0.2);
let v_sixth = f32x4_splat(1.0 / 6.0);
let v127 = i32x4_splat(127);
let vneg126 = i32x4_splat(-126);
let v_bias = i32x4_splat(127);
let mut i = 0;
let chunks = x.len() / 4;
for _ in 0..chunks {
let vx = v128_load(x.as_ptr().add(i).cast());
let vn_f = f32x4_nearest(f32x4_mul(vx, v_inv_ln2));
let vn_i = i32x4_trunc_sat_f32x4(vn_f);
let vg = f32x4_sub(
f32x4_sub(vx, f32x4_mul(vn_f, v_ln2_hi)),
f32x4_mul(vn_f, v_ln2_lo),
);
let p6 = f32x4_add(v_one, f32x4_mul(vg, v_sixth)); let p5 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_fifth), p6)); let p4 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_quarter), p5)); let p3 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_third), p4)); let p2 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_half), p3)); let q = f32x4_add(v_one, f32x4_mul(vg, p2));
let vn_clamped = i32x4_max(i32x4_min(vn_i, v127), vneg126);
let v_shifted = i32x4_shl(i32x4_add(vn_clamped, v_bias), 23);
let vresult = f32x4_mul(v_shifted, q);
v128_store(x.as_mut_ptr().add(i).cast(), vresult);
i += 4;
}
while i < x.len() {
*x.get_unchecked_mut(i) = cephes_exp_scalar(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_exp_sum_inplace(x: &mut [f32]) -> f32 {
use core::arch::wasm32::{
f32x4_add, f32x4_extract_lane, f32x4_mul, f32x4_nearest, f32x4_splat, f32x4_sub, i32x4_add,
i32x4_max, i32x4_min, i32x4_shl, i32x4_splat, i32x4_trunc_sat_f32x4, v128_load, v128_store,
};
unsafe {
let v_inv_ln2 = f32x4_splat(CEPHES_INV_LN2);
let v_ln2_hi = f32x4_splat(CEPHES_LN2_HI);
let v_ln2_lo = f32x4_splat(CEPHES_LN2_LO);
let v_one = f32x4_splat(1.0);
let v_half = f32x4_splat(0.5);
let v_third = f32x4_splat(1.0 / 3.0);
let v_quarter = f32x4_splat(0.25);
let v_fifth = f32x4_splat(0.2);
let v_sixth = f32x4_splat(1.0 / 6.0);
let v127 = i32x4_splat(127);
let vneg126 = i32x4_splat(-126);
let v_bias = i32x4_splat(127);
let mut acc0 = f32x4_splat(0.0);
let mut acc1 = f32x4_splat(0.0);
let mut acc2 = f32x4_splat(0.0);
let mut acc3 = f32x4_splat(0.0);
let mut i = 0;
let len = x.len();
let chunks4 = len / 16;
for _ in 0..chunks4 {
macro_rules! step {
($acc:expr, $off:expr) => {{
let vx = v128_load(x.as_ptr().add(i + $off).cast());
let vn_f = f32x4_nearest(f32x4_mul(vx, v_inv_ln2));
let vn_i = i32x4_trunc_sat_f32x4(vn_f);
let vg = f32x4_sub(
f32x4_sub(vx, f32x4_mul(vn_f, v_ln2_hi)),
f32x4_mul(vn_f, v_ln2_lo),
);
let p6 = f32x4_add(v_one, f32x4_mul(vg, v_sixth));
let p5 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_fifth), p6));
let p4 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_quarter), p5));
let p3 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_third), p4));
let p2 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_half), p3));
let q = f32x4_add(v_one, f32x4_mul(vg, p2));
let vn_clamped = i32x4_max(i32x4_min(vn_i, v127), vneg126);
let v_shifted = i32x4_shl(i32x4_add(vn_clamped, v_bias), 23);
let r = f32x4_mul(v_shifted, q);
v128_store(x.as_mut_ptr().add(i + $off).cast(), r);
$acc = f32x4_add($acc, r);
}};
}
step!(acc0, 0);
step!(acc1, 4);
step!(acc2, 8);
step!(acc3, 12);
i += 16;
}
let s01 = f32x4_add(acc0, acc1);
let s23 = f32x4_add(acc2, acc3);
let s = f32x4_add(s01, s23);
let mut sum = f32x4_extract_lane::<0>(s)
+ f32x4_extract_lane::<1>(s)
+ f32x4_extract_lane::<2>(s)
+ f32x4_extract_lane::<3>(s);
let mut acc_rem = f32x4_splat(0.0);
let remaining = (len - i) / 4;
for _ in 0..remaining {
let vx = v128_load(x.as_ptr().add(i).cast());
let vn_f = f32x4_nearest(f32x4_mul(vx, v_inv_ln2));
let vn_i = i32x4_trunc_sat_f32x4(vn_f);
let vg = f32x4_sub(
f32x4_sub(vx, f32x4_mul(vn_f, v_ln2_hi)),
f32x4_mul(vn_f, v_ln2_lo),
);
let p6 = f32x4_add(v_one, f32x4_mul(vg, v_sixth));
let p5 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_fifth), p6));
let p4 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_quarter), p5));
let p3 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_third), p4));
let p2 = f32x4_add(v_one, f32x4_mul(f32x4_mul(vg, v_half), p3));
let q = f32x4_add(v_one, f32x4_mul(vg, p2));
let vn_clamped = i32x4_max(i32x4_min(vn_i, v127), vneg126);
let v_shifted = i32x4_shl(i32x4_add(vn_clamped, v_bias), 23);
let r = f32x4_mul(v_shifted, q);
v128_store(x.as_mut_ptr().add(i).cast(), r);
acc_rem = f32x4_add(acc_rem, r);
i += 4;
}
sum += f32x4_extract_lane::<0>(acc_rem)
+ f32x4_extract_lane::<1>(acc_rem)
+ f32x4_extract_lane::<2>(acc_rem)
+ f32x4_extract_lane::<3>(acc_rem);
while i < len {
let e = cephes_exp_scalar(*x.get_unchecked(i));
*x.get_unchecked_mut(i) = e;
sum += e;
i += 1;
}
sum
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_sigmoid_inplace(x: &mut [f32]) {
use core::arch::wasm32::{
f32x4_add, f32x4_div, f32x4_mul, f32x4_nearest, f32x4_neg, f32x4_splat, f32x4_sub,
i32x4_add, i32x4_max, i32x4_min, i32x4_shl, i32x4_splat, i32x4_trunc_sat_f32x4, v128_load,
v128_store,
};
unsafe {
let v_inv_ln2 = f32x4_splat(CEPHES_INV_LN2);
let v_ln2_hi = f32x4_splat(CEPHES_LN2_HI);
let v_ln2_lo = f32x4_splat(CEPHES_LN2_LO);
let v_one = f32x4_splat(1.0);
let v_half = f32x4_splat(0.5);
let v_third = f32x4_splat(1.0 / 3.0);
let v_quarter = f32x4_splat(0.25);
let v_fifth = f32x4_splat(0.2);
let v_sixth = f32x4_splat(1.0 / 6.0);
let v127 = i32x4_splat(127);
let vneg126 = i32x4_splat(-126);
let v_bias = i32x4_splat(127);
let mut i = 0;
let chunks = x.len() / 4;
for _ in 0..chunks {
let vx = f32x4_neg(v128_load(x.as_ptr().add(i).cast()));
let vn_f = f32x4_nearest(f32x4_mul(vx, v_inv_ln2));
let vn_i = i32x4_trunc_sat_f32x4(vn_f);
let vg = f32x4_sub(
f32x4_sub(vx, f32x4_mul(vn_f, v_ln2_hi)),
f32x4_mul(vn_f, v_ln2_lo),
);
let gc_sixth = f32x4_mul(vg, v_sixth);
let p6 = f32x4_add(v_one, gc_sixth);
let gc_fifth = f32x4_mul(vg, v_fifth);
let p5 = f32x4_add(v_one, f32x4_mul(gc_fifth, p6));
let gc_quarter = f32x4_mul(vg, v_quarter);
let p4 = f32x4_add(v_one, f32x4_mul(gc_quarter, p5));
let gc_third = f32x4_mul(vg, v_third);
let p3 = f32x4_add(v_one, f32x4_mul(gc_third, p4));
let gc_half = f32x4_mul(vg, v_half);
let p2 = f32x4_add(v_one, f32x4_mul(gc_half, p3));
let qpoly = f32x4_add(v_one, f32x4_mul(vg, p2));
let vn_clamped = i32x4_max(i32x4_min(vn_i, v127), vneg126);
let v_shifted = i32x4_shl(i32x4_add(vn_clamped, v_bias), 23);
let exp_neg_x = f32x4_mul(v_shifted, qpoly);
let denom = f32x4_add(v_one, exp_neg_x);
let sigma = f32x4_div(v_one, denom);
v128_store(x.as_mut_ptr().add(i).cast(), sigma);
i += 4;
}
while i < x.len() {
*x.get_unchecked_mut(i) = fast_sigmoid(*x.get_unchecked(i));
i += 1;
}
}
}
#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[inline]
unsafe fn wasm32_sigmoid_tanh_clamp(out: &mut [f32], a: &[f32], q: &[f32], clamp: f32) {
use core::arch::wasm32::{
f32x4_add, f32x4_div, f32x4_max, f32x4_min, f32x4_mul, f32x4_nearest, f32x4_neg,
f32x4_splat, f32x4_sub, i32x4_add, i32x4_max, i32x4_min, i32x4_shl, i32x4_splat,
i32x4_trunc_sat_f32x4, v128_load, v128_store,
};
unsafe {
let v_inv_ln2 = f32x4_splat(CEPHES_INV_LN2);
let v_ln2_hi = f32x4_splat(CEPHES_LN2_HI);
let v_ln2_lo = f32x4_splat(CEPHES_LN2_LO);
let v_one = f32x4_splat(1.0);
let v_half = f32x4_splat(0.5);
let v_third = f32x4_splat(1.0 / 3.0);
let v_quarter = f32x4_splat(0.25);
let v_fifth = f32x4_splat(0.2);
let v_sixth = f32x4_splat(1.0 / 6.0);
let v_two = f32x4_splat(2.0);
let v_clamp = f32x4_splat(clamp);
let v_neg_clamp = f32x4_neg(v_clamp);
let v127 = i32x4_splat(127);
let vneg126 = i32x4_splat(-126);
let v_bias = i32x4_splat(127);
let mut i = 0;
let chunks = out.len() / 4;
for _ in 0..chunks {
let va = v128_load(a.as_ptr().add(i).cast());
let vq = v128_load(q.as_ptr().add(i).cast());
let vy = f32x4_add(va, vq);
let vx = f32x4_neg(vy);
let vn_f = f32x4_nearest(f32x4_mul(vx, v_inv_ln2));
let vn_i = i32x4_trunc_sat_f32x4(vn_f);
let vg = f32x4_sub(
f32x4_sub(vx, f32x4_mul(vn_f, v_ln2_hi)),
f32x4_mul(vn_f, v_ln2_lo),
);
let gc_sixth = f32x4_mul(vg, v_sixth);
let p6 = f32x4_add(v_one, gc_sixth);
let gc_fifth = f32x4_mul(vg, v_fifth);
let p5 = f32x4_add(v_one, f32x4_mul(gc_fifth, p6));
let gc_quarter = f32x4_mul(vg, v_quarter);
let p4 = f32x4_add(v_one, f32x4_mul(gc_quarter, p5));
let gc_third = f32x4_mul(vg, v_third);
let p3 = f32x4_add(v_one, f32x4_mul(gc_third, p4));
let gc_half = f32x4_mul(vg, v_half);
let p2 = f32x4_add(v_one, f32x4_mul(gc_half, p3));
let qpoly = f32x4_add(v_one, f32x4_mul(vg, p2));
let vn_clamped = i32x4_max(i32x4_min(vn_i, v127), vneg126);
let v_shifted = i32x4_shl(i32x4_add(vn_clamped, v_bias), 23);
let exp_neg_y = f32x4_mul(v_shifted, qpoly);
let denom = f32x4_add(v_one, exp_neg_y);
let sigma = f32x4_div(v_one, denom);
let tanh_like = f32x4_sub(f32x4_mul(v_two, sigma), v_one);
let clamped = f32x4_max(f32x4_min(tanh_like, v_clamp), v_neg_clamp);
v128_store(out.as_mut_ptr().add(i).cast(), clamped);
i += 4;
}
while i < out.len() {
let s = fast_sigmoid(*a.get_unchecked(i) + *q.get_unchecked(i));
let v = 2.0 * s - 1.0;
*out.get_unchecked_mut(i) = v.clamp(-clamp, clamp);
i += 1;
}
}
}