Marat Dukhan | 36173d2 | 2020-10-15 17:14:26 -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/common.h> |
| 12 | #include <xnnpack/math-stubs.h> |
| 13 | |
| 14 | |
| 15 | // Table of exp2(k / 64) values, k = 0..63 |
| 16 | extern XNN_INTERNAL const float xnn_table_exp2_k_over_64[64]; |
| 17 | |
| 18 | void xnn_math_f32_sigmoid__avx512f_rr1_lut64_p2_gather_scalef_nr1fma1adj( |
| 19 | size_t n, |
| 20 | const float* input, |
| 21 | float* output) |
| 22 | { |
| 23 | assert(n % (16 * sizeof(float)) == 0); |
| 24 | |
| 25 | // Floating-point mask with only the sign bit set |
| 26 | const __m512i vsign_mask = _mm512_set1_epi32(0x80000000); |
| 27 | // Large number such that ulp(magic bias) == exp2(-6) |
| 28 | const __m512 vmagic_bias = _mm512_set1_ps(0x1.800000p17f); |
| 29 | const __m512 vlog2e = _mm512_set1_ps(0x1.715476p0f); |
| 30 | // Mask for the lowest 6 bits |
| 31 | const __m512i vindex_mask = _mm512_set1_epi32(INT32_C(0x3F)); |
| 32 | const __m512 vminus_ln2 = _mm512_set1_ps(-0x1.62e43p-1f); |
| 33 | // Coefficient of polynomial approximation of exp(t) ~ 1 + t * (1 + t * c2) on [-log(2)/128, log(2)/128] |
| 34 | const __m512 vc2 = _mm512_set1_ps(0x1.FFFF0Ap-2f); |
| 35 | const __m512 vone = _mm512_set1_ps(1.0f); |
| 36 | |
| 37 | for (; n != 0; n -= 16 * sizeof(float)) { |
| 38 | const __m512 vx = _mm512_loadu_ps(input); |
| 39 | |
| 40 | // General structure of the algorithm: |
| 41 | // |
| 42 | // / exp(x) / (1 + exp(x)) if x <= 0 |
| 43 | // f[x] := |
| 44 | // \ 1 - f[-x] if x >= 0 |
| 45 | // |
| 46 | // First we compute f[z] := exp(z) / (1 + exp(z)) where z = -abs(x), then replace result with 1 - f[z] if x >= 0. |
| 47 | const __m512 vz = _mm512_castsi512_ps(_mm512_or_epi32(_mm512_castps_si512(vx), vsign_mask)); |
| 48 | |
| 49 | // Compute reduced argument n := round(z / log(2), 6). |
| 50 | // We do it by adding a large number (magic bias), which cause rounding of the result to 6 fractional bits, then |
| 51 | // subtracing the large number back. The addition is combined with multiplication by log2e into a single FMA |
| 52 | // instruction. The trick with adding large number is valid only within certain bounds (|z / log(2)| <= 2**16, i.e. |
| 53 | // |z| <= 0x1.62E43p+15 = 45426.09375), but that is acceptable, because inputs x outside of [-87.336544, 17.328678] |
| 54 | // (i.e. z outsize [87.336544, 0]) underflow or saturate sigmoidf(x). We fixup the result for such inputs at the |
| 55 | // very end of the algorithm. |
| 56 | __m512 vn = _mm512_fmadd_ps(vz, vlog2e, vmagic_bias); |
| 57 | |
| 58 | // Use the low 6 bits of n (as integer) for table lookup. |
| 59 | const __m512i vidx = _mm512_and_epi32(_mm512_castps_si512(vn), vindex_mask); |
| 60 | const __m512 vl = _mm512_i32gather_ps(vidx, xnn_table_exp2_k_over_64, sizeof(float)); |
| 61 | |
| 62 | // Subtract the large number back to get the final n := round(z / log(2), 6) as a floating-point number. |
| 63 | vn = _mm512_sub_ps(vn, vmagic_bias); |
| 64 | |
| 65 | // Compute reduced argument t := z - n * log(2). |
| 66 | const __m512 vt = _mm512_fmadd_ps(vn, vminus_ln2, vz); |
| 67 | |
| 68 | // Compute degree-2 polynomial approximation for exp(t) on [-log(2)/128, log(2)/128]. |
| 69 | // P(t) = 1 + t * (1 + t * c2) = 1 + (t + t * (t * c2)) |
| 70 | // p = l * P(t) |
| 71 | // = l + l * (t + t * (t * c2)) |
| 72 | __m512 vp = _mm512_mul_ps(vt, vc2); |
| 73 | vp = _mm512_fmadd_ps(vt, vp, vt); |
| 74 | vp = _mm512_fmadd_ps(vl, vp, vl); |
| 75 | |
| 76 | // Reconstruct the exp(z) value: e = exp2(floor(n)) * p. |
| 77 | const __m512 ve = _mm512_scalef_ps(vp, vn); |
| 78 | |
| 79 | // Denominator of the sigmoid fraction: 1.0 + exp(z) |
| 80 | const __m512 vd = _mm512_add_ps(ve, vone); |
| 81 | |
| 82 | // Use Newton-Raphson method (1 iteration) to compute reciprocal of denominator. |
| 83 | // Note: 1 < d <= 2, because z >= 0.0 and 0 < exp(-z) <= 1.0. |
| 84 | // Thus the reciprocal of the denominator never overflows. |
| 85 | __m512 vr = _mm512_rcp14_ps(vd); |
| 86 | vr = _mm512_fmadd_ps(_mm512_fnmadd_ps(vr, vd, vone), vr, vr); |
| 87 | |
| 88 | // Reconstruct sigmoid(z) = exp(z) / (1.0 + exp(z)) with adjustment to match IEEE division result |
| 89 | __m512 vf = _mm512_mul_ps(ve, vr); |
| 90 | vf = _mm512_fmadd_ps(_mm512_fnmadd_ps(vf, vd, ve), vr, vf); |
| 91 | |
| 92 | // Reconstruct sigmoid(x) = x < 0 ? sigmoid(z) : 1.0 - sigmoid(z) |
| 93 | vf = _mm512_mask_sub_ps(vf, _mm512_testn_epi32_mask(_mm512_castps_si512(vx), vsign_mask), vone, vf); |
| 94 | |
| 95 | _mm512_storeu_ps(output, vf); |
| 96 | |
| 97 | input += 16; |
| 98 | output += 16; |
| 99 | } |
| 100 | } |