xref: /aosp_15_r20/external/XNNPACK/src/qs8-vmulc/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()}_vmulc_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 __m128 vscale = _mm_load_ps(params->${PARAMS_STRUCT}.scale);
39  const __m128i voutput_zero_point = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_zero_point);
40  const __m128i voutput_min = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_min);
41  const __m128i voutput_max = _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.output_max);
42
43  __m128i vxb = _mm_sub_epi16(
44    _mm_shuffle_epi32(_mm_cvtsi32_si128(UINT32_C(0x00010001) * (uint32_t) (uint16_t) (int16_t) *input_b), 0),
45    _mm_load_si128((const __m128i*) params->${PARAMS_STRUCT}.b_zero_point));
46  for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
47    $if SSE == 4:
48      const __m128i va${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_a));
49      $for N in range(8, BATCH_TILE, 8):
50        const __m128i va${ABC[N:N+8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) (input_a + ${N})));
51    $else:
52      __m128i va${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_a);
53      $for N in range(8, BATCH_TILE, 8):
54        __m128i va${ABC[N:N+8]} = _mm_loadl_epi64((const __m128i*) (input_a + ${N}));
55    input_a += ${BATCH_TILE};
56
57    $if SSE < 4:
58      $if DATATYPE == "QU8":
59        const __m128i vzero = _mm_setzero_si128();
60        $for N in range(0, BATCH_TILE, 8):
61          va${ABC[N:N+8]} = _mm_unpacklo_epi8(va${ABC[N:N+8]}, vzero);
62      $else:
63        $for N in range(0, BATCH_TILE, 8):
64          va${ABC[N:N+8]} = _mm_srai_epi16(_mm_unpacklo_epi8(va${ABC[N:N+8]}, va${ABC[N:N+8]}), 8);
65
66    $for N in range(0, BATCH_TILE, 8):
67      const __m128i vxa${ABC[N:N+8]} = _mm_sub_epi16(va${ABC[N:N+8]}, va_zero_point);
68
69    $for N in range(0, BATCH_TILE, 8):
70      const __m128i vprod${ABC[N:N+8]}lo = _mm_mullo_epi16(vxa${ABC[N:N+8]}, vxb);
71      const __m128i vprod${ABC[N:N+8]}hi = _mm_mulhi_epi16(vxa${ABC[N:N+8]}, vxb);
72
73    $for N in range(0, BATCH_TILE, 8):
74      const __m128i vprod${ABC[N:N+4]} = _mm_unpacklo_epi16(vprod${ABC[N:N+8]}lo, vprod${ABC[N:N+8]}hi);
75      const __m128i vprod${ABC[N+4:N+8]} = _mm_unpackhi_epi16(vprod${ABC[N:N+8]}lo, vprod${ABC[N:N+8]}hi);
76
77    $for N in range(0, BATCH_TILE, 4):
78      __m128 vfpacc${ABC[N:N+4]} = _mm_cvtepi32_ps(vprod${ABC[N:N+4]});
79
80    $for N in range(0, BATCH_TILE, 4):
81      vfpacc${ABC[N:N+4]} = _mm_mul_ps(vfpacc${ABC[N:N+4]}, vscale);
82
83    $for N in range(0, BATCH_TILE, 4):
84      const __m128i vacc${ABC[N:N+4]} = _mm_cvtps_epi32(vfpacc${ABC[N:N+4]});
85
86    $for N in range(0, BATCH_TILE, 8):
87      __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);
88
89    $if DATATYPE == "QS8" and SSE < 4:
90      $for N in range(0, BATCH_TILE, 8):
91        vout${ABC[N:N+8]} = _mm_max_epi16(vout${ABC[N:N+8]}, voutput_min);
92
93      $for N in range(0, BATCH_TILE, 8):
94        vout${ABC[N:N+8]} = _mm_min_epi16(vout${ABC[N:N+8]}, voutput_max);
95
96    $for N in range(0, BATCH_TILE, 16):
97      $if N + 8 < BATCH_TILE:
98        __m128i vout${ABC[N:N+16]} = ${_MM_PACKXS_EPI16}(vout${ABC[N:N+8]}, vout${ABC[N+8:N+16]});
99      $else:
100        __m128i vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_PACKXS_EPI16}(vout${ABC[N:N+8]}, vout${ABC[N:N+8]});
101
102    $if DATATYPE == "QU8" or SSE == 4:
103      $for N in range(0, BATCH_TILE, 16):
104        $if N + 8 < BATCH_TILE:
105          vout${ABC[N:N+16]} = ${_MM_MAX_EPX8}(vout${ABC[N:N+16]}, voutput_min);
106        $else:
107          vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_MAX_EPX8}(vout${ABC[N:N+8]}${ABC[N:N+8]}, voutput_min);
108
109      $for N in range(0, BATCH_TILE, 16):
110        $if N + 8 < BATCH_TILE:
111          vout${ABC[N:N+16]} = ${_MM_MIN_EPX8}(vout${ABC[N:N+16]}, voutput_max);
112        $else:
113          vout${ABC[N:N+8]}${ABC[N:N+8]} = ${_MM_MIN_EPX8}(vout${ABC[N:N+8]}${ABC[N:N+8]}, voutput_max);
114
115    $if BATCH_TILE >= 16:
116      _mm_storeu_si128((__m128i*) output, vout${ABC[0:16]});
117    $else:
118      _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
119    $for N in range(16, BATCH_TILE, 16):
120      $if N + 8 < BATCH_TILE:
121        _mm_storeu_si128((__m128i*) (output + ${N}), vout${ABC[N:N+16]});
122      $else:
123        _mm_storel_epi64((__m128i*) (output + ${N}), vout${ABC[N:N+8]}${ABC[N:N+8]});
124    output += ${BATCH_TILE};
125  }
126  if XNN_UNLIKELY(n != 0) {
127    ${"do " if BATCH_TILE > 8 else ""}{
128      $if SSE == 4:
129        const __m128i va${ABC[0:8]} = ${_MM_CVTEPX8_EPI16}(_mm_loadl_epi64((const __m128i*) input_a));
130      $else:
131        __m128i va${ABC[0:8]} = _mm_loadl_epi64((const __m128i*) input_a);
132      $if BATCH_TILE > 8:
133        input_a += 8;
134
135      $if SSE < 4:
136        $if DATATYPE == "QU8":
137          const __m128i vzero = _mm_setzero_si128();
138          va${ABC[0:8]} = _mm_unpacklo_epi8(va${ABC[0:8]}, vzero);
139        $else:
140          va${ABC[0:8]} = _mm_srai_epi16(_mm_unpacklo_epi8(va${ABC[0:8]}, va${ABC[0:8]}), 8);
141
142      const __m128i vxa${ABC[0:8]} = _mm_sub_epi16(va${ABC[0:8]}, va_zero_point);
143
144      const __m128i vprod${ABC[0:8]}lo = _mm_mullo_epi16(vxa${ABC[0:8]}, vxb);
145      const __m128i vprod${ABC[0:8]}hi = _mm_mulhi_epi16(vxa${ABC[0:8]}, vxb);
146
147      const __m128i vprod${ABC[0:4]} = _mm_unpacklo_epi16(vprod${ABC[0:8]}lo, vprod${ABC[0:8]}hi);
148      const __m128i vprod${ABC[4:8]} = _mm_unpackhi_epi16(vprod${ABC[0:8]}lo, vprod${ABC[0:8]}hi);
149
150      __m128 vfpacc${ABC[0:4]} = _mm_cvtepi32_ps(vprod${ABC[0:4]});
151      __m128 vfpacc${ABC[4:8]} = _mm_cvtepi32_ps(vprod${ABC[4:8]});
152
153      vfpacc${ABC[0:4]} = _mm_mul_ps(vfpacc${ABC[0:4]}, vscale);
154      vfpacc${ABC[4:8]} = _mm_mul_ps(vfpacc${ABC[4:8]}, vscale);
155
156      const __m128i vacc${ABC[0:4]} = _mm_cvtps_epi32(vfpacc${ABC[0:4]});
157      const __m128i vacc${ABC[4:8]} = _mm_cvtps_epi32(vfpacc${ABC[4:8]});
158
159      __m128i vout${ABC[0:8]} = _mm_adds_epi16(_mm_packs_epi32(vacc${ABC[0:4]}, vacc${ABC[4:8]}), voutput_zero_point);
160      $if DATATYPE == "QS8" and SSE < 4:
161        vout${ABC[0:8]} = _mm_max_epi16(vout${ABC[0:8]}, voutput_min);
162        vout${ABC[0:8]} = _mm_min_epi16(vout${ABC[0:8]}, voutput_max);
163
164      __m128i vout${ABC[0:8]}${ABC[0:8]} = ${_MM_PACKXS_EPI16}(vout${ABC[0:8]}, vout${ABC[0:8]});
165      $if DATATYPE == "QU8" or SSE == 4:
166        vout${ABC[0:8]}${ABC[0:8]} = ${_MM_MAX_EPX8}(vout${ABC[0:8]}${ABC[0:8]}, voutput_min);
167        vout${ABC[0:8]}${ABC[0:8]} = ${_MM_MIN_EPX8}(vout${ABC[0:8]}${ABC[0:8]}, voutput_max);
168
169      $if BATCH_TILE > 8:
170        if XNN_LIKELY(n >= (8 * sizeof(${XINT8_T}))) {
171          _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
172          output += 8;
173          n -= 8 * sizeof(${XINT8_T});
174        } else {
175          if (n & (4 * sizeof(${XINT8_T}))) {
176            unaligned_store_u32(output, (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
177            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
178            output += 4;
179          }
180          if (n & (2 * sizeof(${XINT8_T}))) {
181            $if SSE == 4:
182              unaligned_store_u16(output, (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0));
183            $else:
184              unaligned_store_u16(output, (uint16_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
185            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
186            output += 2;
187          }
188          if (n & (1 * sizeof(${XINT8_T}))) {
189            $if SSE == 4:
190              *output = (${XINT8_T}) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
191            $else:
192              *output = (int32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
193          }
194          n = 0;
195        }
196      $else:
197        if (n & (4 * sizeof(${XINT8_T}))) {
198          unaligned_store_u32(output, (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
199          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
200          output += 4;
201        }
202        if (n & (2 * sizeof(${XINT8_T}))) {
203          $if SSE == 4:
204            unaligned_store_u16(output, (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0));
205          $else:
206            unaligned_store_u16(output, (uint16_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]}));
207          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
208          output += 2;
209        }
210        if (n & (1 * sizeof(${XINT8_T}))) {
211          $if SSE == 4:
212            *output = (${XINT8_T}) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
213          $else:
214            *output = (${XINT8_T}) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
215        }
216    }${" while (n != 0);" if BATCH_TILE > 8 else ""}
217  }
218}
219