xref: /aosp_15_r20/external/XNNPACK/src/qs8-vmul/sse-mul16-ld64.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36) 
1// Copyright 2021 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 DATATYPE in ["QS8", "QU8"]
7$assert REQUANTIZATION == "FP32"
8$assert SSE in [2, 4]
9$assert not AVX or SSE == 4
10$SSE_HEADER = {2: "emmintrin.h", 4: "smmintrin.h"}[SSE]
11$assert BATCH_TILE % 8 == 0
12$assert BATCH_TILE >= 8
13$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
14#include <assert.h>
15
16#include <${SSE_HEADER}>
17
18#include <xnnpack/unaligned.h>
19#include <xnnpack/vmul.h>
20
21
22$PARAMS_STRUCT = REQUANTIZATION.lower() + "_" + ("sse4" if SSE == 4 and DATATYPE == "QS8" else "sse2")
23$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
24$_MM_CVTEPX8_EPI16 = {"QS8": "_mm_cvtepi8_epi16", "QU8": "_mm_cvtepu8_epi16"}[DATATYPE]
25$_MM_PACKXS_EPI16 = {"QS8": "_mm_packs_epi16", "QU8": "_mm_packus_epi16"}[DATATYPE]
26$_MM_MIN_EPX8 = {"QS8": "_mm_min_epi8", "QU8": "_mm_min_epu8"}[DATATYPE]
27$_MM_MAX_EPX8 = {"QS8": "_mm_max_epi8", "QU8": "_mm_max_epu8"}[DATATYPE]
28$ISA = "avx" if AVX else {2: "sse2", 4: "sse41"}[SSE]
29void xnn_${DATATYPE.lower()}_vmul_minmax_${REQUANTIZATION.lower()}_ukernel__${ISA}_mul16_ld64_x${BATCH_TILE}(
30    size_t n,
31    const ${XINT8_T}* input_a,
32    const ${XINT8_T}* input_b,
33    ${XINT8_T}* output,
34    const union xnn_${DATATYPE.lower()}_mul_minmax_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
35
36{
37  const __m128i va_zero_point = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.a_zero_point);
38  const __m128i vb_zero_point = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.b_zero_point);
39  const __m128 vscale = _mm_load_ps(params->${PARAMS_STRUCT}.scale);
40  const __m128i voutput_zero_point = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_zero_point);
41  const __m128i voutput_min = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_min);
42  const __m128i voutput_max = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_max);
43
44  for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
45    $if SSE == 4:
46      const __m128i va${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_a));
47      const __m128i vb${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_b));
48      $for N in range(8, BATCH_TILE, 8):
49        const __m128i va${ABC[N:N+8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) (input_a + ${N})));
50        const __m128i vb${ABC[N:N+8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) (input_b + ${N})));
51    $else:
52      __m128i va${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_a);
53      __m128i vb${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_b);
54      $for N in range(8, BATCH_TILE, 8):
55        __m128i va${ABC[N:N+8]} = _mm_loadl_epi64((const __m128i*) (input_a + ${N}));
56        __m128i vb${ABC[N:N+8]} = _mm_loadl_epi64((const __m128i*) (input_b + ${N}));
57    input_a += ${BATCH_TILE};
58    input_b += ${BATCH_TILE};
59
60    $if SSE < 4:
61      $if DATATYPE == "QU8":
62        const __m128i vzero = _mm_setzero_si128();
63        $for N in range(0, BATCH_TILE, 8):
64          va${ABC[N:N+8]} = _mm_unpacklo_epi8(va${ABC[N:N+8]}, vzero);
65          vb${ABC[N:N+8]} = _mm_unpacklo_epi8(vb${ABC[N:N+8]}, vzero);
66      $else:
67        $for N in range(0, BATCH_TILE, 8):
68          va${ABC[N:N+8]} = _mm_srai_epi16(_mm_unpacklo_epi8(va${ABC[N:N+8]}, va${ABC[N:N+8]}), 8);
69          vb${ABC[N:N+8]} = _mm_srai_epi16(_mm_unpacklo_epi8(vb${ABC[N:N+8]}, vb${ABC[N:N+8]}), 8);
70
71    $for N in range(0, BATCH_TILE, 8):
72      const __m128i vxa${ABC[N:N+8]} = _mm_sub_epi16(va${ABC[N:N+8]}, va_zero_point);
73      const __m128i vxb${ABC[N:N+8]} = _mm_sub_epi16(vb${ABC[N:N+8]}, vb_zero_point);
74
75    $for N in range(0, BATCH_TILE, 8):
76      const __m128i vprod${ABC[N:N+8]}lo = _mm_mullo_epi16(vxa${ABC[N:N+8]}, vxb${ABC[N:N+8]});
77      const __m128i vprod${ABC[N:N+8]}hi = _mm_mulhi_epi16(vxa${ABC[N:N+8]}, vxb${ABC[N:N+8]});
78
79    $for N in range(0, BATCH_TILE, 8):
80      const __m128i vprod${ABC[N:N+4]} = _mm_unpacklo_epi16(vprod${ABC[N:N+8]}lo, vprod${ABC[N:N+8]}hi);
81      const __m128i vprod${ABC[N+4:N+8]} = _mm_unpackhi_epi16(vprod${ABC[N:N+8]}lo, vprod${ABC[N:N+8]}hi);
82
83    $for N in range(0, BATCH_TILE, 4):
84      __m128 vfpacc${ABC[N:N+4]} = _mm_cvtepi32_ps(vprod${ABC[N:N+4]});
85
86    $for N in range(0, BATCH_TILE, 4):
87      vfpacc${ABC[N:N+4]} = _mm_mul_ps(vfpacc${ABC[N:N+4]}, vscale);
88
89    $for N in range(0, BATCH_TILE, 4):
90      const __m128i vacc${ABC[N:N+4]} = _mm_cvtps_epi32(vfpacc${ABC[N:N+4]});
91
92    $for N in range(0, BATCH_TILE, 8):
93      __m128i vout${ABC[N:N+8]} = _mm_adds_epi16(_mm_packs_epi32(vacc${ABC[N:N+4]}, vacc${ABC[N+4:N+8]}), voutput_zero_point);
94
95    $if DATATYPE == "QS8" and SSE < 4:
96      $for N in range(0, BATCH_TILE, 8):
97        vout${ABC[N:N+8]} = _mm_max_epi16(vout${ABC[N:N+8]}, voutput_min);
98
99      $for N in range(0, BATCH_TILE, 8):
100        vout${ABC[N:N+8]} = _mm_min_epi16(vout${ABC[N:N+8]}, voutput_max);
101
102    $for N in range(0, BATCH_TILE, 16):
103      $if N + 8 < BATCH_TILE:
104        __m128i vout${ABC[N:N+16]} = ${_MM_PACKXS_EPI16}(vout${ABC[N:N+8]}, vout${ABC[N+8:N+16]});
105      $else:
106        __m128i vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_PACKXS_EPI16}(vout${ABC[N:N+8]}, vout${ABC[N:N+8]});
107
108    $if DATATYPE == "QU8" or SSE == 4:
109      $for N in range(0, BATCH_TILE, 16):
110        $if N + 8 < BATCH_TILE:
111          vout${ABC[N:N+16]} = ${_MM_MAX_EPX8}(vout${ABC[N:N+16]}, voutput_min);
112        $else:
113          vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_MAX_EPX8}(vout${ABC[N:N+8]}${ABC[N:N+8]}, voutput_min);
114
115      $for N in range(0, BATCH_TILE, 16):
116        $if N + 8 < BATCH_TILE:
117          vout${ABC[N:N+16]} = ${_MM_MIN_EPX8}(vout${ABC[N:N+16]}, voutput_max);
118        $else:
119          vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_MIN_EPX8}(vout${ABC[N:N+8]}${ABC[N:N+8]}, voutput_max);
120
121    $if BATCH_TILE >= 16:
122      _mm_storeu_si128((__m128i*) output, vout${ABC[0:16]});
123    $else:
124      _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
125    $for N in range(16, BATCH_TILE, 16):
126      $if N + 8 < BATCH_TILE:
127        _mm_storeu_si128((__m128i*) (output + ${N}), vout${ABC[N:N+16]});
128      $else:
129        _mm_storel_epi64((__m128i*) (output + ${N}), vout${ABC[N:N+8]}${ABC[N:N+8]});
130    output += ${BATCH_TILE};
131  }
132  if XNN_UNLIKELY(n != 0) {
133    ${"do " if BATCH_TILE > 8 else ""}{
134      $if SSE == 4:
135        const __m128i va${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_a));
136        const __m128i vb${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_b));
137      $else:
138        __m128i va${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_a);
139        __m128i vb${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_b);
140      $if BATCH_TILE > 8:
141        input_a += 8;
142        input_b += 8;
143
144      $if SSE < 4:
145        $if DATATYPE == "QU8":
146          const __m128i vzero = _mm_setzero_si128();
147          va${ABC[0:8]} = _mm_unpacklo_epi8(va${ABC[0:8]}, vzero);
148          vb${ABC[0:8]} = _mm_unpacklo_epi8(vb${ABC[0:8]}, vzero);
149        $else:
150          va${ABC[0:8]} = _mm_srai_epi16(_mm_unpacklo_epi8(va${ABC[0:8]}, va${ABC[0:8]}), 8);
151          vb${ABC[0:8]} = _mm_srai_epi16(_mm_unpacklo_epi8(vb${ABC[0:8]}, vb${ABC[0:8]}), 8);
152
153      const __m128i vxa${ABC[0:8]} = _mm_sub_epi16(va${ABC[0:8]}, va_zero_point);
154      const __m128i vxb${ABC[0:8]} = _mm_sub_epi16(vb${ABC[0:8]}, vb_zero_point);
155
156      const __m128i vprod${ABC[0:8]}lo = _mm_mullo_epi16(vxa${ABC[0:8]}, vxb${ABC[0:8]});
157      const __m128i vprod${ABC[0:8]}hi = _mm_mulhi_epi16(vxa${ABC[0:8]}, vxb${ABC[0:8]});
158
159      const __m128i vprod${ABC[0:4]} = _mm_unpacklo_epi16(vprod${ABC[0:8]}lo, vprod${ABC[0:8]}hi);
160      const __m128i vprod${ABC[4:8]} = _mm_unpackhi_epi16(vprod${ABC[0:8]}lo, vprod${ABC[0:8]}hi);
161
162      __m128 vfpacc${ABC[0:4]} = _mm_cvtepi32_ps(vprod${ABC[0:4]});
163      __m128 vfpacc${ABC[4:8]} = _mm_cvtepi32_ps(vprod${ABC[4:8]});
164
165      vfpacc${ABC[0:4]} = _mm_mul_ps(vfpacc${ABC[0:4]}, vscale);
166      vfpacc${ABC[4:8]} = _mm_mul_ps(vfpacc${ABC[4:8]}, vscale);
167
168      const __m128i vacc${ABC[0:4]} = _mm_cvtps_epi32(vfpacc${ABC[0:4]});
169      const __m128i vacc${ABC[4:8]} = _mm_cvtps_epi32(vfpacc${ABC[4:8]});
170
171      __m128i vout${ABC[0:8]} = _mm_adds_epi16(_mm_packs_epi32(vacc${ABC[0:4]}, vacc${ABC[4:8]}), voutput_zero_point);
172      $if DATATYPE == "QS8" and SSE < 4:
173        vout${ABC[0:8]} = _mm_max_epi16(vout${ABC[0:8]}, voutput_min);
174        vout${ABC[0:8]} = _mm_min_epi16(vout${ABC[0:8]}, voutput_max);
175
176      __m128i vout${ABC[0:8]}${ABC[0:8]} = ${_MM_PACKXS_EPI16}(vout${ABC[0:8]}, vout${ABC[0:8]});
177      $if DATATYPE == "QU8" or SSE == 4:
178        vout${ABC[0:8]}${ABC[0:8]} = ${_MM_MAX_EPX8}(vout${ABC[0:8]}${ABC[0:8]}, voutput_min);
179        vout${ABC[0:8]}${ABC[0:8]} = ${_MM_MIN_EPX8}(vout${ABC[0:8]}${ABC[0:8]}, voutput_max);
180
181      $if BATCH_TILE > 8:
182        if XNN_LIKELY(n >= (8 * sizeof(${XINT8_T}))) {
183          _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
184          output += 8;
185          n -= 8 * sizeof(${XINT8_T});
186        } else {
187          if (n & (4 * sizeof(${XINT8_T}))) {
188            unaligned_store_u32(output, (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
189            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
190            output += 4;
191          }
192          if (n & (2 * sizeof(${XINT8_T}))) {
193            $if SSE == 4:
194              unaligned_store_u16(output, (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0));
195            $else:
196              unaligned_store_u16(output, (uint16_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
197            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
198            output += 2;
199          }
200          if (n & (1 * sizeof(${XINT8_T}))) {
201            $if SSE == 4:
202              *output = (${XINT8_T}) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
203            $else:
204              *output = (int32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
205          }
206          n = 0;
207        }
208      $else:
209        if (n & (4 * sizeof(${XINT8_T}))) {
210          unaligned_store_u32(output, (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
211          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
212          output += 4;
213        }
214        if (n & (2 * sizeof(${XINT8_T}))) {
215          $if SSE == 4:
216            unaligned_store_u16(output, (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0));
217          $else:
218            unaligned_store_u16(output, (uint16_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
219          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
220          output += 2;
221        }
222        if (n & (1 * sizeof(${XINT8_T}))) {
223          $if SSE == 4:
224            *output = (${XINT8_T}) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
225          $else:
226            *output = (${XINT8_T}) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
227        }
228    }${" while (n != 0);" if BATCH_TILE > 8 else ""}
229  }
230}
231