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$SSE_HEADER = {2: "emmintrin.h", 3: "tmmintrin.h", 4: "smmintrin.h", 5: "ammintrin.h"}[SSE] 7$assert VARIANT in ["LD64", "LD128", "EXTENDED"] 8$assert MR <= 4 9#include <assert.h> 10 11$if SSE == 5: 12 #ifdef __GNUC__ 13 #include <x86intrin.h> 14 #else 15 #include <immintrin.h> 16 #include <${SSE_HEADER}> 17 #endif 18$else: 19 #include <${SSE_HEADER}> 20 21#include <xnnpack/gemm.h> 22#include <xnnpack/math.h> 23 24 25$LOAD_SUFFIX = {"LD128": "_ld128", "LD64": "_ld64", "EXTENDED": ""}[VARIANT] 26$GEMM_SUFFIX = "_xw" if VARIANT == "EXTENDED" else "" 27$ISA = {2: "sse2", 3: "ssse3", 4: "sse41", 5: "xop"}[SSE] 28void xnn_qs8_gemm${GEMM_SUFFIX}_minmax_ukernel_${MR}x4c8__${ISA}${LOAD_SUFFIX}( 29 size_t mr, 30 size_t nc, 31 size_t kc, 32 const int8_t* restrict a, 33 size_t a_stride, 34 const void* restrict w, 35 int8_t* restrict c, 36 size_t cm_stride, 37 size_t cn_stride, 38 const union xnn_qs8_gemm${GEMM_SUFFIX}_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_DISABLE_TSAN 39{ 40 assert(mr != 0); 41 assert(mr <= ${MR}); 42 assert(nc != 0); 43 assert(kc != 0); 44 assert(kc % sizeof(int8_t) == 0); 45 assert(a != NULL); 46 assert(w != NULL); 47 assert(c != NULL); 48 49 kc = round_up_po2(kc, 8); 50 const int8_t* a0 = a; 51 int8_t* c0 = c; 52 $for M in range(1, MR): 53 const int8_t* a${M} = (const int8_t*) ((uintptr_t) a${M-1} + a_stride); 54 int8_t* c${M} = (int8_t*) ((uintptr_t) c${M-1} + cm_stride); 55 $if M % 2 == 0: 56 if XNN_UNPREDICTABLE(mr <= ${M}) { 57 a${M} = a${M-1}; 58 c${M} = c${M-1}; 59 } 60 $elif M + 1 == MR: 61 if XNN_UNPREDICTABLE(mr != ${M+1}) { 62 a${M} = a${M-1}; 63 c${M} = c${M-1}; 64 } 65 $else: 66 if XNN_UNPREDICTABLE(mr < ${M+1}) { 67 a${M} = a${M-1}; 68 c${M} = c${M-1}; 69 } 70 71 do { 72 $for N in range(4): 73 __m128i vacc0x${N} = _mm_cvtsi32_si128((int) ((const int32_t*) w)[${N}]); 74 $for M in range(1, MR): 75 $for N in range(4): 76 __m128i vacc${M}x${N} = vacc0x${N}; 77 w = (const void*) ((uintptr_t) w + 4 * sizeof(int32_t)); 78 79 size_t k = 0; 80 while (k < kc) { 81 $for M in range(MR): 82 const __m128i va${M} = _mm_loadl_epi64((const __m128i*) a${M}); 83 $if SSE >= 4: 84 const __m128i vxa${M} = _mm_cvtepi8_epi16(va${M}); 85 $else: 86 const __m128i vxa${M} = _mm_unpacklo_epi8(va${M}, _mm_cmpgt_epi8(_mm_setzero_si128(), va${M})); 87 a${M} += 8; 88 89 $if VARIANT == "LD128": 90 $for N in range(0, 4, 2): 91 $if N == 0: 92 const __m128i vb${N}${N+1} = _mm_load_si128((const __m128i*) w); 93 $else: 94 const __m128i vb${N}${N+1} = _mm_load_si128((const __m128i*) ((uintptr_t) w + ${N * 8} * sizeof(int8_t))); 95 const __m128i vsb${N}${N+1} = _mm_cmpgt_epi8(_mm_setzero_si128(), vb${N}${N+1}); 96 const __m128i vxb${N} = _mm_unpacklo_epi8(vb${N}${N+1}, vsb${N}${N+1}); 97 const __m128i vxb${N+1} = _mm_unpackhi_epi8(vb${N}${N+1}, vsb${N}${N+1}); 98 99 $for M in range(MR): 100 $if SSE == 5: 101 vacc${M}x${N} = _mm_maddd_epi16(vxa${M}, vxb${N}, vacc${M}x${N}); 102 vacc${M}x${N+1} = _mm_maddd_epi16(vxa${M}, vxb${N+1}, vacc${M}x${N+1}); 103 $else: 104 vacc${M}x${N} = _mm_add_epi32(vacc${M}x${N}, _mm_madd_epi16(vxa${M}, vxb${N})); 105 vacc${M}x${N+1} = _mm_add_epi32(vacc${M}x${N+1}, _mm_madd_epi16(vxa${M}, vxb${N+1})); 106 $else: 107 $for N in range(4): 108 $if VARIANT == "LD64": 109 $if N == 0: 110 const __m128i vb${N} = _mm_loadl_epi64((const __m128i*) w); 111 $else: 112 const __m128i vb${N} = _mm_loadl_epi64((const __m128i*) ((uintptr_t) w + ${N * 8} * sizeof(int8_t))); 113 $if SSE >= 4: 114 const __m128i vxb${N} = _mm_cvtepi8_epi16(vb${N}); 115 $else: 116 const __m128i vxb${N} = _mm_unpacklo_epi8(vb${N}, _mm_cmpgt_epi8(_mm_setzero_si128(), vb${N})); 117 $elif VARIANT == "EXTENDED": 118 $if N == 0: 119 const __m128i vxb${N} = _mm_load_si128((const __m128i*) w); 120 $else: 121 const __m128i vxb${N} = _mm_load_si128((const __m128i*) ((uintptr_t) w + ${N * 8} * sizeof(int16_t))); 122 123 $for M in range(MR): 124 $if SSE == 5: 125 vacc${M}x${N} = _mm_maddd_epi16(vxa${M}, vxb${N}, vacc${M}x${N}); 126 $else: 127 vacc${M}x${N} = _mm_add_epi32(vacc${M}x${N}, _mm_madd_epi16(vxa${M}, vxb${N})); 128 129 $if VARIANT == "EXTENDED": 130 w = (const void*) ((uintptr_t) w + 32 * sizeof(int16_t)); 131 $else: 132 w = (const void*) ((uintptr_t) w + 32 * sizeof(int8_t)); 133 k += 8 * sizeof(int8_t); 134 } 135 136 $if SSE >= 3: 137 $for M in range(MR): 138 const __m128i vacc${M}x01 = _mm_hadd_epi32(vacc${M}x0, vacc${M}x1); 139 const __m128i vacc${M}x23 = _mm_hadd_epi32(vacc${M}x2, vacc${M}x3); 140 141 $for M in range(MR): 142 __m128i vacc${M}x0123 = _mm_hadd_epi32(vacc${M}x01, vacc${M}x23); 143 $else: 144 $for M in range(MR): 145 const __m128i vacc${M}x02 = _mm_add_epi32(_mm_unpacklo_epi32(vacc${M}x0, vacc${M}x2), _mm_unpackhi_epi32(vacc${M}x0, vacc${M}x2)); 146 const __m128i vacc${M}x13 = _mm_add_epi32(_mm_unpacklo_epi32(vacc${M}x1, vacc${M}x3), _mm_unpackhi_epi32(vacc${M}x1, vacc${M}x3)); 147 148 $for M in range(MR): 149 __m128i vacc${M}x0123 = _mm_add_epi32(_mm_unpacklo_epi32(vacc${M}x02, vacc${M}x13), _mm_unpackhi_epi32(vacc${M}x02, vacc${M}x13)); 150 151 const __m128i vmultiplier = _mm_load_si128((const __m128i*) params->sse2.multiplier); 152 const __m128i vrounding = _mm_load_si128((const __m128i*) params->sse2.rounding); 153 154 $if SSE >= 4: 155 $for M in range(MR): 156 const __m128i vacc${M}x1133 = _mm_shuffle_epi32(vacc${M}x0123, _MM_SHUFFLE(3, 3, 1, 1)); 157 158 $for M in range(MR): 159 const __m128i vprod${M}x02 = _mm_add_epi64(_mm_mul_epi32(vacc${M}x0123, vmultiplier), vrounding); 160 161 $for M in range(MR): 162 const __m128i vprod${M}x13 = _mm_add_epi64(_mm_mul_epi32(vacc${M}x1133, vmultiplier), vrounding); 163 164 $for M in range(MR): 165 const __m128i vq31prod${M}x02 = _mm_srli_epi64(vprod${M}x02, 31); 166 const __m128i vq31prod${M}x13 = _mm_add_epi64(vprod${M}x13, vprod${M}x13); 167 168 $for M in range(MR): 169 const __m128i vq31prod${M}x0123 = _mm_blend_epi16(vq31prod${M}x02, vq31prod${M}x13, 0xCC); 170 $else: 171 $for M in range(MR): 172 const __m128i vnmask${M}x0123 = _mm_cmpgt_epi32(_mm_setzero_si128(), vacc${M}x0123); 173 174 $for M in range(MR): 175 $if SSE >= 3: 176 const __m128i vabsacc${M}x0123 = _mm_abs_epi32(vacc${M}x0123); 177 $else: 178 const __m128i vabsacc${M}x0123 = _mm_sub_epi32(_mm_xor_si128(vacc${M}x0123, vnmask${M}x0123), vnmask${M}x0123); 179 180 $for M in range(MR): 181 const __m128i vabsacc${M}x1133 = _mm_shuffle_epi32(vabsacc${M}x0123, _MM_SHUFFLE(3, 3, 1, 1)); 182 183 $for M in range(MR): 184 const __m128i vabsprod${M}x02 = _mm_mul_epu32(vabsacc${M}x0123, vmultiplier); 185 186 $for M in range(MR): 187 const __m128i vnmask${M}x02 = _mm_shuffle_epi32(vnmask${M}x0123, _MM_SHUFFLE(2, 2, 0, 0)); 188 189 $for M in range(MR): 190 const __m128i vprod${M}x02 = _mm_sub_epi64(_mm_xor_si128(vabsprod${M}x02, vnmask${M}x02), vnmask${M}x02); 191 192 $for M in range(MR): 193 const __m128i vq31prod${M}x02 = _mm_srli_epi64(_mm_add_epi64(vprod${M}x02, vrounding), 31); 194 195 $for M in range(MR): 196 const __m128i vabsprod${M}x13 = _mm_mul_epu32(vabsacc${M}x1133, vmultiplier); 197 198 $for M in range(MR): 199 const __m128i vnmask${M}x13 = _mm_shuffle_epi32(vnmask${M}x0123, _MM_SHUFFLE(3, 3, 1, 1)); 200 201 $for M in range(MR): 202 const __m128i vprod${M}x13 = _mm_sub_epi64(_mm_xor_si128(vabsprod${M}x13, vnmask${M}x13), vnmask${M}x13); 203 204 $for M in range(MR): 205 const __m128i vq31prod${M}x13 = _mm_srli_epi64(_mm_add_epi64(vprod${M}x13, vrounding), 31); 206 207 $for M in range(MR): 208 const __m128i vq31prod${M}x0213 = _mm_castps_si128(_mm_shuffle_ps( 209 _mm_castsi128_ps(vq31prod${M}x02), _mm_castsi128_ps(vq31prod${M}x13), _MM_SHUFFLE(2, 0, 2, 0))); 210 211 $for M in range(MR): 212 const __m128i vq31prod${M}x0123 = _mm_shuffle_epi32(vq31prod${M}x0213, _MM_SHUFFLE(3, 1, 2, 0)); 213 214 const __m128i vremainder_mask = _mm_load_si128((const __m128i*) params->sse2.remainder_mask); 215 $for M in range(MR): 216 const __m128i vrem${M}x0123 = 217 _mm_add_epi32(_mm_and_si128(vq31prod${M}x0123, vremainder_mask), _mm_cmpgt_epi32(_mm_setzero_si128(), vq31prod${M}x0123)); 218 219 const __m128i vremainder_threshold = _mm_load_si128((const __m128i*) params->sse2.remainder_threshold); 220 const __m128i vshift = _mm_load_si128((const __m128i*) params->sse2.shift); 221 $for M in range(MR): 222 vacc${M}x0123 = 223 _mm_sub_epi32(_mm_sra_epi32(vq31prod${M}x0123, vshift), _mm_cmpgt_epi32(vrem${M}x0123, vremainder_threshold)); 224 225 const __m128i voutput_zero_point = _mm_load_si128((const __m128i*) params->sse2.output_zero_point); 226 $for M in range(0, MR, 2): 227 __m128i vacc${M}${min(M+1, MR-1)}x0123 = _mm_adds_epi16(_mm_packs_epi32(vacc${M}x0123, vacc${min(M+1, MR-1)}x0123), voutput_zero_point); 228 229 const __m128i voutput_min = _mm_load_si128((const __m128i*) params->sse2.output_min); 230 const __m128i voutput_max = _mm_load_si128((const __m128i*) params->sse2.output_max); 231 $for M in range(0, MR, 2): 232 vacc${M}${min(M+1, MR-1)}x0123 = _mm_min_epi16(_mm_max_epi16(vacc${M}${min(M+1, MR-1)}x0123, voutput_min), voutput_max); 233 234 $if MR > 2: 235 __m128i vout = _mm_packs_epi16(vacc0${min(1, MR-1)}x0123, vacc${min(2, MR-1)}${min(3, MR-1)}x0123); 236 $else: 237 __m128i vout = _mm_packs_epi16(vacc0${min(1, MR-1)}x0123, vacc0${min(1, MR-1)}x0123); 238 239 if (nc >= 4) { 240 *((uint32_t*) c0) = (uint32_t) _mm_cvtsi128_si32(vout); 241 $for M in range(1, MR): 242 $if SSE >= 4: 243 *((uint32_t*) c${M}) = (uint32_t) _mm_extract_epi32(vout, ${M}); 244 $else: 245 vout = _mm_srli_si128(vout, 4); 246 *((uint32_t*) c${M}) = (uint32_t) _mm_cvtsi128_si32(vout); 247 248 $for M in range(MR): 249 c${M} = (int8_t*) ((uintptr_t) c${M} + cn_stride); 250 251 $for M in range(MR): 252 a${M} = (const int8_t*) ((uintptr_t) a${M} - kc); 253 254 nc -= 4; 255 } else { 256 if (nc & 2) { 257 $for M in range(MR): 258 *((uint16_t*) c${M}) = (uint16_t) _mm_extract_epi16(vout, ${M * 2}); 259 c${M} += 2; 260 vout = _mm_srli_epi32(vout, 16); 261 } 262 if (nc & 1) { 263 $if SSE >= 4: 264 $for M in range(MR): 265 *((int8_t*) c${M}) = (int8_t) _mm_extract_epi8(vout, ${M * 4}); 266 $else: 267 *((int8_t*) c0) = (int8_t) _mm_cvtsi128_si32(vout); 268 $for M in range(1, MR): 269 *((int8_t*) c${M}) = (int8_t) _mm_extract_epi16(vout, ${M * 2}); 270 } 271 272 nc = 0; 273 } 274 } while (nc != 0); 275} 276