Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 1 | // Copyright 2020 Google LLC |
| 2 | // |
| 3 | // This source code is licensed under the BSD-style license found in the |
| 4 | // LICENSE file in the root directory of this source tree. |
| 5 | |
| 6 | #include <assert.h> |
| 7 | #include <stddef.h> |
| 8 | |
| 9 | #include <immintrin.h> |
| 10 | |
| 11 | #include <xnnpack/math-stubs.h> |
| 12 | |
| 13 | |
| 14 | void xnn_math_f32_sigmoid__avx512f_rr1_lut32_p2_perm2_scalef_nr1fma( |
| 15 | size_t n, |
| 16 | const float* input, |
| 17 | float* output) |
| 18 | { |
| 19 | assert(n % (16 * sizeof(float)) == 0); |
| 20 | |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 21 | // Floating-point mask with only the sign bit set |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 22 | const __m512i vsign_mask = _mm512_set1_epi32(0x80000000); |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 23 | // Large number such that ulp(magic bias) == exp2(-5) |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 24 | const __m512 vmagic_bias = _mm512_set1_ps(0x1.800000p18f); |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 25 | const __m512 vlog2e = _mm512_set1_ps(0x1.715476p0f); |
| 26 | // Table of exp2(k / 32) values, k = 0..31 |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 27 | const __m512 vtable_hi = _mm512_set_ps( |
| 28 | 0x1.F50766p+0f, 0x1.EA4AFAp+0f, 0x1.DFC974p+0f, 0x1.D5818Ep+0f, |
| 29 | 0x1.CB720Ep+0f, 0x1.C199BEp+0f, 0x1.B7F770p+0f, 0x1.AE89FAp+0f, |
| 30 | 0x1.A5503Cp+0f, 0x1.9C4918p+0f, 0x1.93737Cp+0f, 0x1.8ACE54p+0f, |
| 31 | 0x1.82589Ap+0f, 0x1.7A1148p+0f, 0x1.71F75Ep+0f, 0x1.6A09E6p+0f); |
| 32 | const __m512 vtable_lo = _mm512_set_ps( |
| 33 | 0x1.6247ECp+0f, 0x1.5AB07Ep+0f, 0x1.5342B6p+0f, 0x1.4BFDAEp+0f, |
| 34 | 0x1.44E086p+0f, 0x1.3DEA64p+0f, 0x1.371A74p+0f, 0x1.306FE0p+0f, |
| 35 | 0x1.29E9E0p+0f, 0x1.2387A6p+0f, 0x1.1D4874p+0f, 0x1.172B84p+0f, |
| 36 | 0x1.11301Ep+0f, 0x1.0B5586p+0f, 0x1.059B0Ep+0f, 0x1.000000p+0f); |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 37 | const __m512 vminus_ln2 = _mm512_set1_ps(-0x1.62E43p-1f); |
| 38 | // Coefficient of polynomial approximation of |
| 39 | // exp(t) ~ 1 + t * (c1 + t * c2) on [-log(2)/64, log(2)/64] |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 40 | const __m512 vc2 = _mm512_set1_ps(0x1.000000p-1f); |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 41 | const __m512 vc1 = _mm512_set1_ps(0x1.0000F6p-0f); |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 42 | const __m512 vone = _mm512_set1_ps(1.0f); |
| 43 | |
| 44 | for (; n != 0; n -= 16 * sizeof(float)) { |
| 45 | const __m512 vx = _mm512_loadu_ps(input); |
| 46 | |
| 47 | // General structure of the algorithm: |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 48 | // |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 49 | // / exp(x) / (1 + exp(x)) if x <= 0 |
| 50 | // f[x] := |
| 51 | // \ 1 - f[-x] if x >= 0 |
| 52 | // |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 53 | // First we compute f[z] := exp(z) / (1 + exp(z)) where z = -abs(x), then replace result with 1 - f[z] if x >= 0. |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 54 | const __m512 vz = _mm512_castsi512_ps(_mm512_or_epi32(_mm512_castps_si512(vx), vsign_mask)); |
| 55 | |
| 56 | // Compute reduced argument n := round(z / log(2), 5). |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 57 | // We do it by adding a large number (magic bias), which cause rounding of the result to 5 fractional bits, then |
| 58 | // subtracing the large number back. The addition is combined with multiplication by log2e into a single FMA |
| 59 | // instruction. The trick with adding large number is valid only within certain bounds (|z / log(2)| <= 2**17, |
| 60 | // i.e. |z| <= 0x1.62E43p+17 = 181704.375), but that is acceptable, because inputs x outside of |
| 61 | // [-87.336544, 17.328678] (i.e. z outsize [87.336544, 0]) underflow or saturate sigmoidf(x). We fixup the result |
| 62 | // for such inputs at the very end of the algorithm. |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 63 | __m512 vn = _mm512_fmadd_ps(vz, vlog2e, vmagic_bias); |
| 64 | |
| 65 | // Use the low 5 bits of n (as integer) for table lookup. |
| 66 | const __m512 vl = _mm512_permutex2var_ps(vtable_lo, _mm512_castps_si512(vn), vtable_hi); |
| 67 | |
| 68 | // Subtract the large number back to get final n := round(z / log(2), 5). |
| 69 | vn = _mm512_sub_ps(vn, vmagic_bias); |
| 70 | |
| 71 | // Compute reduced argument t := z - n * log(2). |
| 72 | __m512 vt = _mm512_fmadd_ps(vn, vminus_ln2, vz); |
| 73 | |
| 74 | // Compute degree-2 polynomial approximation for exp(t) on [-log(2)/64, log(2)/64]. |
Marat Dukhan | 9501298 | 2020-09-28 01:17:52 -0700 | [diff] [blame^] | 75 | // P(t) = 1 + t * (c1 + t * c2) |
| 76 | // p = l * P(t) |
Marat Dukhan | 6a34c5f | 2020-09-22 21:44:15 -0700 | [diff] [blame] | 77 | // = l + l * t * (c1 + t * c2) |
| 78 | __m512 vp = _mm512_fmadd_ps(vt, vc2, vc1); |
| 79 | vt = _mm512_mul_ps(vt, vl); |
| 80 | vp = _mm512_fmadd_ps(vt, vp, vl); |
| 81 | |
| 82 | // Reconstruct the exp(z) value: e = exp2(floor(n)) * p. |
| 83 | const __m512 ve = _mm512_scalef_ps(vp, vn); |
| 84 | |
| 85 | // Denominator of the sigmoid fraction: 1.0 + exp(z) |
| 86 | const __m512 vd = _mm512_add_ps(ve, vone); |
| 87 | |
| 88 | // Use Newton-Raphson method (1 iteration) to compute reciprocal of denominator. |
| 89 | // Note: 1 < d <= 2, because z >= 0.0 and 0 < exp(-z) <= 1.0. |
| 90 | // Thus the reciprocal of the denominator never overflows. |
| 91 | __m512 vr = _mm512_rcp14_ps(vd); |
| 92 | vr = _mm512_fmadd_ps(_mm512_fnmadd_ps(vr, vd, vone), vr, vr); |
| 93 | |
| 94 | // Reconstruct sigmoid(z) = exp(z) / (1.0 + exp(z)) |
| 95 | __m512 vf = _mm512_mul_ps(ve, vr); |
| 96 | |
| 97 | // Reconstruct sigmoid(x) = x < 0 ? sigmoid(z) : 1.0 - sigmoid(z) |
| 98 | vf = _mm512_mask_sub_ps(vf, _mm512_testn_epi32_mask(_mm512_castps_si512(vx), vsign_mask), vone, vf); |
| 99 | |
| 100 | _mm512_storeu_ps(output, vf); |
| 101 | |
| 102 | input += 16; |
| 103 | output += 16; |
| 104 | } |
| 105 | } |