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 % 8 == 0 7$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 8#include <assert.h> 9 10#include <immintrin.h> 11 12#include <xnnpack/igemm.h> 13 14 15$ISA = {0: "avx", 3: "fma3"}[FMA] 16void xnn_f32_igemm_minmax_ukernel_${MR}x${NR}s4__${ISA}_broadcast( 17 size_t mr, 18 size_t nc, 19 size_t kc, 20 size_t ks, 21 const float**restrict a, 22 const float*restrict w, 23 float*restrict c, 24 size_t cm_stride, 25 size_t cn_stride, 26 size_t a_offset, 27 const float* zero, 28 const union xnn_f32_minmax_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS 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(ks != 0); 36 assert(ks % (${MR} * sizeof(void*)) == 0); 37 assert(a_offset % sizeof(float) == 0); 38 assert(a != NULL); 39 assert(w != NULL); 40 assert(c != NULL); 41 42 float* c0 = c; 43 $for M in range(1, MR): 44 float* c${M} = (float*) ((uintptr_t) c${M-1} + cm_stride); 45 $if M % 2 == 0: 46 if XNN_UNPREDICTABLE(mr <= ${M}) { 47 c${M} = c${M-1}; 48 } 49 $elif M + 1 == MR: 50 if XNN_UNPREDICTABLE(mr != ${M+1}) { 51 c${M} = c${M-1}; 52 } 53 $else: 54 if XNN_UNPREDICTABLE(mr < ${M+1}) { 55 c${M} = c${M-1}; 56 } 57 58 do { 59 __m256 vacc0x${ABC[0:8]} = _mm256_load_ps(w); 60 $for N in range(8, NR, 8): 61 __m256 vacc0x${ABC[N:N+8]} = _mm256_load_ps(w + ${N}); 62 $for M in range(1, MR): 63 $for N in range(0, NR, 8): 64 __m256 vacc${M}x${ABC[N:N+8]} = vacc0x${ABC[N:N+8]}; 65 w += ${NR}; 66 67 size_t p = ks; 68 do { 69 $for M in range(MR): 70 const float* restrict a${M} = a[${M}]; 71 assert(a${M} != NULL); 72 if XNN_UNPREDICTABLE(a${M} != zero) { 73 a${M} = (const float*) ((uintptr_t) a${M} + a_offset); 74 } 75 a += ${MR}; 76 77 size_t k = kc; 78 while (k >= 4 * sizeof(float)) { 79 $for M in range(MR): 80 __m256 va${M} = _mm256_broadcast_ps((const __m128*) a${M}); 81 a${M} += 4; 82 83 $for L in range(4): 84 85 $for N in range(0, NR, 8): 86 const __m256 vb${ABC[N:N+8]}c${L} = _mm256_load_ps(w + ${L * NR + N}); 87 88 $for N in range(0, NR, 8): 89 $for M in range(MR): 90 $if FMA == 3: 91 vacc${M}x${ABC[N:N+8]} = _mm256_fmadd_ps(va${M}, vb${ABC[N:N+8]}c${L}, vacc${M}x${ABC[N:N+8]}); 92 $else: 93 vacc${M}x${ABC[N:N+8]} = _mm256_add_ps(vacc${M}x${ABC[N:N+8]}, _mm256_mul_ps(va${M}, vb${ABC[N:N+8]}c${L})); 94 95 $if L + 1 != 4: 96 $for M in range(MR): 97 va${M} = _mm256_permute_ps(va${M}, _MM_SHUFFLE(0, 3, 2, 1)); 98 99 w += ${4 * NR}; 100 k -= 4 * sizeof(float); 101 } 102 if XNN_UNLIKELY(k != 0) { 103 $for M in range(MR): 104 __m256 va${M} = _mm256_broadcast_ps((const __m128*) a${M}); 105 a${M} = (const float*) ((uintptr_t) a${M} + k); 106 107 const __m256 vzero = _mm256_setzero_ps(); 108 $for L in range(4): 109 110 $for N in range(0, NR, 8): 111 const __m256 vb${ABC[N:N+8]}c${L} = _mm256_load_ps(w + ${L * NR + N}); 112 113 $for N in range(0, NR, 8): 114 $for M in range(MR): 115 $if FMA == 3: 116 vacc${M}x${ABC[N:N+8]} = _mm256_fmadd_ps(_mm256_and_ps(va${M}, _mm256_cmp_ps(vb${ABC[N:N+8]}c${L}, vzero, _CMP_NEQ_OQ)), vb${ABC[N:N+8]}c${L}, vacc${M}x${ABC[N:N+8]}); 117 $else: 118 vacc${M}x${ABC[N:N+8]} = _mm256_add_ps(vacc${M}x${ABC[N:N+8]}, _mm256_mul_ps(_mm256_and_ps(va${M}, _mm256_cmp_ps(vb${ABC[N:N+8]}c${L}, vzero, _CMP_NEQ_OQ)), vb${ABC[N:N+8]}c${L})); 119 120 $if L + 1 != 4: 121 $for M in range(MR): 122 va${M} = _mm256_permute_ps(va${M}, _MM_SHUFFLE(0, 3, 2, 1)); 123 124 w += ${4 * NR}; 125 } 126 p -= ${MR} * sizeof(void*); 127 } while (p != 0); 128 129 const __m256 vmin = _mm256_load_ps(params->avx.min); 130 $for N in range(0, NR, 8): 131 $for M in range(MR): 132 vacc${M}x${ABC[N:N+8]} = _mm256_max_ps(vacc${M}x${ABC[N:N+8]}, vmin); 133 134 const __m256 vmax = _mm256_load_ps(params->avx.max); 135 $for N in range(0, NR, 8): 136 $for M in range(MR): 137 vacc${M}x${ABC[N:N+8]} = _mm256_min_ps(vacc${M}x${ABC[N:N+8]}, vmax); 138 139 if XNN_LIKELY(nc >= ${NR}) { 140 $for M in reversed(range(MR)): 141 _mm256_storeu_ps(c${M}, vacc${M}x${ABC[0:8]}); 142 $for N in range(8, NR, 8): 143 _mm256_storeu_ps(c${M} + ${N}, vacc${M}x${ABC[N:N+8]}); 144 c${M} = (float*) ((uintptr_t) c${M} + cn_stride); 145 146 a = (const float**restrict) ((uintptr_t) a - ks); 147 nc -= ${NR}; 148 } else { 149 $for LOG2N in reversed(range(NR.bit_length())): 150 $if NR != 1 << LOG2N: 151 if (nc & ${1 << LOG2N}) { 152 $if LOG2N >= 3: 153 $for M in reversed(range(MR)): 154 _mm256_storeu_ps(c${M}, vacc${M}x${ABC[0:8]}); 155 $for N in range(8, 1 << LOG2N, 8): 156 _mm256_storeu_ps(c${M} + ${N}, vacc${M}x${ABC[N:N+8]}); 157 158 $for M in reversed(range(MR)): 159 $for N in range(0, 1 << (LOG2N - 1), 8): 160 vacc${M}x${ABC[N:N+8]} = vacc${M}x${ABC[N + (1 << LOG2N):N + (1 << LOG2N)+8]}; 161 162 $for M in reversed(range(MR)): 163 c${M} += ${1 << LOG2N}; 164 $elif LOG2N == 2: 165 $for M in reversed(range(MR)): 166 _mm_storeu_ps(c${M}, vacc${M}x${ABC[0:4]}); 167 168 $for M in reversed(range(MR)): 169 vacc${M}x${ABC[0:4]} = _mm256_extractf128_ps(vacc${M}x${ABC[0:8]}, 1); 170 171 $for M in reversed(range(MR)): 172 c${M} += 4; 173 $elif LOG2N == 1: 174 $for M in reversed(range(MR)): 175 _mm_storel_pi((__m64*) c${M}, vacc${M}x${ABC[0:4]}); 176 177 $for M in reversed(range(MR)): 178 vacc${M}x${ABC[0:4]} = _mm_movehl_ps(vacc${M}x${ABC[0:4]}, vacc${M}x${ABC[0:4]}); 179 180 $for M in reversed(range(MR)): 181 c${M} += 2; 182 $elif LOG2N == 0: 183 $for M in reversed(range(MR)): 184 _mm_store_ss(c${M}, vacc${M}x${ABC[0:4]}); 185 } 186 $if LOG2N == 3: 187 $for M in reversed(range(MR)): 188 __m128 vacc${M}x${ABC[0:4]} = _mm256_castps256_ps128(vacc${M}x${ABC[0:8]}); 189 190 nc = 0; 191 } 192 } while (nc != 0); 193} 194