1// Copyright 2019 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$assert ELEMENTS_TILE % 8 == 0 7$assert ELEMENTS_TILE >= 8 8$SIMD_TILE = ELEMENTS_TILE // 8 9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 10#include <assert.h> 11 12#include <immintrin.h> 13 14#include <xnnpack/raddexpminusmax.h> 15 16 17static const int32_t mask_table[14] = {-1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0}; 18 19void xnn_f32_raddexpminusmax_ukernel__avx2_p5_x${ELEMENTS_TILE}${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}( 20 size_t elements, 21 const float* input, 22 float* sum, 23 float max) 24{ 25 assert(elements % sizeof(float) == 0); 26 27 const __m256 vmagic_bias = _mm256_set1_ps(0x1.8000FEp23f); 28 // The smallest x for which expf(x) is normalized. 29 const __m256 vdenorm_cutoff = _mm256_set1_ps(-0x1.5D589Ep6f); 30 const __m256 vlog2e = _mm256_set1_ps(0x1.715476p+0f); 31 const __m256 vminus_ln2_hi = _mm256_set1_ps(-0x1.62E43p-1f); 32 const __m256 vminus_ln2_lo = _mm256_set1_ps(0x1.05C61p-29f); 33 34 const __m256 vc1 = _mm256_set1_ps(0x1.FFFFF6p-1f); 35 const __m256 vc2 = _mm256_set1_ps(0x1.FFFDC6p-2f); 36 const __m256 vc3 = _mm256_set1_ps(0x1.555A80p-3f); 37 const __m256 vc4 = _mm256_set1_ps(0x1.573A1Ap-5f); 38 const __m256 vc5 = _mm256_set1_ps(0x1.0F9F9Cp-7f); 39 40 const __m256 vi_max = _mm256_set1_ps(max); 41 42 $for K in range(ACCUMULATORS): 43 __m256 vacc${K} = _mm256_setzero_ps(); 44 for (; elements >= ${ELEMENTS_TILE} * sizeof(float); elements -= ${ELEMENTS_TILE} * sizeof(float)) { 45 // Load ${ELEMENTS_TILE} (${SIMD_TILE}x8) inputs at a time. 46 const __m256 vi0 = _mm256_loadu_ps(input); 47 $for N in range(1, SIMD_TILE): 48 const __m256 vi${N} = _mm256_loadu_ps(input + ${N * 8}); 49 input += ${ELEMENTS_TILE}; 50 51 // Subtract maximum input x := i - i_max. This implies x <= 0. 52 $for N in range(SIMD_TILE): 53 const __m256 vx${N} = _mm256_sub_ps(vi${N}, vi_max); 54 55 // Compute reduced argument elements := round(x / log(2)). 56 $for N in range(SIMD_TILE): 57 __m256 vn${N} = _mm256_fmadd_ps(vx${N}, vlog2e, vmagic_bias); 58 59 // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e. 60 // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly. 61 $for N in range(SIMD_TILE): 62 const __m256 vs${N} = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(vn${N}), 23)); 63 64 // Subtract the large number back to get final elements := round(x / log(2)). 65 $for N in range(SIMD_TILE): 66 vn${N} = _mm256_sub_ps(vn${N}, vmagic_bias); 67 68 // Compute reduced argument t := x - elements * log(2). 69 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 70 $for N in range(SIMD_TILE): 71 __m256 vt${N} = _mm256_fmadd_ps(vn${N}, vminus_ln2_hi, vx${N}); 72 73 $for N in range(SIMD_TILE): 74 vt${N} = _mm256_fmadd_ps(vn${N}, vminus_ln2_lo, vt${N}); 75 76 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 77 $for N in range(SIMD_TILE): 78 __m256 vp${N} = _mm256_fmadd_ps(vc5, vt${N}, vc4); 79 80 $for N in range(SIMD_TILE): 81 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc3); 82 83 $for N in range(SIMD_TILE): 84 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc2); 85 86 $for N in range(SIMD_TILE): 87 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc1); 88 89 // Reconstruct the final f value: 90 // f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))) 91 // = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))) 92 // = s + (t * s) * p 93 $for N in range(SIMD_TILE): 94 vt${N} = _mm256_mul_ps(vt${N}, vs${N}); 95 96 $for N in range(SIMD_TILE): 97 __m256 vf${N} = _mm256_fmadd_ps(vt${N}, vp${N}, vs${N}); 98 99 // For inputs below zero cutoff, replace output with +0.0f. 100 // Note that for NaN inputs, comparison result is false, and outputs are left unchanged. 101 $for N in range(SIMD_TILE): 102 vf${N} = _mm256_andnot_ps(_mm256_cmp_ps(vx${N}, vdenorm_cutoff, _CMP_LT_OS), vf${N}); 103 104 // Accumulate computed exponents. 105 $for N in range(SIMD_TILE): 106 vacc${N % ACCUMULATORS} = _mm256_add_ps(vacc${N % ACCUMULATORS}, vf${N}); 107 } 108 $if ACCUMULATORS > 1: 109 // Add up all accumulators to vacc0 110 $ACC_SLICE = 1 111 $while ACC_SLICE < ACCUMULATORS: 112 $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): 113 $if A + ACC_SLICE < ACCUMULATORS: 114 vacc${A} = _mm256_add_ps(vacc${A}, vacc${A + ACC_SLICE}); 115 $ACC_SLICE *= 2 116 117 __m256 vacc = vacc0; 118 for (; elements >= 8 * sizeof(float); elements -= 8 * sizeof(float)) { 119 // Load 8 inputs at a time. 120 const __m256 vi = _mm256_loadu_ps(input); 121 input += 8; 122 123 // Subtract maximum input x := i - i_max. This implies x <= 0. 124 const __m256 vx = _mm256_sub_ps(vi, vi_max); 125 126 // Compute reduced argument elements := round(x / log(2)). 127 __m256 vn = _mm256_fmadd_ps(vx, vlog2e, vmagic_bias); 128 129 // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e. 130 // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly. 131 const __m256 vs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(vn), 23)); 132 133 // Subtract the large number back to get final elements := round(x / log(2)). 134 vn = _mm256_sub_ps(vn, vmagic_bias); 135 136 // Compute reduced argument t := x - elements * log(2). 137 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 138 __m256 vt = _mm256_fmadd_ps(vn, vminus_ln2_hi, vx); 139 vt = _mm256_fmadd_ps(vn, vminus_ln2_lo, vt); 140 141 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 142 __m256 vp = _mm256_fmadd_ps(vc5, vt, vc4); 143 vp = _mm256_fmadd_ps(vp, vt, vc3); 144 vp = _mm256_fmadd_ps(vp, vt, vc2); 145 vp = _mm256_fmadd_ps(vp, vt, vc1); 146 147 // Reconstruct the final f value: 148 // f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))) 149 // = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))) 150 // = s + (t * s) * p 151 vt = _mm256_mul_ps(vt, vs); 152 __m256 vf = _mm256_fmadd_ps(vt, vp, vs); 153 154 // For inputs below zero cutoff, replace output with +0.0f. 155 // Note that for NaN inputs, comparison result is false, and outputs are left unchanged. 156 vf = _mm256_andnot_ps(_mm256_cmp_ps(vx, vdenorm_cutoff, _CMP_LT_OS), vf); 157 158 // Accumulate computed exponents. 159 vacc = _mm256_add_ps(vacc, vf); 160 } 161 if (elements != 0) { 162 assert(elements >= 1 * sizeof(float)); 163 assert(elements <= 7 * sizeof(float)); 164 const __m256i vmask = _mm256_loadu_si256((const __m256i*) ((uintptr_t) &mask_table[7] - elements)); 165 166 // Load up to 7 inputs at a time. 167 const __m256 vi = _mm256_maskload_ps(input, vmask); 168 169 // Subtract maximum input x := i - i_max. This implies x <= 0. 170 const __m256 vx = _mm256_sub_ps(vi, vi_max); 171 172 // Compute reduced argument elements := round(x / log(2)). 173 __m256 vn = _mm256_fmadd_ps(vx, vlog2e, vmagic_bias); 174 175 // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e. 176 // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly. 177 const __m256 vs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(vn), 23)); 178 179 // Subtract the large number back to get final elements := round(x / log(2)). 180 vn = _mm256_sub_ps(vn, vmagic_bias); 181 182 // Compute reduced argument t := x - elements * log(2). 183 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 184 __m256 vt = _mm256_fmadd_ps(vn, vminus_ln2_hi, vx); 185 vt = _mm256_fmadd_ps(vn, vminus_ln2_lo, vt); 186 187 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 188 __m256 vp = _mm256_fmadd_ps(vc5, vt, vc4); 189 vp = _mm256_fmadd_ps(vp, vt, vc3); 190 vp = _mm256_fmadd_ps(vp, vt, vc2); 191 vp = _mm256_fmadd_ps(vp, vt, vc1); 192 193 // Reconstruct the final f value: 194 // f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))) 195 // = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))) 196 // = s + (t * s) * p 197 vt = _mm256_mul_ps(vt, vs); 198 __m256 vf = _mm256_fmadd_ps(vt, vp, vs); 199 200 // For inputs below zero cutoff, replace output with +0.0f. 201 // Note that for NaN inputs, comparison result is false, and outputs are left unchanged. 202 vf = _mm256_andnot_ps(_mm256_cmp_ps(vx, vdenorm_cutoff, _CMP_LT_OS), vf); 203 204 // Accumulate computed exponents. And addend with mask to leave unmasked 32-bit lanes unchanged. 205 vacc = _mm256_add_ps(vacc, _mm256_and_ps(vf, _mm256_castsi256_ps(vmask))); 206 } 207 // Reduce 8 elements in the SIMD register 208 __m128 vacc_lo = _mm_add_ps(_mm256_castps256_ps128(vacc), _mm256_extractf128_ps(vacc, 1)); 209 vacc_lo = _mm_add_ps(vacc_lo, _mm_movehl_ps(vacc_lo, vacc_lo)); 210 vacc_lo = _mm_add_ss(vacc_lo, _mm_movehdup_ps(vacc_lo)); 211 _mm_store_ss(sum, vacc_lo); 212 _mm256_zeroupper(); 213} 214