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