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