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$assert BATCH_TILE % 8 == 0 7$assert BATCH_TILE >= 8 8$ABC = "01234567456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 9$assert OP in ["ADD", "DIV", "RDIV", "MAX", "MIN", "MUL", "SUB", "RSUB", "SQRDIFF"] 10$assert ACTIVATION in ["LINEAR", "MINMAX"] 11#include <assert.h> 12 13#include <immintrin.h> 14 15#include <xnnpack/common.h> 16#include <xnnpack/vbinary.h> 17 18 19$_MM256_OP_PS = { 20$ "ADD": lambda x: "_mm256_add_ps(%s, vb)" % x, 21$ "DIV": lambda x: "_mm256_div_ps(%s, vb)" % x, 22$ "RDIV": lambda x: "_mm256_div_ps(vb, %s)" % x, 23$ "MAX": lambda x: "_mm256_max_ps(%s, vb)" % x, 24$ "MIN": lambda x: "_mm256_min_ps(%s, vb)" % x, 25$ "MUL": lambda x: "_mm256_mul_ps(%s, vb)" % x, 26$ "SUB": lambda x: "_mm256_sub_ps(%s, vb)" % x, 27$ "RSUB": lambda x: "_mm256_sub_ps(vb, %s)" % x, 28$ "SQRDIFF": lambda x: "_mm256_sub_ps(%s, vb)" % x, 29$}[OP] 30$SUFFIX = {"LINEAR": "", "MINMAX": "_minmax"}[ACTIVATION] 31$PARAMS = {"LINEAR": "xnn_f16_default_params", "MINMAX": "xnn_f16_minmax_params"}[ACTIVATION] 32void xnn_f16_v${OP.lower()}c${SUFFIX}_ukernel__f16c_x${BATCH_TILE}( 33 size_t n, 34 const void* restrict a_ptr, 35 const void* restrict b_ptr, 36 void* restrict y_ptr, 37 const union ${PARAMS} params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS 38{ 39 assert(n != 0); 40 assert(n % sizeof(uint16_t) == 0); 41 assert(a_ptr != NULL); 42 assert(b_ptr != NULL); 43 assert(y_ptr != NULL); 44 45 const uint16_t* a = (const uint16_t*) a_ptr; 46 const uint16_t* b = (const uint16_t*) b_ptr; 47 uint16_t* y = (uint16_t*) y_ptr; 48 49 $if ACTIVATION == "MINMAX": 50 const __m256 vy_min = _mm256_load_ps(params->avx.min); 51 const __m256 vy_max = _mm256_load_ps(params->avx.max); 52 53 const __m256 vb = _mm256_cvtph_ps(_mm_set1_epi16((short) *b)); 54 $if BATCH_TILE > 8: 55 for (; n >= ${BATCH_TILE} * sizeof(uint16_t); n -= ${BATCH_TILE} * sizeof(uint16_t)) { 56 const __m256 va${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a)); 57 $for N in range(8, BATCH_TILE, 8): 58 const __m256 va${ABC[N:N+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (a + ${N}))); 59 a += ${BATCH_TILE}; 60 61 $for N in range(0, BATCH_TILE, 8): 62 __m256 vy${ABC[N:N+8]} = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va" + ABC[N:N+8])}, _MM_FROUND_NO_EXC)); 63 64 $if OP == "SQRDIFF": 65 $for N in range(0, BATCH_TILE, 8): 66 vy${ABC[N:N+8]} = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy${ABC[N:N+8]}, vy${ABC[N:N+8]}), _MM_FROUND_NO_EXC)); 67 68 $if ACTIVATION == "MINMAX": 69 $for N in range(0, BATCH_TILE, 8): 70 vy${ABC[N:N+8]} = _mm256_max_ps(vy${ABC[N:N+8]}, vy_min); 71 72 $for N in range(0, BATCH_TILE, 8): 73 vy${ABC[N:N+8]} = _mm256_min_ps(vy${ABC[N:N+8]}, vy_max); 74 75 _mm_storeu_si128((__m128i*) y, _mm256_cvtps_ph(vy${ABC[0:8]}, _MM_FROUND_NO_EXC)); 76 $for N in range(8, BATCH_TILE, 8): 77 _mm_storeu_si128((__m128i*) (y + ${N}), _mm256_cvtps_ph(vy${ABC[N:N+8]}, _MM_FROUND_NO_EXC)); 78 y += ${BATCH_TILE}; 79 } 80 for (; n >= 8 * sizeof(uint16_t); n -= 8 * sizeof(uint16_t)) { 81 const __m256 va = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a)); 82 a += 8; 83 84 __m256 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va")}, _MM_FROUND_NO_EXC)); 85 $if OP == "SQRDIFF": 86 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy, vy), _MM_FROUND_NO_EXC)); 87 88 $if ACTIVATION == "MINMAX": 89 vy = _mm256_max_ps(vy, vy_min); 90 vy = _mm256_min_ps(vy, vy_max); 91 92 _mm_storeu_si128((__m128i*) y, _mm256_cvtps_ph(vy, _MM_FROUND_NO_EXC)); 93 y += 8; 94 } 95 if XNN_UNLIKELY(n != 0) { 96 const __m256 va = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) a)); 97 98 __m256 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(${_MM256_OP_PS("va")}, _MM_FROUND_NO_EXC)); 99 $if OP == "SQRDIFF": 100 vy = _mm256_cvtph_ps(_mm256_cvtps_ph(_mm256_mul_ps(vy, vy), _MM_FROUND_NO_EXC)); 101 102 $if ACTIVATION == "MINMAX": 103 vy = _mm256_max_ps(vy, vy_min); 104 vy = _mm256_min_ps(vy, vy_max); 105 106 __m128i vh = _mm256_cvtps_ph(vy, _MM_FROUND_NO_EXC); 107 if (n & (4 * sizeof(uint16_t))) { 108 _mm_storel_epi64((__m128i*) y, vh); 109 vh = _mm_unpackhi_epi64(vh, vh); 110 y += 4; 111 } 112 if (n & (2 * sizeof(uint16_t))) { 113 *((uint32_t*) y) = (uint32_t) _mm_cvtsi128_si32(vh); 114 vh = _mm_srli_epi64(vh, 32); 115 y += 2; 116 } 117 if (n & (1 * sizeof(uint16_t))) { 118 *y = (uint16_t) _mm_extract_epi16(vh, 0); 119 } 120 } 121} 122