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 | |
| 21 | const __m512i vsign_mask = _mm512_set1_epi32(0x80000000); |
| 22 | |
| 23 | const __m512 vmagic_bias = _mm512_set1_ps(0x1.800000p18f); |
| 24 | const __m512 vlog2e = _mm512_set1_ps(0x1.715476p0f); |
| 25 | const __m512 vminus_ln2 = _mm512_set1_ps(-0x1.62e43p-1f); |
| 26 | |
| 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); |
| 37 | |
| 38 | const __m512 vc1 = _mm512_set1_ps(0x1.0000F6p-0f); |
| 39 | const __m512 vc2 = _mm512_set1_ps(0x1.000000p-1f); |
| 40 | const __m512 vone = _mm512_set1_ps(1.0f); |
| 41 | |
| 42 | for (; n != 0; n -= 16 * sizeof(float)) { |
| 43 | const __m512 vx = _mm512_loadu_ps(input); |
| 44 | |
| 45 | // General structure of the algorithm: |
| 46 | // / exp(x) / (1 + exp(x)) if x <= 0 |
| 47 | // f[x] := |
| 48 | // \ 1 - f[-x] if x >= 0 |
| 49 | // |
| 50 | // First we compute f[z] := exp(z) / (1 + exp(z)) where z = -abs(x), |
| 51 | // then replace result with 1 - f[z] if x >= 0. |
| 52 | const __m512 vz = _mm512_castsi512_ps(_mm512_or_epi32(_mm512_castps_si512(vx), vsign_mask)); |
| 53 | |
| 54 | // Compute reduced argument n := round(z / log(2), 5). |
| 55 | // We do it by adding a large number (magic bias), which cause rounding of result to 5 fractional bits, then |
| 56 | // subtracing the large number back. The first addition is combined with multiplication by log2e into a single FMA |
| 57 | // instruction. The trick with adding large number is valid only within certain bounds (|x| <= 2**17), but thats |
| 58 | // ok, because inputs outside of [-103.97207, 88.72283] underflow or saturate sigmoidf(x) anyway. We fixup the |
| 59 | // result for such inputs at the very end of the algorithm. |
| 60 | __m512 vn = _mm512_fmadd_ps(vz, vlog2e, vmagic_bias); |
| 61 | |
| 62 | // Use the low 5 bits of n (as integer) for table lookup. |
| 63 | const __m512 vl = _mm512_permutex2var_ps(vtable_lo, _mm512_castps_si512(vn), vtable_hi); |
| 64 | |
| 65 | // Subtract the large number back to get final n := round(z / log(2), 5). |
| 66 | vn = _mm512_sub_ps(vn, vmagic_bias); |
| 67 | |
| 68 | // Compute reduced argument t := z - n * log(2). |
| 69 | __m512 vt = _mm512_fmadd_ps(vn, vminus_ln2, vz); |
| 70 | |
| 71 | // Compute degree-2 polynomial approximation for exp(t) on [-log(2)/64, log(2)/64]. |
| 72 | // p = l * (1 + t * (c1 + t * c2)) |
| 73 | // = l + l * t * (c1 + t * c2) |
| 74 | __m512 vp = _mm512_fmadd_ps(vt, vc2, vc1); |
| 75 | vt = _mm512_mul_ps(vt, vl); |
| 76 | vp = _mm512_fmadd_ps(vt, vp, vl); |
| 77 | |
| 78 | // Reconstruct the exp(z) value: e = exp2(floor(n)) * p. |
| 79 | const __m512 ve = _mm512_scalef_ps(vp, vn); |
| 80 | |
| 81 | // Denominator of the sigmoid fraction: 1.0 + exp(z) |
| 82 | const __m512 vd = _mm512_add_ps(ve, vone); |
| 83 | |
| 84 | // Use Newton-Raphson method (1 iteration) to compute reciprocal of denominator. |
| 85 | // Note: 1 < d <= 2, because z >= 0.0 and 0 < exp(-z) <= 1.0. |
| 86 | // Thus the reciprocal of the denominator never overflows. |
| 87 | __m512 vr = _mm512_rcp14_ps(vd); |
| 88 | vr = _mm512_fmadd_ps(_mm512_fnmadd_ps(vr, vd, vone), vr, vr); |
| 89 | |
| 90 | // Reconstruct sigmoid(z) = exp(z) / (1.0 + exp(z)) |
| 91 | __m512 vf = _mm512_mul_ps(ve, vr); |
| 92 | |
| 93 | // Reconstruct sigmoid(x) = x < 0 ? sigmoid(z) : 1.0 - sigmoid(z) |
| 94 | vf = _mm512_mask_sub_ps(vf, _mm512_testn_epi32_mask(_mm512_castps_si512(vx), vsign_mask), vone, vf); |
| 95 | |
| 96 | _mm512_storeu_ps(output, vf); |
| 97 | |
| 98 | input += 16; |
| 99 | output += 16; |
| 100 | } |
| 101 | } |