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 #include <assert.h>
7
8 #include <immintrin.h>
9
10 #include <xnnpack/rmax.h>
11
12
xnn_f32_rmax_ukernel__avx(size_t n,const float * x,float * y)13 void xnn_f32_rmax_ukernel__avx(
14 size_t n,
15 const float* x,
16 float* y)
17 {
18 assert(n != 0);
19 assert(n % sizeof(float) == 0);
20
21 __m256 vmax0 = _mm256_broadcast_ss(x);
22 __m256 vmax1 = vmax0;
23 __m256 vmax2 = vmax0;
24 __m256 vmax3 = vmax0;
25 for (; n >= 128; n -= 128) {
26 const __m256 vx0 = _mm256_loadu_ps(x);
27 const __m256 vx1 = _mm256_loadu_ps(x + 8);
28 const __m256 vx2 = _mm256_loadu_ps(x + 16);
29 const __m256 vx3 = _mm256_loadu_ps(x + 24);
30 x += 32;
31
32 vmax0 = _mm256_max_ps(vmax0, vx0);
33 vmax1 = _mm256_max_ps(vmax1, vx1);
34 vmax2 = _mm256_max_ps(vmax2, vx2);
35 vmax3 = _mm256_max_ps(vmax3, vx3);
36 }
37 __m256 vmax = _mm256_max_ps(_mm256_max_ps(vmax0, vmax1), _mm256_max_ps(vmax2, vmax3));
38 for (; n >= 32; n -= 32) {
39 const __m256 vx = _mm256_loadu_ps(x);
40 vmax = _mm256_max_ps(vmax, vx);
41 x += 8;
42 }
43 __m128 vmax_lo = _mm_max_ps(_mm256_castps256_ps128(vmax), _mm256_extractf128_ps(vmax, 1));
44 vmax_lo = _mm_max_ps(vmax_lo, _mm_movehl_ps(vmax_lo, vmax_lo));
45 vmax_lo = _mm_max_ss(vmax_lo, _mm_shuffle_ps(vmax_lo, vmax_lo, _MM_SHUFFLE(3, 3, 1, 1)));
46 if XNN_UNLIKELY(n != 0) {
47 do {
48 vmax_lo = _mm_max_ss(vmax_lo, _mm_load_ss(x));
49 x += 1;
50 n -= 4;
51 } while (n != 0);
52 }
53 _mm_store_ss(y, vmax_lo);
54 }
55