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#include <math.h> 12 13#include <immintrin.h> 14 15#include <xnnpack/raddextexp.h> 16 17 18static const int32_t mask_table[14] = {-1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0}; 19 20void xnn_f32_raddextexp_ukernel__avx2_p5_x${ELEMENTS_TILE}${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}( 21 size_t elements, 22 const float* x, 23 float* sum) 24{ 25 assert(elements % sizeof(float) == 0); 26 27 const __m256 vlog2e = _mm256_set1_ps(0x1.715476p+0f); 28 const __m256 vminus_ln2_hi = _mm256_set1_ps(-0x1.62E43p-1f); 29 const __m256 vminus_ln2_lo = _mm256_set1_ps(0x1.05C61p-29f); 30 31 // The smallest elements such that 2**elements is considered non-negligible. 32 // For smaller elements, 2**elements is replaced with zero. 33 const __m256 vmin_exponent = _mm256_set1_ps(-127.0f); 34 const __m256 vmagic_bias = _mm256_set1_ps(0x1.8000FEp23f); 35 const __m256 vminus_inf = _mm256_set1_ps(-INFINITY); 36 37 const __m256 vc0 = _mm256_set1_ps(1.0f); 38 const __m256 vc1 = _mm256_set1_ps(0x1.FFFFF6p-1f); 39 const __m256 vc2 = _mm256_set1_ps(0x1.FFFDC6p-2f); 40 const __m256 vc3 = _mm256_set1_ps(0x1.555A80p-3f); 41 const __m256 vc4 = _mm256_set1_ps(0x1.573A1Ap-5f); 42 const __m256 vc5 = _mm256_set1_ps(0x1.0F9F9Cp-7f); 43 44 $for K in range(ACCUMULATORS): 45 __m256 vaccv${K} = _mm256_setzero_ps(); 46 $for K in range(ACCUMULATORS): 47 __m256 vacce${K} = vminus_inf; 48 for (; elements >= ${ELEMENTS_TILE} * sizeof(float); elements -= ${ELEMENTS_TILE} * sizeof(float)) { 49 // Load ${ELEMENTS_TILE} (${SIMD_TILE}x8) inputs at a time. 50 const __m256 vx0 = _mm256_loadu_ps(x); 51 $for N in range(1, SIMD_TILE): 52 const __m256 vx${N} = _mm256_loadu_ps(x + ${N * 8}); 53 x += ${ELEMENTS_TILE}; 54 55 // Compute reduced argument elements := round(x / log(2)). 56 $for N in range(SIMD_TILE): 57 const __m256 vn${N} = _mm256_round_ps(_mm256_mul_ps(vx${N}, vlog2e), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); 58 59 // Compute reduced argument t := x - elements * log(2). 60 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 61 $for N in range(SIMD_TILE): 62 __m256 vt${N} = _mm256_fmadd_ps(vn${N}, vminus_ln2_hi, vx${N}); 63 64 $for N in range(SIMD_TILE): 65 vt${N} = _mm256_fmadd_ps(vn${N}, vminus_ln2_lo, vt${N}); 66 67 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 68 $for N in range(SIMD_TILE): 69 __m256 vp${N} = _mm256_fmadd_ps(vc5, vt${N}, vc4); 70 71 $for N in range(SIMD_TILE): 72 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc3); 73 74 $for N in range(SIMD_TILE): 75 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc2); 76 77 $for N in range(SIMD_TILE): 78 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc1); 79 80 $for N in range(SIMD_TILE): 81 vp${N} = _mm256_fmadd_ps(vp${N}, vt${N}, vc0); 82 83 // Accumulate "extended" floating-point numbers in ("mantissa", "exponent") representation where 84 // - vnX is "exponent" 85 // - vpX is "mantissa" 86 // 87 // exp2(ae) * av + exp2(be) * bv = 88 // = exp2(max(ae, be)) * exp2(ae - max(ae, be)) * av + exp2(max(ae, be)) * exp2(be - max(ae, be)) * bv 89 // = exp2(max_e) * (exp2(ae - max_e) * av + exp2(be - max_e) * bv) 90 // = exp2(max_e) * (exp2(delta_ae) * av + exp2(delta_be) * bv) 91 // 92 // For computational efficiency we may add several "extended" floating-point numbers at a time. 93 $for N in range(SIMD_TILE): 94 $if N < ACCUMULATORS: 95 __m256 vmax_e${N} = _mm256_max_ps(vacce${N}, vn${N}); 96 $else: 97 vmax_e${N % ACCUMULATORS} = _mm256_max_ps(vmax_e${N % ACCUMULATORS}, vn${N}); 98 99 // For computational efficiency, replace exp2(delta_e) with 0.0f when delta_e <= -127.0. 100 // This replacement is done in two steps: 101 // 1. Clamp minimum delta_e at -127.0. 102 // 2. Map delta_e to scale factor 0.0 when delta_e == -127.0 103 $for K in range(ACCUMULATORS): 104 const __m256 vdelta_acce${K} = _mm256_max_ps(_mm256_sub_ps(vacce${K}, vmax_e${K}), vmin_exponent); 105 $for N in range(SIMD_TILE): 106 const __m256 vdelta_e${N} = _mm256_max_ps(_mm256_sub_ps(vn${N}, vmax_e${N % ACCUMULATORS}), vmin_exponent); 107 108 // Convert delta-exponents into scale factors: 109 // - s = exp2(delta_e) when delta_e > -127.0 110 // - s = 0.0 when delta_e <= -127.0 111 // 112 // Note: delta-exponents can not exceed 0.0, thus scale factors can not exceed 1.0. 113 $for K in range(ACCUMULATORS): 114 const __m256 vaccs${K} = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_acce${K}, vmagic_bias)), 23)); 115 $for N in range(SIMD_TILE): 116 const __m256 vs${N} = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_e${N}, vmagic_bias)), 23)); 117 118 // Update accumulated "mantissa" and "exponent" values 119 $for K in range(ACCUMULATORS): 120 vaccv${K} = _mm256_mul_ps(vaccv${K}, vaccs${K}); 121 $for N in range(SIMD_TILE): 122 vaccv${N % ACCUMULATORS} = _mm256_fmadd_ps(vp${N}, vs${N}, vaccv${N % ACCUMULATORS}); 123 124 $for K in range(ACCUMULATORS): 125 vacce${K} = vmax_e${K}; 126 } 127 128 // Reduce partial sums of "extended" floating-point numbers into a single "extended" SIMD vector of sums. 129 $if ACCUMULATORS > 1: 130 $for A in range(0, ACCUMULATORS, 2): 131 $if A + 1 < ACCUMULATORS: 132 const __m256 vmax_acce${ABC[A:A+2]} = _mm256_max_ps(vacce${A}, vacce${A+1}); 133 $else: 134 const __m256 vmax_acce${ABC[A]} = vacce${A}; 135 $ACC_SLICE = 2 136 $while ACC_SLICE < ACCUMULATORS: 137 $for A in range(0, ACCUMULATORS, ACC_SLICE * 2): 138 $if A + ACC_SLICE < ACCUMULATORS: 139 const __m256 vmax_acce${ABC[A:min(A+ACC_SLICE*2, ACCUMULATORS)]} = _mm256_max_ps(vmax_acce${ABC[A:A+ACC_SLICE]}, vmax_acce${ABC[A+ACC_SLICE:min(ACCUMULATORS,A+ACC_SLICE*2)]}); 140 $ACC_SLICE *= 2 141 142 $for K in range(ACCUMULATORS): 143 const __m256 vdelta_acce${K} = _mm256_max_ps(_mm256_sub_ps(vacce${K}, vmax_acce${ABC[0:ACCUMULATORS]}), vmin_exponent); 144 145 $for K in range(ACCUMULATORS): 146 const __m256 vaccs${K} = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_acce${K}, vmagic_bias)), 23)); 147 148 __m256 vaccv = _mm256_mul_ps(vaccv0, vaccs0); 149 $for K in range(1, ACCUMULATORS): 150 vaccv = _mm256_fmadd_ps(vaccv${K}, vaccs${K}, vaccv); 151 __m256 vacce = vmax_acce${ABC[0:ACCUMULATORS]}; 152 $else: 153 __m256 vaccv = vaccv0; 154 __m256 vacce = vacce0; 155 156 for (; elements >= 8 * sizeof(float); elements -= 8 * sizeof(float)) { 157 // Load 8 inputs at a time. 158 const __m256 vx = _mm256_loadu_ps(x); 159 x += 8; 160 161 // Compute reduced argument elements := round(x / log(2)). 162 const __m256 vn = _mm256_round_ps(_mm256_mul_ps(vx, vlog2e), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); 163 164 // Compute reduced argument t := x - elements * log(2). 165 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 166 __m256 vt = _mm256_fmadd_ps(vn, vminus_ln2_hi, vx); 167 vt = _mm256_fmadd_ps(vn, vminus_ln2_lo, vt); 168 169 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 170 __m256 vp = _mm256_fmadd_ps(vc5, vt, vc4); 171 vp = _mm256_fmadd_ps(vp, vt, vc3); 172 vp = _mm256_fmadd_ps(vp, vt, vc2); 173 vp = _mm256_fmadd_ps(vp, vt, vc1); 174 vp = _mm256_fmadd_ps(vp, vt, vc0); 175 176 // Accumulate "extended" floating-point numbers in ("mantissa", "exponent") representation. 177 const __m256 vmax_e = _mm256_max_ps(vacce, vn); 178 179 // For computational efficiency, clamp minimum exp2(delta_e) at -127.0. It will be mapped to 0.0 scale factor later. 180 const __m256 vdelta_acce = _mm256_max_ps(_mm256_sub_ps(vacce, vmax_e), vmin_exponent); 181 const __m256 vdelta_e = _mm256_max_ps(_mm256_sub_ps(vn, vmax_e), vmin_exponent); 182 183 // Convert exponents into scale factors. 184 const __m256 vaccs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_acce, vmagic_bias)), 23)); 185 const __m256 vs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_e, vmagic_bias)), 23)); 186 187 // Update accumulated "mantissa" and "exponent" values. 188 vaccv = _mm256_mul_ps(vaccv, vaccs); 189 vaccv = _mm256_fmadd_ps(vp, vs, vaccv); 190 191 vacce = vmax_e; 192 } 193 if XNN_UNLIKELY(elements != 0) { 194 assert(elements >= 1 * sizeof(float)); 195 assert(elements <= 7 * sizeof(float)); 196 const __m256i vmask = _mm256_loadu_si256((const __m256i*) ((uintptr_t) &mask_table[7] - elements)); 197 198 // Load up to 7 inputs at a time. 199 const __m256 vx = _mm256_maskload_ps(x, vmask); 200 201 // Compute reduced argument elements := round(x / log(2)). 202 __m256 vn = _mm256_round_ps(_mm256_mul_ps(vx, vlog2e), _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC); 203 204 // Compute reduced argument t := x - elements * log(2). 205 // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy. 206 __m256 vt = _mm256_fmadd_ps(vn, vminus_ln2_hi, vx); 207 vt = _mm256_fmadd_ps(vn, vminus_ln2_lo, vt); 208 209 // Correct reduced argument elements for masked out elements. 210 vn = _mm256_blendv_ps(vacce, vn, _mm256_castsi256_ps(vmask)); 211 212 // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2]. 213 __m256 vp = _mm256_fmadd_ps(vc5, vt, vc4); 214 vp = _mm256_fmadd_ps(vp, vt, vc3); 215 vp = _mm256_fmadd_ps(vp, vt, vc2); 216 vp = _mm256_fmadd_ps(vp, vt, vc1); 217 vp = _mm256_fmadd_ps(vp, vt, vc0); 218 vp = _mm256_and_ps(vp, _mm256_castsi256_ps(vmask)); 219 220 // Accumulate "extended" floating-point numbers in ("mantissa", "exponent") representation. 221 const __m256 vmax_e = _mm256_max_ps(vacce, vn); 222 223 // For computational efficiency, clamp minimum exp2(delta_e) at -127.0. It will be mapped to 0.0 scale factor later. 224 const __m256 vdelta_e = _mm256_max_ps(_mm256_sub_ps(vn, vmax_e), vmin_exponent); 225 const __m256 vdelta_acce = _mm256_max_ps(_mm256_sub_ps(vacce, vmax_e), vmin_exponent); 226 227 // Convert exponents into scale factors. 228 const __m256 vs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_e, vmagic_bias)), 23)); 229 const __m256 vaccs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_acce, vmagic_bias)), 23)); 230 231 // Update accumulated "mantissa" and "exponent" values. 232 vaccv = _mm256_mul_ps(vaccv, vaccs); 233 vaccv = _mm256_fmadd_ps(vp, vs, vaccv); 234 235 vacce = vmax_e; 236 } 237 238 // Reduce partial sums of "extended" floating-point numbers into a single "extended" floating-point sum. 239 __m256 vmax_acce = _mm256_max_ps(vacce, _mm256_permute2f128_ps(vacce, vacce, 1)); 240 vmax_acce = _mm256_max_ps(vmax_acce, _mm256_shuffle_ps(vmax_acce, vmax_acce, _MM_SHUFFLE(1, 0, 3, 2))); 241 vmax_acce = _mm256_max_ps(vmax_acce, _mm256_shuffle_ps(vmax_acce, vmax_acce, _MM_SHUFFLE(2, 3, 0, 1))); 242 const __m256 vdelta_acce = _mm256_max_ps(_mm256_sub_ps(vacce, vmax_acce), vmin_exponent); 243 const __m256 vaccs = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_castps_si256(_mm256_add_ps(vdelta_acce, vmagic_bias)), 23)); 244 245 vaccv = _mm256_mul_ps(vaccv, vaccs); 246 __m128 vaccv_sum = _mm_add_ps(_mm256_castps256_ps128(vaccv), _mm256_extractf128_ps(vaccv, 1)); 247 vaccv_sum = _mm_add_ps(vaccv_sum, _mm_movehl_ps(vaccv_sum, vaccv_sum)); 248 vaccv_sum = _mm_add_ss(vaccv_sum, _mm_movehdup_ps(vaccv_sum)); 249 250 _mm_store_ss(&sum[0], vaccv_sum); 251 _mm_store_ss(&sum[1], _mm256_castps256_ps128(vmax_acce)); 252 253 _mm256_zeroupper(); 254} 255