1 // Copyright 2022 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/common.h>
11 #include <xnnpack/rmax.h>
12
13
xnn_f16_rmax_ukernel__f16c(size_t batch,const void * input,void * output)14 void xnn_f16_rmax_ukernel__f16c(
15 size_t batch,
16 const void* input,
17 void* output) XNN_OOB_READS
18 {
19 assert(batch != 0);
20 assert(batch % sizeof(uint16_t) == 0);
21
22 const uint16_t* i = (const uint16_t*) input;
23 __m128i vmax_init = _mm_shufflelo_epi16(_mm_loadl_epi64((const __m128i*) i), _MM_SHUFFLE(0, 0, 0, 0));
24 vmax_init = _mm_unpacklo_epi64(vmax_init, vmax_init);
25 __m256 vmax0 = _mm256_cvtph_ps(vmax_init);
26 __m256 vmax1 = vmax0;
27 __m256 vmax2 = vmax0;
28 __m256 vmax3 = vmax0;
29 for (; batch >= 32 * sizeof(uint16_t); batch -= 32 * sizeof(uint16_t)) {
30 const __m256 vx0 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i));
31 const __m256 vx1 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (i + 8)));
32 const __m256 vx2 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (i + 16)));
33 const __m256 vx3 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (i + 24)));
34 i += 32;
35
36 vmax0 = _mm256_max_ps(vmax0, vx0);
37 vmax1 = _mm256_max_ps(vmax1, vx1);
38 vmax2 = _mm256_max_ps(vmax2, vx2);
39 vmax3 = _mm256_max_ps(vmax3, vx3);
40 }
41 __m256 vmax = _mm256_max_ps(_mm256_max_ps(vmax0, vmax1), _mm256_max_ps(vmax2, vmax3));
42 for (; batch >= 8 * sizeof(uint16_t); batch -= 8 * sizeof(uint16_t)) {
43 const __m256 vx = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i));
44 i += 8;
45 vmax = _mm256_max_ps(vmax, vx);
46 }
47 __m128 vmax_lo = _mm_max_ps(_mm256_castps256_ps128(vmax), _mm256_extractf128_ps(vmax, 1));
48 if XNN_UNLIKELY(batch != 0) {
49 const __m256 vx = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i));
50 __m128 vx_lo = _mm256_castps256_ps128(vx);
51 if (batch & (4 * sizeof(uint16_t))) {
52 vmax_lo = _mm_max_ps(vmax_lo, vx_lo);
53 vx_lo = _mm256_extractf128_ps(vx, 1);
54 }
55 if (batch & (2 * sizeof(uint16_t))) {
56 vmax_lo = _mm_blend_ps(_mm_max_ps(vmax_lo, vx_lo), vmax_lo, 0xC);
57 vx_lo = _mm_movehl_ps(vx_lo, vx_lo);
58 }
59 if (batch & (1 * sizeof(uint16_t))) {
60 vmax_lo = _mm_max_ss(vmax_lo, vx_lo);
61 }
62 }
63 vmax_lo = _mm_max_ps(vmax_lo, _mm_movehl_ps(vmax_lo, vmax_lo));
64 vmax_lo = _mm_max_ss(vmax_lo, _mm_movehdup_ps(vmax_lo));
65 *((uint16_t*) output) = (uint16_t) _mm_extract_epi16(_mm_cvtps_ph(vmax_lo, _MM_FROUND_NO_EXC), 0);
66 }
67