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$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" 7$assert REQUANTIZATION == "FP32" 8$assert DATATYPE in ["QC8", "QS8", "QU8"] 9$assert VARIANT in ["LD256", "EXTENDED"] 10$assert MR <= 4 11#include <assert.h> 12 13#include <immintrin.h> 14 15#include <xnnpack/igemm.h> 16#include <xnnpack/intrinsics-polyfill.h> 17#include <xnnpack/math.h> 18 19 20$GEMM_SUFFIX = "_xw" if VARIANT == "EXTENDED" else "" 21$PARAMS_STRUCT = REQUANTIZATION.lower() + "_avx512" 22$PARAMS_UNION = "xnn_%s_conv_minmax_params" % DATATYPE.lower() 23$XINT8_T = "uint8_t" if DATATYPE == "QU8" else "int8_t" 24void xnn_${DATATYPE.lower()}_igemm${GEMM_SUFFIX}_minmax_fp32_ukernel_${MR}x16c8__avx512skx( 25 size_t mr, 26 size_t nc, 27 size_t kc, 28 size_t ks, 29 const ${XINT8_T}** restrict a, 30 const void* restrict w, 31 ${XINT8_T}* restrict c, 32 size_t cm_stride, 33 size_t cn_stride, 34 size_t a_offset, 35 const ${XINT8_T}* zero, 36 const union ${PARAMS_UNION} params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS 37{ 38 assert(mr != 0); 39 assert(mr <= ${MR}); 40 assert(nc != 0); 41 assert(kc != 0); 42 assert(kc % sizeof(${XINT8_T}) == 0); 43 assert(a != NULL); 44 assert(w != NULL); 45 assert(c != NULL); 46 47 kc = round_up_po2(kc, 8); 48 ${XINT8_T}* c0 = c; 49 $for M in range(1, MR): 50 ${XINT8_T}* c${M} = (${XINT8_T}*) ((uintptr_t) c${M-1} + cm_stride); 51 $if M % 2 == 0: 52 if XNN_UNPREDICTABLE(mr <= ${M}) { 53 c${M} = c${M-1}; 54 } 55 $elif M + 1 == MR: 56 if XNN_UNPREDICTABLE(mr != ${M+1}) { 57 c${M} = c${M-1}; 58 } 59 $else: 60 if XNN_UNPREDICTABLE(mr < ${M+1}) { 61 c${M} = c${M-1}; 62 } 63 64 const __mmask16 vbias_mask = _cvtu32_mask16(0x1111); 65 $if DATATYPE != "QC8": 66 const __m512 vscale = _mm512_load_ps(params->${PARAMS_STRUCT}.scale); 67 const __m512 voutput_max_less_zero_point = _mm512_load_ps(params->${PARAMS_STRUCT}.output_max_less_zero_point); 68 $if MR > 1: 69 const __m512i voutput_zero_point = _mm512_load_si512(params->${PARAMS_STRUCT}.output_zero_point); 70 $else: 71 const __m256i voutput_zero_point = _mm256_load_si256((const __m256i*) params->${PARAMS_STRUCT}.output_zero_point); 72 $if MR > 2: 73 const __m512i voutput_min = _mm512_load_si512(params->${PARAMS_STRUCT}.output_min); 74 $elif MR == 2: 75 const __m256i voutput_min = _mm256_load_si256((const __m256i*) params->${PARAMS_STRUCT}.output_min); 76 $else: 77 const __m128i voutput_min = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_min); 78 do { 79 __m512i vacc0x0123 = _mm512_maskz_expandloadu_epi32(vbias_mask, w); 80 $for N in range(4, 16, 4): 81 __m512i vacc0x${ABC[N:N+4]} = _mm512_maskz_expandloadu_epi32(vbias_mask, (const void*) ((const int32_t*) w + ${N})); 82 $for M in range(1, MR): 83 $for N in range(0, 16, 4): 84 __m512i vacc${M}x${ABC[N:N+4]} = vacc0x${ABC[N:N+4]}; 85 w = (const void*) ((const int32_t*) w + 16); 86 87 size_t p = ks; 88 do { 89 $for M in range(MR): 90 const ${XINT8_T}* restrict a${M} = a[${M}]; 91 if XNN_UNPREDICTABLE(a${M} != zero) { 92 a${M} = (const ${XINT8_T}*) ((uintptr_t) a${M} + a_offset); 93 } 94 a += ${MR}; 95 96 size_t k = 0; 97 $if DATATYPE == "QU8": 98 const __m512i vb_zero_point = _mm512_load_si512(params->${PARAMS_STRUCT}.kernel_zero_point); 99 while (k < kc) { 100 $for M in range(MR): 101 $if DATATYPE == "QU8": 102 const __m512i va${M} = _mm512_broadcast_i32x4(_mm_cvtepu8_epi16(_mm_loadl_epi64((const __m128i*) a${M}))); 103 $else: 104 const __m512i va${M} = _mm512_broadcast_i32x4(_mm_cvtepi8_epi16(_mm_loadl_epi64((const __m128i*) a${M}))); 105 a${M} += 8; 106 107 $for N in range(0, 16, 4): 108 $if VARIANT == "EXTENDED": 109 $if N == 0: 110 const __m512i vb${ABC[N:N+4]} = _mm512_load_si512((const __m512i*) w); 111 $else: 112 const __m512i vb${ABC[N:N+4]} = _mm512_load_si512((const __m512i*) ((const int16_t*) w + ${N * 8})); 113 $else: 114 $if DATATYPE == "QU8": 115 $if N == 0: 116 const __m512i vb${ABC[N:N+4]} = _mm512_sub_epi16(_mm512_cvtepu8_epi16(_mm256_load_si256((const __m256i*) w)), vb_zero_point); 117 $else: 118 const __m512i vb${ABC[N:N+4]} = _mm512_sub_epi16(_mm512_cvtepu8_epi16(_mm256_load_si256((const __m256i*) ((const ${XINT8_T}*) w + ${N * 8}))), vb_zero_point); 119 $else: 120 $if N == 0: 121 const __m512i vb${ABC[N:N+4]} = _mm512_cvtepi8_epi16(_mm256_load_si256((const __m256i*) w)); 122 $else: 123 const __m512i vb${ABC[N:N+4]} = _mm512_cvtepi8_epi16(_mm256_load_si256((const __m256i*) ((const ${XINT8_T}*) w + ${N * 8}))); 124 125 $for M in range(MR): 126 vacc${M}x${ABC[N:N+4]} = _mm512_add_epi32(vacc${M}x${ABC[N:N+4]}, _mm512_madd_epi16(va${M}, vb${ABC[N:N+4]})); 127 128 $if VARIANT == "EXTENDED": 129 w = (const void*) ((const int16_t*) w + 128); 130 $else: 131 w = (const void*) ((const ${XINT8_T}*) w + 128); 132 k += 8 * sizeof(${XINT8_T}); 133 } 134 p -= ${MR} * sizeof(void*); 135 } while (p != 0); 136 137 $for M in range(MR): 138 const __m512i vacc${M}x04152637 = _mm512_add_epi32(_mm512_unpacklo_epi32(vacc${M}x0123, vacc${M}x4567), _mm512_unpackhi_epi32(vacc${M}x0123, vacc${M}x4567)); 139 const __m512i vacc${M}x8C9DAEBF = _mm512_add_epi32(_mm512_unpacklo_epi32(vacc${M}x89AB, vacc${M}xCDEF), _mm512_unpackhi_epi32(vacc${M}x89AB, vacc${M}xCDEF)); 140 141 $for M in range(MR): 142 __m512i vacc${M}x084C195D2A6E3B7F = _mm512_add_epi32(_mm512_unpacklo_epi32(vacc${M}x04152637, vacc${M}x8C9DAEBF), _mm512_unpackhi_epi32(vacc${M}x04152637, vacc${M}x8C9DAEBF)); 143 144 $for M in range(MR): 145 __m512 vscaled${M}x084C195D2A6E3B7F = _mm512_cvtepi32_ps(vacc${M}x084C195D2A6E3B7F); 146 147 $if DATATYPE == "QC8": 148 const __m512 vscale012345678ABCDEF = _mm512_load_ps(w); 149 w = (const void*) ((const float*) w + 16); 150 const __m512 vscale084C195D2A6E3B7F = _mm512_permutexvar_ps(_mm512_set_epi32(15, 7, 11, 3, 14, 6, 10, 2, 13, 5, 9, 1, 12, 4, 8, 0), vscale012345678ABCDEF); 151 $for M in range(MR): 152 vscaled${M}x084C195D2A6E3B7F = _mm512_mul_ps(vscaled${M}x084C195D2A6E3B7F, vscale084C195D2A6E3B7F); 153 $else: 154 $for M in range(MR): 155 vscaled${M}x084C195D2A6E3B7F = _mm512_mul_ps(vscaled${M}x084C195D2A6E3B7F, vscale); 156 157 $for M in range(MR): 158 vscaled${M}x084C195D2A6E3B7F = _mm512_min_ps(vscaled${M}x084C195D2A6E3B7F, voutput_max_less_zero_point); 159 160 $for M in range(MR): 161 vacc${M}x084C195D2A6E3B7F = _mm512_cvtps_epi32(vscaled${M}x084C195D2A6E3B7F); 162 163 $if MR == 1: 164 const __m256i vacc0x084C2A6E195D3B7F = _mm256_adds_epi16(_mm256_packs_epi32(_mm512_castsi512_si256(vacc0x084C195D2A6E3B7F), _mm512_extracti32x8_epi32(vacc0x084C195D2A6E3B7F, 1)), voutput_zero_point); 165 $else: 166 $for M in range(0, MR, 2): 167 const __m512i vacc${M}${min(M+1, MR-1)}x084Cx195Dx2A6Ex3B7F = _mm512_adds_epi16(_mm512_packs_epi32(vacc${M}x084C195D2A6E3B7F, vacc${min(M+1, MR-1)}x084C195D2A6E3B7F), voutput_zero_point); 168 169 $if MR > 2: 170 $if DATATYPE == "QU8": 171 __m512i vout012${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F = _mm512_packus_epi16(vacc01x084Cx195Dx2A6Ex3B7F, vacc2${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F); 172 $else: 173 __m512i vout012${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F = _mm512_packs_epi16(vacc01x084Cx195Dx2A6Ex3B7F, vacc2${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F); 174 vout012${min(M+3, MR-1)}x084Cx195Dx2A6Ex3B7F = _mm512_permutexvar_epi32(_mm512_set_epi32(15, 11, 7, 3, 14, 10, 6, 2, 13, 9, 5, 1, 12, 8, 4, 0), vout012${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F); 175 __m512i vout012${min(3, MR-1)}x0123456789ABCDEF = _mm512_shuffle_epi8(vout012${min(3, MR-1)}x084Cx195Dx2A6Ex3B7F, _mm512_set_epi8(15, 11, 7, 3, 13, 9, 5, 1, 14, 10, 6, 2, 12, 8, 4, 0, 15, 11, 7, 3, 13, 9, 5, 1, 14, 10, 6, 2, 12, 8, 4, 0, 15, 11, 7, 3, 13, 9, 5, 1, 14, 10, 6, 2, 12, 8, 4, 0, 15, 11, 7, 3, 13, 9, 5, 1, 14, 10, 6, 2, 12, 8, 4, 0)); 176 $if DATATYPE == "QU8": 177 vout012${min(3, MR-1)}x0123456789ABCDEF = _mm512_max_epu8(vout012${min(3, MR-1)}x0123456789ABCDEF, voutput_min); 178 $else: 179 vout012${min(3, MR-1)}x0123456789ABCDEF = _mm512_max_epi8(vout012${min(3, MR-1)}x0123456789ABCDEF, voutput_min); 180 $elif MR == 2: 181 $if DATATYPE == "QU8": 182 const __m256i vout01x084Cx2A6Ex195Dx3B7F = _mm256_packus_epi16(_mm512_castsi512_si256(vacc01x084Cx195Dx2A6Ex3B7F), _mm512_extracti32x8_epi32(vacc01x084Cx195Dx2A6Ex3B7F, 1)); 183 $else: 184 const __m256i vout01x084Cx2A6Ex195Dx3B7F = _mm256_packs_epi16(_mm512_castsi512_si256(vacc01x084Cx195Dx2A6Ex3B7F), _mm512_extracti32x8_epi32(vacc01x084Cx195Dx2A6Ex3B7F, 1)); 185 const __m256i vout01x084C2A6E195D3B7F = _mm256_permutevar8x32_epi32(vout01x084Cx2A6Ex195Dx3B7F, _mm256_set_epi32(7, 5, 3, 1, 6, 4, 2, 0)); 186 __m256i vout01x0123456789ABCDEF = _mm256_shuffle_epi8(vout01x084C2A6E195D3B7F, _mm256_set_epi8(15, 7, 11, 3, 13, 5, 9, 1, 14, 6, 10, 2, 12, 4, 8, 0, 15, 7, 11, 3, 13, 5, 9, 1, 14, 6, 10, 2, 12, 4, 8, 0)); 187 $if DATATYPE == "QU8": 188 vout01x0123456789ABCDEF = _mm256_max_epu8(vout01x0123456789ABCDEF, voutput_min); 189 $else: 190 vout01x0123456789ABCDEF = _mm256_max_epi8(vout01x0123456789ABCDEF, voutput_min); 191 $elif MR == 1: 192 $if DATATYPE == "QU8": 193 const __m128i vout0x084C2A6E195D3B7F = _mm_packus_epi16(_mm256_castsi256_si128(vacc0x084C2A6E195D3B7F), _mm256_extracti128_si256(vacc0x084C2A6E195D3B7F, 1)); 194 $else: 195 const __m128i vout0x084C2A6E195D3B7F = _mm_packs_epi16(_mm256_castsi256_si128(vacc0x084C2A6E195D3B7F), _mm256_extracti128_si256(vacc0x084C2A6E195D3B7F, 1)); 196 __m128i vout0x0123456789ABCDEF = _mm_shuffle_epi8(vout0x084C2A6E195D3B7F, _mm_set_epi8(15, 7, 11, 3, 13, 5, 9, 1, 14, 6, 10, 2, 12, 4, 8, 0)); 197 $if DATATYPE == "QU8": 198 vout0x0123456789ABCDEF = _mm_max_epu8(vout0x0123456789ABCDEF, voutput_min); 199 $else: 200 vout0x0123456789ABCDEF = _mm_max_epi8(vout0x0123456789ABCDEF, voutput_min); 201 202 $if MR > 2: 203 if (nc >= 16) { 204 $for M in reversed(range(1, MR)): 205 _mm_storeu_si128((__m128i*) c${M}, _mm512_extracti32x4_epi32(vout012${min(M+3, MR-1)}x0123456789ABCDEF, ${M})); 206 _mm_storeu_si128((__m128i*) c0, _mm512_castsi512_si128(vout012${min(M+3, MR-1)}x0123456789ABCDEF)); 207 208 $for M in reversed(range(MR)): 209 c${M} = (${XINT8_T}*) ((uintptr_t) c${M} + cn_stride); 210 211 a = (const ${XINT8_T}**restrict) ((uintptr_t) a - ks); 212 213 nc -= 16; 214 } else { 215 // Prepare mask for valid 8-bit elements (depends on nc). 216 __mmask64 vmask = _cvtu64_mask64((uint64_t) ((UINT64_C(1) << (nc + ${16 * (MR - 1)})) - (UINT64_C(1) << ${16 * (MR - 1)}))); 217 218 $for M in reversed(range(1, MR)): 219 _mm512_mask_storeu_epi8(c${M} - ${M * 16}, vmask, vout012${min(3, MR-1)}x0123456789ABCDEF); 220 vmask = _kshiftri_mask64(vmask, 16); 221 _mm512_mask_storeu_epi8(c0, vmask, vout012${min(M+3, MR-1)}x0123456789ABCDEF); 222 223 nc = 0; 224 } 225 $elif MR == 2: 226 if (nc >= 16) { 227 _mm_storeu_si128((__m128i*) c1, _mm256_extracti128_si256(vout01x0123456789ABCDEF, 1)); 228 _mm_storeu_si128((__m128i*) c0, _mm256_castsi256_si128(vout01x0123456789ABCDEF)); 229 230 $for M in reversed(range(MR)): 231 c${M} = (${XINT8_T}*) ((uintptr_t) c${M} + cn_stride); 232 233 a = (const ${XINT8_T}**restrict) ((uintptr_t) a - ks); 234 235 nc -= 16; 236 } else { 237 // Prepare mask for valid 8-bit elements (depends on nc). 238 __mmask64 vmask = _cvtu64_mask64((uint64_t) ((UINT32_C(1) << (nc + 16)) - (UINT32_C(1) << 16))); 239 240 _mm256_mask_storeu_epi8(c1 - 16, vmask, vout01x0123456789ABCDEF); 241 vmask = _kshiftri_mask64(vmask, 16); 242 _mm256_mask_storeu_epi8(c0, vmask, vout01x0123456789ABCDEF); 243 244 nc = 0; 245 } 246 $elif MR == 1: 247 if (nc >= 16) { 248 _mm_storeu_si128((__m128i*) c0, vout0x0123456789ABCDEF); 249 250 $for M in range(MR): 251 c${M} = (${XINT8_T}*) ((uintptr_t) c${M} + cn_stride); 252 253 a = (const ${XINT8_T}**restrict) ((uintptr_t) a - ks); 254 255 nc -= 16; 256 } else { 257 // Prepare mask for valid 8-bit elements (depends on nc). 258 const __mmask64 vmask = _cvtu64_mask64((uint64_t) ((UINT32_C(1) << nc) - UINT32_C(1))); 259 260 _mm_mask_storeu_epi8(c0, vmask, vout0x0123456789ABCDEF); 261 262 nc = 0; 263 } 264 } while (nc != 0); 265} 266