1// Copyright 2020 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 ["RNDNE", "RNDZ", "RNDU", "RNDD"] 10#include <assert.h> 11 12#include <immintrin.h> 13 14#include <xnnpack/common.h> 15#include <xnnpack/intrinsics-polyfill.h> 16#include <xnnpack/math.h> 17#include <xnnpack/vunary.h> 18 19 20$_MM_FROUND_TO_FLAG = { 21$ "RNDNE": "_MM_FROUND_TO_NEAREST_INT", 22$ "RNDZ": "_MM_FROUND_TO_ZERO", 23$ "RNDU": "_MM_FROUND_TO_POS_INF", 24$ "RNDD": "_MM_FROUND_TO_NEG_INF", 25$}[OP] 26void xnn_f32_v${OP.lower()}_ukernel__avx512f_x${BATCH_TILE}( 27 size_t n, 28 const float* x, 29 float* y, 30 const union xnn_f32_rnd_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_DISABLE_TSAN 31{ 32 assert(n != 0); 33 assert(n % sizeof(float) == 0); 34 35 for (; n >= ${BATCH_TILE} * sizeof(float); n -= ${BATCH_TILE} * sizeof(float)) { 36 const __m512 vx${ABC[0:16]} = _mm512_loadu_ps(x); 37 $for N in range(16, BATCH_TILE, 16): 38 const __m512 vx${ABC[N:N+16]} = _mm512_loadu_ps(x + ${N}); 39 x += ${BATCH_TILE}; 40 41 $for N in range(0, BATCH_TILE, 16): 42 const __m512 vy${ABC[N:N+16]} = _mm512_roundscale_ps(vx${ABC[N:N+16]}, ${_MM_FROUND_TO_FLAG}); 43 44 _mm512_storeu_ps(y, vy${ABC[0:16]}); 45 $for N in range(16, BATCH_TILE, 16): 46 _mm512_storeu_ps(y + ${N}, vy${ABC[N:N+16]}); 47 y += ${BATCH_TILE}; 48 } 49 $if BATCH_TILE > 16: 50 for (; n >= 16 * sizeof(float); n -= 16 * sizeof(float)) { 51 const __m512 vx = _mm512_loadu_ps(x); 52 x += 16; 53 54 const __m512 vy = _mm512_roundscale_ps(vx, ${_MM_FROUND_TO_FLAG}); 55 56 _mm512_storeu_ps(y, vy); 57 y += 16; 58 } 59 if XNN_UNLIKELY(n != 0) { 60 assert(n >= 1 * sizeof(float)); 61 assert(n <= 15 * sizeof(float)); 62 // Prepare mask for valid 32-bit elements (depends on n). 63 n >>= 2 /* log2(sizeof(float)) */; 64 const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << n) - UINT32_C(1))); 65 66 const __m512 vx = _mm512_maskz_loadu_ps(vmask, x); 67 const __m512 vy = _mm512_maskz_roundscale_ps(vmask, vx, ${_MM_FROUND_TO_FLAG}); 68 _mm512_mask_storeu_ps(y, vmask, vy); 69 } 70} 71