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$assert BATCH_TILE % 16 == 0 7$assert BATCH_TILE >= 16 8$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 9$assert OP in ["ADD", "DIV", "RDIV", "MAX", "MIN", "MUL", "SUB", "RSUB"] 10#include <assert.h> 11 12#include <immintrin.h> 13 14#include <xnnpack/common.h> 15#include <xnnpack/intrinsics-polyfill.h> 16#include <xnnpack/vbinary.h> 17 18 19$_MM512_OP_PS = { 20$ "ADD": lambda x: "_mm512_add_ps(%s, vb)" % x, 21$ "DIV": lambda x: "_mm512_div_ps(%s, vb)" % x, 22$ "RDIV": lambda x: "_mm512_div_ps(vb, %s)" % x, 23$ "MAX": lambda x: "_mm512_max_ps(%s, vb)" % x, 24$ "MIN": lambda x: "_mm512_min_ps(%s, vb)" % x, 25$ "MUL": lambda x: "_mm512_mul_ps(%s, vb)" % x, 26$ "SUB": lambda x: "_mm512_sub_ps(%s, vb)" % x, 27$ "RSUB": lambda x: "_mm512_sub_ps(vb, %s)" % x, 28$}[OP] 29void xnn_f32_v${OP.lower()}c_ukernel__avx512f_x${BATCH_TILE}( 30 size_t n, 31 const float* a, 32 const float* b, 33 float* y, 34 const union xnn_f32_output_params params[restrict static 1]) 35{ 36 assert(n != 0); 37 assert(n % sizeof(float) == 0); 38 39 const __m512 vy_min = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.min)); 40 const __m512 vy_max = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.max)); 41 42 const __m512 vb = _mm512_set1_ps(*b); 43 for (; n >= ${BATCH_TILE} * sizeof(float); n -= ${BATCH_TILE} * sizeof(float)) { 44 const __m512 va${ABC[0:16]} = _mm512_loadu_ps(a); 45 $for N in range(16, BATCH_TILE, 16): 46 const __m512 va${ABC[N:N+16]} = _mm512_loadu_ps(a + ${N}); 47 a += ${BATCH_TILE}; 48 49 $for N in range(0, BATCH_TILE, 16): 50 __m512 vy${ABC[N:N+16]} = ${_MM512_OP_PS("va" + ABC[N:N+16])}; 51 52 $for N in range(0, BATCH_TILE, 16): 53 vy${ABC[N:N+16]} = _mm512_max_ps(vy${ABC[N:N+16]}, vy_min); 54 55 $for N in range(0, BATCH_TILE, 16): 56 vy${ABC[N:N+16]} = _mm512_min_ps(vy${ABC[N:N+16]}, vy_max); 57 58 _mm512_storeu_ps(y, vy${ABC[0:16]}); 59 $for N in range(16, BATCH_TILE, 16): 60 _mm512_storeu_ps(y + ${N}, vy${ABC[N:N+16]}); 61 y += ${BATCH_TILE}; 62 } 63 $if BATCH_TILE >= 16: 64 for (; n >= 16 * sizeof(float); n -= 16 * sizeof(float)) { 65 const __m512 va = _mm512_loadu_ps(a); 66 a += 16; 67 68 __m512 vy = ${_MM512_OP_PS("va")}; 69 vy = _mm512_max_ps(vy, vy_min); 70 vy = _mm512_min_ps(vy, vy_max); 71 _mm512_storeu_ps(y, vy); 72 y += 16; 73 } 74 if XNN_UNLIKELY(n != 0) { 75 assert(n >= 1 * sizeof(float)); 76 assert(n <= 15 * sizeof(float)); 77 // Prepare mask for valid 32-bit elements (depends on n). 78 n >>= 2 /* log2(sizeof(float)) */; 79 const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << n) - UINT32_C(1))); 80 81 const __m512 va = _mm512_maskz_loadu_ps(vmask, a); 82 83 __m512 vy = ${_MM512_OP_PS("va")}; 84 vy = _mm512_max_ps(vy, vy_min); 85 vy = _mm512_min_ps(vy, vy_max); 86 _mm512_mask_storeu_ps(y, vmask, vy); 87 } 88} 89