xref: /aosp_15_r20/external/XNNPACK/src/f32-raddstoreexpminusmax/sse2-rr2-p5.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
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 % 4 == 0
7$assert ELEMENTS_TILE >= 4
8$SIMD_TILE = ELEMENTS_TILE // 4
9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
10#include <assert.h>
11
12#include <emmintrin.h>
13
14#include <xnnpack/common.h>
15#include <xnnpack/raddstoreexpminusmax.h>
16
17
18void xnn_f32_raddstoreexpminusmax_ukernel__sse2_rr2_p5_x${ELEMENTS_TILE}${"" if ACCUMULATORS == 1 else "_acc%d" % ACCUMULATORS}(
19    size_t elements,
20    const float* input,
21    const float* max,
22    float* output,
23    float* sum,
24    const union xnn_f32_expminus_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
25{
26  assert(elements % sizeof(float) == 0);
27
28  const __m128 vi_max = _mm_load1_ps(max);
29  const __m128 vlog2e = _mm_load_ps(params->sse2_rr2_p5.log2e);
30  const __m128 vmagic_bias = _mm_load_ps(params->sse2_rr2_p5.magic_bias);
31  const __m128 vminus_ln2_hi = _mm_load_ps(params->sse2_rr2_p5.minus_ln2_hi);
32  const __m128 vminus_ln2_lo = _mm_load_ps(params->sse2_rr2_p5.minus_ln2_lo);
33  const __m128 vc5 = _mm_load_ps(params->sse2_rr2_p5.c5);
34  const __m128 vc4 = _mm_load_ps(params->sse2_rr2_p5.c4);
35  const __m128 vc3 = _mm_load_ps(params->sse2_rr2_p5.c3);
36  const __m128 vc2 = _mm_load_ps(params->sse2_rr2_p5.c2);
37  const __m128 vc1 = _mm_load_ps(params->sse2_rr2_p5.c1);
38  const __m128 vdenorm_cutoff = _mm_load_ps(params->sse2_rr2_p5.denorm_cutoff);
39
40  $for K in range(ACCUMULATORS):
41    __m128 vacc${K} = _mm_setzero_ps();
42  for (; elements >= ${ELEMENTS_TILE} * sizeof(float); elements -= ${ELEMENTS_TILE} * sizeof(float)) {
43    // Load ${ELEMENTS_TILE} (${SIMD_TILE}x4) inputs at a time.
44    const __m128 vi${ABC[0:4]} = _mm_loadu_ps(input);
45    $for N in range(4, ELEMENTS_TILE, 4):
46      const __m128 vi${ABC[N:N+4]} = _mm_loadu_ps(input + ${N});
47    input += ${ELEMENTS_TILE};
48
49    // Subtract maximum input x := i - i_max. This implies x <= 0.
50    $for N in range(0, ELEMENTS_TILE, 4):
51      const __m128 vx${ABC[N:N+4]} = _mm_sub_ps(vi${ABC[N:N+4]}, vi_max);
52
53    // Compute reduced argument elements := round(x / log(2)).
54    $for N in range(0, ELEMENTS_TILE, 4):
55      __m128 vn${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vx${ABC[N:N+4]}, vlog2e), vmagic_bias);
56
57    // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e.
58    // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly.
59    $for N in range(0, ELEMENTS_TILE, 4):
60      const __m128 vs${ABC[N:N+4]} = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(vn${ABC[N:N+4]}), 23));
61
62    // Subtract the large number back to get final elements := round(x / log(2)).
63    $for N in range(0, ELEMENTS_TILE, 4):
64      vn${ABC[N:N+4]} = _mm_sub_ps(vn${ABC[N:N+4]}, vmagic_bias);
65
66    // Compute reduced argument t := x - elements * log(2).
67    // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy.
68    $for N in range(0, ELEMENTS_TILE, 4):
69      __m128 vt${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vn${ABC[N:N+4]}, vminus_ln2_hi), vx${ABC[N:N+4]});
70
71    $for N in range(0, ELEMENTS_TILE, 4):
72      vt${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vn${ABC[N:N+4]}, vminus_ln2_lo), vt${ABC[N:N+4]});
73
74    // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2].
75    $for N in range(0, ELEMENTS_TILE, 4):
76      __m128 vp${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vc5, vt${ABC[N:N+4]}), vc4);
77
78    $for N in range(0, ELEMENTS_TILE, 4):
79      vp${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vp${ABC[N:N+4]}, vt${ABC[N:N+4]}), vc3);
80
81    $for N in range(0, ELEMENTS_TILE, 4):
82      vp${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vp${ABC[N:N+4]}, vt${ABC[N:N+4]}), vc2);
83
84    $for N in range(0, ELEMENTS_TILE, 4):
85      vp${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vp${ABC[N:N+4]}, vt${ABC[N:N+4]}), vc1);
86
87    // Reconstruct the final f value:
88    //   f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))))
89    //     = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))
90    //     = s + (t * s) * p
91    $for N in range(0, ELEMENTS_TILE, 4):
92      vt${ABC[N:N+4]} = _mm_mul_ps(vt${ABC[N:N+4]}, vs${ABC[N:N+4]});
93
94    $for N in range(0, ELEMENTS_TILE, 4):
95      __m128 vf${ABC[N:N+4]} = _mm_add_ps(_mm_mul_ps(vt${ABC[N:N+4]}, vp${ABC[N:N+4]}), vs${ABC[N:N+4]});
96
97    // For inputs below zero cutoff, replace output with +0.0f.
98    // Note that for NaN inputs, comparison result is false, and outputs are left unchanged.
99    $for N in range(0, ELEMENTS_TILE, 4):
100      vf${ABC[N:N+4]} = _mm_andnot_ps(_mm_cmplt_ps(vx${ABC[N:N+4]}, vdenorm_cutoff), vf${ABC[N:N+4]});
101
102    // Store ${ELEMENTS_TILE} (${SIMD_TILE}x4) outputs at a time.
103    _mm_storeu_ps(output, vf${ABC[0:4]});
104    $for N in range(4, ELEMENTS_TILE, 4):
105      _mm_storeu_ps(output + ${N}, vf${ABC[N:N+4]});
106    output += ${ELEMENTS_TILE};
107
108    // Accumulate computed exponents.
109    $for N in range(0, ELEMENTS_TILE, 4):
110      vacc${N % ACCUMULATORS} = _mm_add_ps(vacc${N % ACCUMULATORS}, vf${ABC[N:N+4]});
111  }
112  $if ACCUMULATORS > 1:
113    // Add up all accumulators to vacc0
114    $ACC_SLICE = 1
115    $while ACC_SLICE < ACCUMULATORS:
116      $for A in range(0, ACCUMULATORS, ACC_SLICE * 2):
117        $if A + ACC_SLICE < ACCUMULATORS:
118          vacc${A} = _mm_add_ps(vacc${A}, vacc${A + ACC_SLICE});
119      $ACC_SLICE *= 2
120
121  __m128 vacc = vacc0;
122  for (; elements >= 4 * sizeof(float); elements -= 4 * sizeof(float)) {
123    // Load 4 inputs at a time.
124    const __m128 vi = _mm_loadu_ps(input);
125    input += 4;
126
127    // Subtract maximum input x := i - i_max. This implies x <= 0.
128    const __m128 vx = _mm_sub_ps(vi, vi_max);
129
130    // Compute reduced argument elements := round(x / log(2)).
131    __m128 vn = _mm_add_ps(_mm_mul_ps(vx, vlog2e), vmagic_bias);
132
133    // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e.
134    // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly.
135    const __m128 vs = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(vn), 23));
136
137    // Subtract the large number back to get final elements := round(x / log(2)).
138    vn = _mm_sub_ps(vn, vmagic_bias);
139
140    // Compute reduced argument t := x - elements * log(2).
141    // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy.
142    __m128 vt = _mm_add_ps(_mm_mul_ps(vn, vminus_ln2_hi), vx);
143    vt = _mm_add_ps(_mm_mul_ps(vn, vminus_ln2_lo), vt);
144
145    // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2].
146    __m128 vp = _mm_add_ps(_mm_mul_ps(vc5, vt), vc4);
147    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc3);
148    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc2);
149    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc1);
150
151    // Reconstruct the final f value:
152    //   f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))))
153    //     = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))
154    //     = s + (t * s) * p
155    vt = _mm_mul_ps(vt, vs);
156    __m128 vf = _mm_add_ps(_mm_mul_ps(vt, vp), vs);
157
158    // For inputs below zero cutoff, replace output with +0.0f.
159    // Note that for NaN inputs, comparison result is false, and outputs are left unchanged.
160    vf = _mm_andnot_ps(_mm_cmplt_ps(vx, vdenorm_cutoff), vf);
161
162    // Store 4 outputs at a time.
163    _mm_storeu_ps(output, vf);
164    output += 4;
165
166    // Accumulate computed exponents.
167    vacc = _mm_add_ps(vacc, vf);
168  }
169  if (elements != 0) {
170    assert(elements >= 1 * sizeof(float));
171    assert(elements <= 3 * sizeof(float));
172    // Load 4 inputs at a time.
173    const __m128 vi = _mm_loadu_ps(input);
174
175    // Subtract maximum input x := i - i_max. This implies x <= 0.
176    const __m128 vx = _mm_sub_ps(vi, vi_max);
177
178    // Compute reduced argument elements := round(x / log(2)).
179    __m128 vn = _mm_add_ps(_mm_mul_ps(vx, vlog2e), vmagic_bias);
180
181    // Create a floating-point number s (scale) such that s == 2**elements for inputs which don't cause underflow, i.e.
182    // -87.33642 <= x <= 0.0, and -126 <= elements <= 0 accordingly.
183    const __m128 vs = _mm_castsi128_ps(_mm_slli_epi32(_mm_castps_si128(vn), 23));
184
185    // Subtract the large number back to get final elements := round(x / log(2)).
186    vn = _mm_sub_ps(vn, vmagic_bias);
187
188    // Compute reduced argument t := x - elements * log(2).
189    // Use Cody-Waite range reduction method (note two constants to represent log(2)) to improve accuracy.
190    __m128 vt = _mm_add_ps(_mm_mul_ps(vn, vminus_ln2_hi), vx);
191    vt = _mm_add_ps(_mm_mul_ps(vn, vminus_ln2_lo), vt);
192
193    // Compute degree-5 polynomial approximation for exp(t) on [-log(2)/2, log(2)/2].
194    __m128 vp = _mm_add_ps(_mm_mul_ps(vc5, vt), vc4);
195    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc3);
196    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc2);
197    vp = _mm_add_ps(_mm_mul_ps(vp, vt), vc1);
198
199    // Reconstruct the final f value:
200    //   f = s * (1 + t * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5)))))
201    //     = s + (t * s) * (c1 + t * (c2 + t * (c3 + t * (c4 + t * c5))))
202    //     = s + (t * s) * p
203    vt = _mm_mul_ps(vt, vs);
204    __m128 vf = _mm_add_ps(_mm_mul_ps(vt, vp), vs);
205
206    // For inputs below zero cutoff, replace output with +0.0f.
207    // Note that for NaN inputs, comparison result is false, and outputs are left unchanged.
208    vf = _mm_andnot_ps(_mm_cmplt_ps(vx, vdenorm_cutoff), vf);
209
210    if (elements & (2 * sizeof(float))) {
211      // Store 2 outputs at a time.
212      _mm_storel_pi((__m64*) output, vf);
213      output += 2;
214
215      // Accumulate 2 computed exponents.
216      vacc = _mm_add_ps(vacc, _mm_movelh_ps(vf, _mm_setzero_ps()));
217
218      vf = _mm_movehl_ps(vf, vf);
219    }
220    if (elements & (1 * sizeof(float))) {
221      // Store 1 output at a time.
222      _mm_store_ss(output, vf);
223
224      // Accumulate 1 computed exponent.
225      vacc = _mm_add_ss(vacc, vf);
226    }
227  }
228  // Reduce 4 elements in the SIMD register
229  vacc = _mm_add_ps(vacc, _mm_movehl_ps(vacc, vacc));
230  vacc = _mm_add_ss(vacc, _mm_shuffle_ps(vacc, vacc, _MM_SHUFFLE(2, 3, 0, 1)));
231  _mm_store_ss(sum, vacc);
232}
233