1*4bdc9457SAndroid Build Coastguard Worker// Copyright 2020 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 CHANNEL_TILE % 16 == 0 7*4bdc9457SAndroid Build Coastguard Worker$assert CHANNEL_TILE >= 16 8*4bdc9457SAndroid Build Coastguard Worker$assert ROW_TILE >= 1 9*4bdc9457SAndroid Build Coastguard Worker$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 10*4bdc9457SAndroid Build Coastguard Worker#include <assert.h> 11*4bdc9457SAndroid Build Coastguard Worker 12*4bdc9457SAndroid Build Coastguard Worker#include <immintrin.h> 13*4bdc9457SAndroid Build Coastguard Worker 14*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/intrinsics-polyfill.h> 15*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/math.h> 16*4bdc9457SAndroid Build Coastguard Worker#include <xnnpack/prelu.h> 17*4bdc9457SAndroid Build Coastguard Worker 18*4bdc9457SAndroid Build Coastguard Worker 19*4bdc9457SAndroid Build Coastguard Workervoid xnn_f32_prelu_ukernel__avx512f_${ROW_TILE}x${CHANNEL_TILE}( 20*4bdc9457SAndroid Build Coastguard Worker size_t rows, 21*4bdc9457SAndroid Build Coastguard Worker size_t channels, 22*4bdc9457SAndroid Build Coastguard Worker const float*restrict input, 23*4bdc9457SAndroid Build Coastguard Worker size_t input_stride, 24*4bdc9457SAndroid Build Coastguard Worker const float*restrict weights, 25*4bdc9457SAndroid Build Coastguard Worker float*restrict output, 26*4bdc9457SAndroid Build Coastguard Worker size_t output_stride) 27*4bdc9457SAndroid Build Coastguard Worker{ 28*4bdc9457SAndroid Build Coastguard Worker assert(rows != 0); 29*4bdc9457SAndroid Build Coastguard Worker assert(channels != 0); 30*4bdc9457SAndroid Build Coastguard Worker assert(channels % sizeof(float) == 0); 31*4bdc9457SAndroid Build Coastguard Worker 32*4bdc9457SAndroid Build Coastguard Worker const float* i0 = input; 33*4bdc9457SAndroid Build Coastguard Worker float* o0 = output; 34*4bdc9457SAndroid Build Coastguard Worker $for M in range(1, ROW_TILE): 35*4bdc9457SAndroid Build Coastguard Worker const float* i${M} = (const float*) ((uintptr_t) i${M-1} + input_stride); 36*4bdc9457SAndroid Build Coastguard Worker float* o${M} = (float*) ((uintptr_t) o${M-1} + output_stride); 37*4bdc9457SAndroid Build Coastguard Worker 38*4bdc9457SAndroid Build Coastguard Worker const size_t input_increment = input_stride * ${ROW_TILE} - channels; 39*4bdc9457SAndroid Build Coastguard Worker const size_t output_increment = output_stride * ${ROW_TILE} - channels; 40*4bdc9457SAndroid Build Coastguard Worker 41*4bdc9457SAndroid Build Coastguard Worker const __m512 vzero = _mm512_setzero_ps(); 42*4bdc9457SAndroid Build Coastguard Worker do { 43*4bdc9457SAndroid Build Coastguard Worker $for M in range(1, ROW_TILE): 44*4bdc9457SAndroid Build Coastguard Worker $if M % 2 == 0: 45*4bdc9457SAndroid Build Coastguard Worker if XNN_UNPREDICTABLE(rows <= ${M}) { 46*4bdc9457SAndroid Build Coastguard Worker i${M} = i${M-1}; 47*4bdc9457SAndroid Build Coastguard Worker o${M} = o${M-1}; 48*4bdc9457SAndroid Build Coastguard Worker } 49*4bdc9457SAndroid Build Coastguard Worker $else: 50*4bdc9457SAndroid Build Coastguard Worker if XNN_UNPREDICTABLE(rows < ${M+1}) { 51*4bdc9457SAndroid Build Coastguard Worker i${M} = i${M-1}; 52*4bdc9457SAndroid Build Coastguard Worker o${M} = o${M-1}; 53*4bdc9457SAndroid Build Coastguard Worker } 54*4bdc9457SAndroid Build Coastguard Worker 55*4bdc9457SAndroid Build Coastguard Worker const float* w = weights; 56*4bdc9457SAndroid Build Coastguard Worker size_t c = channels; 57*4bdc9457SAndroid Build Coastguard Worker for (; c >= ${CHANNEL_TILE} * sizeof(float); c -= ${CHANNEL_TILE} * sizeof(float)) { 58*4bdc9457SAndroid Build Coastguard Worker const __m512 vw${ABC[0:16]} = _mm512_load_ps(w); 59*4bdc9457SAndroid Build Coastguard Worker $for C in range(16, CHANNEL_TILE, 16): 60*4bdc9457SAndroid Build Coastguard Worker const __m512 vw${ABC[C:C+16]} = _mm512_load_ps(w + ${C}); 61*4bdc9457SAndroid Build Coastguard Worker w += ${CHANNEL_TILE}; 62*4bdc9457SAndroid Build Coastguard Worker 63*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 64*4bdc9457SAndroid Build Coastguard Worker const __m512 vi${M}x${ABC[0:16]} = _mm512_loadu_ps(i${M}); 65*4bdc9457SAndroid Build Coastguard Worker $for C in range(16, CHANNEL_TILE, 16): 66*4bdc9457SAndroid Build Coastguard Worker const __m512 vi${M}x${ABC[C:C+16]} = _mm512_loadu_ps(i${M} + ${C}); 67*4bdc9457SAndroid Build Coastguard Worker i${M} += ${CHANNEL_TILE}; 68*4bdc9457SAndroid Build Coastguard Worker 69*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 70*4bdc9457SAndroid Build Coastguard Worker $for C in range(0, CHANNEL_TILE, 16): 71*4bdc9457SAndroid Build Coastguard Worker const __mmask16 vsign${M}x${ABC[C:C+16]} = _mm512_cmp_ps_mask(vi${M}x${ABC[C:C+16]}, vzero, _CMP_LT_OQ); 72*4bdc9457SAndroid Build Coastguard Worker const __m512 vacc${M}x${ABC[C:C+16]} = _mm512_mask_mul_ps(vi${M}x${ABC[C:C+16]}, vsign${M}x${ABC[C:C+16]}, vi${M}x${ABC[C:C+16]}, vw${ABC[C:C+16]}); 73*4bdc9457SAndroid Build Coastguard Worker 74*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 75*4bdc9457SAndroid Build Coastguard Worker _mm512_storeu_ps(o${M}, vacc${M}x${ABC[0:16]}); 76*4bdc9457SAndroid Build Coastguard Worker $for C in range(16, CHANNEL_TILE, 16): 77*4bdc9457SAndroid Build Coastguard Worker _mm512_storeu_ps(o${M} + ${C}, vacc${M}x${ABC[C:C+16]}); 78*4bdc9457SAndroid Build Coastguard Worker o${M} += ${CHANNEL_TILE}; 79*4bdc9457SAndroid Build Coastguard Worker } 80*4bdc9457SAndroid Build Coastguard Worker $if CHANNEL_TILE > 16: 81*4bdc9457SAndroid Build Coastguard Worker for (; c >= 16 * sizeof(float); c -= 16 * sizeof(float)) { 82*4bdc9457SAndroid Build Coastguard Worker const __m512 vw = _mm512_load_ps(w); 83*4bdc9457SAndroid Build Coastguard Worker w += 16; 84*4bdc9457SAndroid Build Coastguard Worker 85*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 86*4bdc9457SAndroid Build Coastguard Worker const __m512 vi${M} = _mm512_loadu_ps(i${M}); 87*4bdc9457SAndroid Build Coastguard Worker i${M} += 16; 88*4bdc9457SAndroid Build Coastguard Worker 89*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 90*4bdc9457SAndroid Build Coastguard Worker const __mmask16 vsign${M} = _mm512_cmp_ps_mask(vi${M}, vzero, _CMP_LT_OQ); 91*4bdc9457SAndroid Build Coastguard Worker const __m512 vacc${M} = _mm512_mask_mul_ps(vi${M}, vsign${M}, vi${M}, vw); 92*4bdc9457SAndroid Build Coastguard Worker 93*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 94*4bdc9457SAndroid Build Coastguard Worker _mm512_storeu_ps(o${M}, vacc${M}); 95*4bdc9457SAndroid Build Coastguard Worker o${M} += 16; 96*4bdc9457SAndroid Build Coastguard Worker } 97*4bdc9457SAndroid Build Coastguard Worker if XNN_UNLIKELY(c != 0) { 98*4bdc9457SAndroid Build Coastguard Worker assert(c >= 1 * sizeof(float)); 99*4bdc9457SAndroid Build Coastguard Worker assert(c <= 15 * sizeof(float)); 100*4bdc9457SAndroid Build Coastguard Worker // Prepare mask for valid 32-bit elements (depends on c). 101*4bdc9457SAndroid Build Coastguard Worker const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << (c >> 2 /* log2(sizeof(float))*/)) - UINT32_C(1))); 102*4bdc9457SAndroid Build Coastguard Worker 103*4bdc9457SAndroid Build Coastguard Worker const __m512 vw = _mm512_maskz_loadu_ps(vmask, w); 104*4bdc9457SAndroid Build Coastguard Worker 105*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 106*4bdc9457SAndroid Build Coastguard Worker const __m512 vi${M} = _mm512_maskz_loadu_ps(vmask, i${M}); 107*4bdc9457SAndroid Build Coastguard Worker i${M} = (const float*) ((uintptr_t) i${M} + c); 108*4bdc9457SAndroid Build Coastguard Worker 109*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 110*4bdc9457SAndroid Build Coastguard Worker const __mmask16 vsign${M} = _mm512_cmp_ps_mask(vi${M}, vzero, _CMP_LT_OQ); 111*4bdc9457SAndroid Build Coastguard Worker const __m512 vacc${M} = _mm512_mask_mul_ps(vi${M}, vsign${M}, vi${M}, vw); 112*4bdc9457SAndroid Build Coastguard Worker 113*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 114*4bdc9457SAndroid Build Coastguard Worker _mm512_mask_storeu_ps(o${M}, vmask, vacc${M}); 115*4bdc9457SAndroid Build Coastguard Worker o${M} = (float*) ((uintptr_t) o${M} + c); 116*4bdc9457SAndroid Build Coastguard Worker } 117*4bdc9457SAndroid Build Coastguard Worker $for M in range(ROW_TILE): 118*4bdc9457SAndroid Build Coastguard Worker i${M} = (const float*) ((uintptr_t) i${M} + input_increment); 119*4bdc9457SAndroid Build Coastguard Worker o${M} = (float*) ((uintptr_t) o${M} + output_increment); 120*4bdc9457SAndroid Build Coastguard Worker rows = doz(rows, ${ROW_TILE}); 121*4bdc9457SAndroid Build Coastguard Worker } while (rows != 0); 122*4bdc9457SAndroid Build Coastguard Worker} 123