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 NR % 16 == 0 7$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 8#include <assert.h> 9 10#include <immintrin.h> 11 12#include <xnnpack/gemm.h> 13#include <xnnpack/intrinsics-polyfill.h> 14 15 16void xnn_f32_gemm${"inc" if INC else ""}_ukernel_${MR}x${NR}__avx512f_broadcast( 17 size_t mr, 18 size_t nc, 19 size_t kc, 20 const float*restrict a, 21 size_t a_stride, 22 const float*restrict w, 23 float*restrict c, 24 size_t cm_stride, 25 size_t cn_stride, 26 $if INC: 27 const float*restrict acc, 28 const union xnn_f32_output_params params[restrict static 1]) 29{ 30 assert(mr != 0); 31 assert(mr <= ${MR}); 32 assert(nc != 0); 33 assert(kc != 0); 34 assert(kc % sizeof(float) == 0); 35 assert(a != NULL); 36 assert(w != NULL); 37 assert(c != NULL); 38 $if INC: 39 assert(acc != NULL); 40 41 const float* a0 = a; 42 float* c0 = c; 43 $for M in range(1, MR): 44 const float* a${M} = (const float*) ((uintptr_t) a${M-1} + a_stride); 45 float* c${M} = (float*) ((uintptr_t) c${M-1} + cm_stride); 46 $if M % 2 == 0: 47 if XNN_UNPREDICTABLE(mr <= ${M}) { 48 a${M} = a${M-1}; 49 c${M} = c${M-1}; 50 } 51 $elif M + 1 == MR: 52 if XNN_UNPREDICTABLE(mr != ${M+1}) { 53 a${M} = a${M-1}; 54 c${M} = c${M-1}; 55 } 56 $else: 57 if XNN_UNPREDICTABLE(mr < ${M+1}) { 58 a${M} = a${M-1}; 59 c${M} = c${M-1}; 60 } 61 62 do { 63 $if INC: 64 $for M in range(MR): 65 $for N in range(0, NR, 16): 66 __m512 vacc${M}x${ABC[N:N+16]} = _mm512_load_ps(acc + ${M*NR+N}); 67 acc += ${MR*NR}; 68 $else: 69 __m512 vacc0x${ABC[0:16]} = _mm512_load_ps(w); 70 $for N in range(16, NR, 16): 71 __m512 vacc0x${ABC[N:N+16]} = _mm512_load_ps(w + ${N}); 72 $for M in range(1, MR): 73 $for N in range(0, NR, 16): 74 __m512 vacc${M}x${ABC[N:N+16]} = vacc0x${ABC[N:N+16]}; 75 w += ${NR}; 76 77 size_t k = kc; 78 do { 79 const __m512 vb${ABC[0:16]} = _mm512_load_ps(w); 80 $for N in range(16, NR, 16): 81 const __m512 vb${ABC[N:N+16]} = _mm512_load_ps(w + ${N}); 82 w += ${NR}; 83 84 $for N in range(0, NR, 16): 85 $for M in range(MR): 86 vacc${M}x${ABC[N:N+16]} = _mm512_fmadd_ps(_mm512_set1_ps(*a${M}), vb${ABC[N:N+16]}, vacc${M}x${ABC[N:N+16]}); 87 88 $for M in range(MR): 89 a${M} += 1; 90 91 k -= sizeof(float); 92 } while (k != 0); 93 94 const __m512 vmax = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.max)); 95 $for N in range(0, NR, 16): 96 $for M in range(MR): 97 vacc${M}x${ABC[N:N+16]} = _mm512_min_ps(vacc${M}x${ABC[N:N+16]}, vmax); 98 99 const __m512 vmin = _mm512_broadcast_f32x4(_mm_load_ps(params->sse.min)); 100 $for N in range(0, NR, 16): 101 $for M in range(MR): 102 vacc${M}x${ABC[N:N+16]} = _mm512_max_ps(vacc${M}x${ABC[N:N+16]}, vmin); 103 104 if XNN_LIKELY(nc >= ${NR}) { 105 $for M in reversed(range(MR)): 106 _mm512_storeu_ps(c${M}, vacc${M}x${ABC[0:16]}); 107 $for N in range(16, NR, 16): 108 _mm512_storeu_ps(c${M} + ${N}, vacc${M}x${ABC[N:N+16]}); 109 c${M} = (float*) ((uintptr_t) c${M} + cn_stride); 110 111 $for M in reversed(range(MR)): 112 a${M} = (const float*) ((uintptr_t) a${M} - kc); 113 114 nc -= ${NR}; 115 } else { 116 $for LOG2N in reversed(range(4, NR.bit_length())): 117 $if NR != 1 << LOG2N: 118 if (nc & ${1 << LOG2N}) { 119 $if LOG2N >= 4: 120 $for M in reversed(range(MR)): 121 _mm512_storeu_ps(c${M}, vacc${M}x${ABC[0:16]}); 122 $for N in range(16, 1 << LOG2N, 16): 123 _mm512_storeu_ps(c${M} + ${N}, vacc${M}x${ABC[N:N+16]}); 124 125 $for M in reversed(range(MR)): 126 $for N in range(0, 1 << (LOG2N - 1), 16): 127 vacc${M}x${ABC[N:N+16]} = vacc${M}x${ABC[N + (1 << LOG2N):N + (1 << LOG2N)+16]}; 128 129 $for M in reversed(range(MR)): 130 c${M} += ${1 << LOG2N}; 131 } 132 $if LOG2N == 4: 133 if (nc & 15) { 134 // Prepare mask for valid 32-bit elements (depends on nc). 135 const __mmask16 vmask = _cvtu32_mask16((uint16_t) ((uint32_t) (UINT32_C(1) << nc) - UINT32_C(1))); 136 137 $for M in reversed(range(MR)): 138 _mm512_mask_storeu_ps(c${M}, vmask, vacc${M}x${ABC[0:16]}); 139 } 140 141 nc = 0; 142 } 143 } while (nc != 0); 144} 145