Skip to main content

tract_linalg/x86_64/
ln.rs

1//! f32 ln kernels: the [`crate::generic::ln`] fit, eight or sixteen lanes at a time.
2
3routine_ew_rust!(x86_64;
4    f32,
5    x86_64_fma_ln_f32_32n,
6    32,
7    8,
8    #[inline(never)]
9    fn run(buf: &mut [f32], _: ()) {
10        debug_assert!(buf.len() % Self::nr() == 0);
11        debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
12        unsafe { x86_64_fma_ln_f32_32n_run(buf) }
13    },
14    func(Ln),
15    isa(X86_64Avx2, X86_64Fma)
16);
17
18#[cfg(target_arch = "x86_64")]
19#[target_feature(enable = "avx2,fma")]
20unsafe fn x86_64_fma_ln_f32_32n_run(buf: &mut [f32]) {
21    use crate::generic::ln::{LN2_HI, LN2_LO, POLY, SPLIT, SUBNORMAL_SCALE, SUBNORMAL_SHIFT};
22    use std::arch::x86_64::*;
23    #[inline(always)]
24    unsafe fn ln8(x: __m256) -> __m256 {
25        unsafe {
26            let one = _mm256_set1_ps(1.0);
27            let zero = _mm256_setzero_ps();
28            let subnormal = _mm256_cmp_ps::<_CMP_LT_OQ>(x, _mm256_set1_ps(f32::MIN_POSITIVE));
29            let scaled =
30                _mm256_blendv_ps(x, _mm256_mul_ps(x, _mm256_set1_ps(SUBNORMAL_SCALE)), subnormal);
31            let bits = _mm256_castps_si256(scaled);
32            let mut exponent =
33                _mm256_sub_epi32(_mm256_srli_epi32::<23>(bits), _mm256_set1_epi32(127));
34            exponent = _mm256_add_epi32(
35                exponent,
36                _mm256_and_si256(
37                    _mm256_castps_si256(subnormal),
38                    _mm256_set1_epi32(-SUBNORMAL_SHIFT),
39                ),
40            );
41            let mantissa = _mm256_or_ps(
42                _mm256_and_ps(scaled, _mm256_castsi256_ps(_mm256_set1_epi32(0x007fffff))),
43                one,
44            );
45            let split = _mm256_cmp_ps::<_CMP_GT_OQ>(mantissa, _mm256_set1_ps(SPLIT));
46            let mantissa =
47                _mm256_blendv_ps(mantissa, _mm256_mul_ps(mantissa, _mm256_set1_ps(0.5)), split);
48            exponent = _mm256_add_epi32(
49                exponent,
50                _mm256_and_si256(_mm256_castps_si256(split), _mm256_set1_epi32(1)),
51            );
52            let e = _mm256_cvtepi32_ps(exponent);
53            let f = _mm256_sub_ps(mantissa, one);
54            let mut p = _mm256_set1_ps(POLY[0]);
55            for c in &POLY[1..] {
56                p = _mm256_fmadd_ps(p, f, _mm256_set1_ps(*c));
57            }
58            let f2 = _mm256_mul_ps(f, f);
59            let mut y = _mm256_mul_ps(_mm256_mul_ps(p, f2), f);
60            y = _mm256_fmadd_ps(e, _mm256_set1_ps(LN2_LO), y);
61            y = _mm256_add_ps(_mm256_fnmadd_ps(_mm256_set1_ps(0.5), f2, y), f);
62            y = _mm256_fmadd_ps(e, _mm256_set1_ps(LN2_HI), y);
63            // A NaN compares false whichever way round, so `not greater than zero` is what
64            // gathers the negatives, the zeros and the NaNs in one mask.
65            let outside = _mm256_cmp_ps::<_CMP_NGT_UQ>(x, zero);
66            let special = _mm256_blendv_ps(
67                _mm256_set1_ps(f32::NAN),
68                _mm256_set1_ps(f32::NEG_INFINITY),
69                _mm256_cmp_ps::<_CMP_EQ_OQ>(x, zero),
70            );
71            y = _mm256_blendv_ps(y, special, outside);
72            _mm256_blendv_ps(
73                y,
74                _mm256_set1_ps(f32::INFINITY),
75                _mm256_cmp_ps::<_CMP_EQ_OQ>(x, _mm256_set1_ps(f32::INFINITY)),
76            )
77        }
78    }
79    unsafe {
80        let p = buf.as_mut_ptr();
81        for i in (0..buf.len()).step_by(8) {
82            _mm256_store_ps(p.add(i), ln8(_mm256_load_ps(p.add(i))));
83        }
84    }
85}
86
87routine_ew_rust!(x86_64;
88    f32,
89    x86_64_avx512_ln_f32_64n,
90    64,
91    16,
92    #[inline(never)]
93    fn run(buf: &mut [f32], _: ()) {
94        debug_assert!(buf.len() % Self::nr() == 0);
95        debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
96        unsafe { x86_64_avx512_ln_f32_64n_run(buf) }
97    },
98    func(Ln),
99    isa(X86_64Avx512f)
100);
101
102#[cfg(target_arch = "x86_64")]
103#[target_feature(enable = "avx512f")]
104unsafe fn x86_64_avx512_ln_f32_64n_run(buf: &mut [f32]) {
105    use crate::generic::ln::{LN2_HI, LN2_LO, POLY, SPLIT, SUBNORMAL_SCALE, SUBNORMAL_SHIFT};
106    use std::arch::x86_64::*;
107    #[inline(always)]
108    unsafe fn ln16(x: __m512) -> __m512 {
109        unsafe {
110            let one = _mm512_set1_ps(1.0);
111            let zero = _mm512_setzero_ps();
112            let subnormal = _mm512_cmp_ps_mask::<_CMP_LT_OQ>(x, _mm512_set1_ps(f32::MIN_POSITIVE));
113            let scaled = _mm512_mask_mul_ps(x, subnormal, x, _mm512_set1_ps(SUBNORMAL_SCALE));
114            let bits = _mm512_castps_si512(scaled);
115            let mut exponent =
116                _mm512_sub_epi32(_mm512_srli_epi32::<23>(bits), _mm512_set1_epi32(127));
117            exponent = _mm512_mask_sub_epi32(
118                exponent,
119                subnormal,
120                exponent,
121                _mm512_set1_epi32(SUBNORMAL_SHIFT),
122            );
123            let mantissa = _mm512_or_ps(
124                _mm512_and_ps(scaled, _mm512_castsi512_ps(_mm512_set1_epi32(0x007fffff))),
125                one,
126            );
127            let split = _mm512_cmp_ps_mask::<_CMP_GT_OQ>(mantissa, _mm512_set1_ps(SPLIT));
128            let mantissa = _mm512_mask_mul_ps(mantissa, split, mantissa, _mm512_set1_ps(0.5));
129            exponent = _mm512_mask_add_epi32(exponent, split, exponent, _mm512_set1_epi32(1));
130            let e = _mm512_cvtepi32_ps(exponent);
131            let f = _mm512_sub_ps(mantissa, one);
132            let mut p = _mm512_set1_ps(POLY[0]);
133            for c in &POLY[1..] {
134                p = _mm512_fmadd_ps(p, f, _mm512_set1_ps(*c));
135            }
136            let f2 = _mm512_mul_ps(f, f);
137            let mut y = _mm512_mul_ps(_mm512_mul_ps(p, f2), f);
138            y = _mm512_fmadd_ps(e, _mm512_set1_ps(LN2_LO), y);
139            y = _mm512_add_ps(_mm512_fnmadd_ps(_mm512_set1_ps(0.5), f2, y), f);
140            y = _mm512_fmadd_ps(e, _mm512_set1_ps(LN2_HI), y);
141            let outside = _mm512_cmp_ps_mask::<_CMP_NGT_UQ>(x, zero);
142            let special = _mm512_mask_blend_ps(
143                _mm512_cmp_ps_mask::<_CMP_EQ_OQ>(x, zero),
144                _mm512_set1_ps(f32::NAN),
145                _mm512_set1_ps(f32::NEG_INFINITY),
146            );
147            y = _mm512_mask_blend_ps(outside, y, special);
148            _mm512_mask_blend_ps(
149                _mm512_cmp_ps_mask::<_CMP_EQ_OQ>(x, _mm512_set1_ps(f32::INFINITY)),
150                y,
151                _mm512_set1_ps(f32::INFINITY),
152            )
153        }
154    }
155    unsafe {
156        let p = buf.as_mut_ptr();
157        for i in (0..buf.len()).step_by(16) {
158            _mm512_store_ps(p.add(i), ln16(_mm512_load_ps(p.add(i))));
159        }
160    }
161}