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$assert BATCH_TILE % 8 == 0
7$assert BATCH_TILE >= 8
8$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
9#include <assert.h>
10
11#include <immintrin.h>
12
13#include <xnnpack/intrinsics-polyfill.h>
14#include <xnnpack/vadd.h>
15
16
17void xnn_qs8_vaddc_minmax_ukernel__avx2_mul32_ld64_x${BATCH_TILE}(
18    size_t n,
19    const int8_t* input_x,
20    const int8_t* input_y,
21    int8_t* output,
22    const union xnn_qs8_add_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_DISABLE_TSAN
23{
24  const __m256i vx_multiplier = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.x_multiplier));
25  const __m256i vremainder_mask = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.remainder_mask));
26  const __m256i vremainder_threshold = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.remainder_threshold));
27  const __m128i vshift = _mm_cvtsi32_si128((int) params->sse2.shift);
28  $if BATCH_TILE > 8:
29    const __m256i voutput_zero_point = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.output_zero_point));
30    const __m256i voutput_min = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.output_min));
31    const __m256i voutput_max = _mm256_broadcastsi128_si256(_mm_load_si128((const __m128i*) params->sse2.output_max));
32  $else:
33    const __m128i voutput_zero_point = _mm_load_si128((const __m128i*) params->sse2.output_zero_point);
34    const __m128i voutput_min = _mm_load_si128((const __m128i*) params->sse2.output_min);
35    const __m128i voutput_max = _mm_load_si128((const __m128i*) params->sse2.output_max);
36
37  __m256i vzero_point_product = _mm256_broadcastsi128_si256(_mm_add_epi32(
38    _mm_broadcastd_epi32(_mm_cvtsi32_si128(params->sse2.y_multiplier[0] * (int32_t) *input_y)),
39    _mm_load_si128((const __m128i*) params->sse2.zero_point_product)));
40  for (; n >= ${BATCH_TILE} * sizeof(int8_t); n -= ${BATCH_TILE} * sizeof(int8_t)) {
41    const __m256i vx${ABC[0:8]} = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i*) input_x));
42    $for N in range(8, BATCH_TILE, 8):
43      const __m256i vx${ABC[N:N+8]} = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i*) (input_x + ${N})));
44    input_x += ${BATCH_TILE};
45
46    $for N in range(0, BATCH_TILE, 8):
47      __m256i vacc${ABC[N:N+8]} = _mm256_add_epi32(vzero_point_product, _mm256_mullo_epi32(vx${ABC[N:N+8]}, vx_multiplier));
48
49    $for N in range(0, BATCH_TILE, 8):
50      const __m256i vrem${ABC[N:N+8]} = _mm256_add_epi32(_mm256_and_si256(vacc${ABC[N:N+8]}, vremainder_mask), _mm256_srai_epi32(vacc${ABC[N:N+8]}, 31));
51
52    $for N in range(0, BATCH_TILE, 8):
53      vacc${ABC[N:N+8]} = _mm256_sub_epi32(_mm256_sra_epi32(vacc${ABC[N:N+8]}, vshift), _mm256_cmpgt_epi32(vrem${ABC[N:N+8]}, vremainder_threshold));
54
55    $for N in range(0, BATCH_TILE, 16):
56      $if N + 8 < BATCH_TILE:
57        __m256i vout${ABC[N:N+4]}${ABC[N+8:N+12]}${ABC[N+4:N+8]}${ABC[N+12:N+16]} = _mm256_adds_epi16(_mm256_packs_epi32(vacc${ABC[N:N+8]}, vacc${ABC[N+8:N+16]}), voutput_zero_point);
58      $elif BATCH_TILE > 8:
59        __m128i vout${ABC[N:N+8]} = _mm_adds_epi16(_mm_packs_epi32(_mm256_castsi256_si128(vacc${ABC[N:N+8]}), _mm256_extracti128_si256(vacc${ABC[N:N+8]}, 1)), _mm256_castsi256_si128(voutput_zero_point));
60      $else:
61        __m128i vout${ABC[N:N+8]} = _mm_adds_epi16(_mm_packs_epi32(_mm256_castsi256_si128(vacc${ABC[N:N+8]}), _mm256_extracti128_si256(vacc${ABC[N:N+8]}, 1)), voutput_zero_point);
62
63    $for N in range(0, BATCH_TILE, 16):
64      $if N + 8 < BATCH_TILE:
65        vout${ABC[N:N+4]}${ABC[N+8:N+12]}${ABC[N+4:N+8]}${ABC[N+12:N+16]} = _mm256_min_epi16(_mm256_max_epi16(vout${ABC[N:N+4]}${ABC[N+8:N+12]}${ABC[N+4:N+8]}${ABC[N+12:N+16]}, voutput_min), voutput_max);
66      $elif BATCH_TILE > 8:
67        vout${ABC[N:N+8]} = _mm_min_epi16(_mm_max_epi16(vout${ABC[N:N+8]}, _mm256_castsi256_si128(voutput_min)), _mm256_castsi256_si128(voutput_max));
68      $else:
69        vout${ABC[N:N+8]} = _mm_min_epi16(_mm_max_epi16(vout${ABC[N:N+8]}, voutput_min), voutput_max);
70
71    $for N in range(0, BATCH_TILE, 16):
72      $if N + 8 < BATCH_TILE:
73        __m128i vout${ABC[N:N+16]} = _mm_shuffle_epi32(_mm_packs_epi16(_mm256_castsi256_si128(vout${ABC[N:N+4]}${ABC[N+8:N+12]}${ABC[N+4:N+8]}${ABC[N+12:N+16]}), _mm256_extracti128_si256(vout${ABC[N:N+4]}${ABC[N+8:N+12]}${ABC[N+4:N+8]}${ABC[N+12:N+16]}, 1)), _MM_SHUFFLE(3, 1, 2, 0));
74      $else:
75        __m128i vout${ABC[N:N+8]}${ABC[N:N+8]} = _mm_packs_epi16(vout${ABC[N:N+8]}, vout${ABC[N:N+8]});
76
77    $if BATCH_TILE >= 16:
78      _mm_storeu_si128((__m128i*) output, vout${ABC[0:16]});
79    $else:
80      _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
81    $for N in range(16, BATCH_TILE, 16):
82      $if N + 8 < BATCH_TILE:
83        _mm_storeu_si128((__m128i*) (output + ${N}), vout${ABC[N:N+16]});
84      $else:
85        _mm_storel_epi64((__m128i*) (output + ${N}), vout${ABC[N:N+8]}${ABC[N:N+8]});
86    output += ${BATCH_TILE};
87  }
88  if XNN_UNLIKELY(n != 0) {
89    ${"do " if BATCH_TILE > 8 else ""}{
90      const __m256i vx${ABC[0:8]} = _mm256_cvtepi8_epi32(_mm_loadl_epi64((const __m128i*) input_x));
91      $if BATCH_TILE > 8:
92        input_x += 8;
93
94      __m256i vacc${ABC[0:8]} = _mm256_add_epi32(vzero_point_product, _mm256_mullo_epi32(vx${ABC[0:8]}, vx_multiplier));
95
96      const __m256i vrem${ABC[0:8]} = _mm256_add_epi32(_mm256_and_si256(vacc${ABC[0:8]}, vremainder_mask), _mm256_srai_epi32(vacc${ABC[0:8]}, 31));
97
98      vacc${ABC[0:8]} = _mm256_sub_epi32(_mm256_sra_epi32(vacc${ABC[0:8]}, vshift), _mm256_cmpgt_epi32(vrem${ABC[0:8]}, vremainder_threshold));
99
100      $if BATCH_TILE > 8:
101        __m128i vout${ABC[0:8]} = _mm_adds_epi16(_mm_packs_epi32(_mm256_castsi256_si128(vacc${ABC[0:8]}), _mm256_extracti128_si256(vacc${ABC[0:8]}, 1)), _mm256_castsi256_si128(voutput_zero_point));
102        vout${ABC[0:8]} = _mm_min_epi16(_mm_max_epi16(vout${ABC[0:8]}, _mm256_castsi256_si128(voutput_min)), _mm256_castsi256_si128(voutput_max));
103      $else:
104        __m128i vout${ABC[0:8]} = _mm_adds_epi16(_mm_packs_epi32(_mm256_castsi256_si128(vacc${ABC[0:8]}), _mm256_extracti128_si256(vacc${ABC[0:8]}, 1)), voutput_zero_point);
105        vout${ABC[0:8]} = _mm_min_epi16(_mm_max_epi16(vout${ABC[0:8]}, voutput_min), voutput_max);
106      __m128i vout${ABC[0:8]}${ABC[0:8]} = _mm_packs_epi16(vout${ABC[0:8]}, vout${ABC[0:8]});
107
108      $if BATCH_TILE > 8:
109        if XNN_LIKELY(n >= (8 * sizeof(int8_t))) {
110          _mm_storel_epi64((__m128i*) output, vout${ABC[0:8]}${ABC[0:8]});
111          output += 8;
112          n -= 8 * sizeof(int8_t);
113        } else {
114          if (n & (4 * sizeof(int8_t))) {
115            *((uint32_t*) output) = (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
116            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
117            output += 4;
118          }
119          if (n & (2 * sizeof(int8_t))) {
120            *((uint16_t*) output) = (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0);
121            vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
122            output += 2;
123          }
124          if (n & (1 * sizeof(int8_t))) {
125            *output = (int8_t) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
126          }
127          n = 0;
128        }
129      $else:
130        if (n & (4 * sizeof(int8_t))) {
131          *((uint32_t*) output) = (uint32_t) _mm_cvtsi128_si32(vout${ABC[0:8]}${ABC[0:8]});
132          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi64(vout${ABC[0:8]}${ABC[0:8]}, 32);
133          output += 4;
134        }
135        if (n & (2 * sizeof(int8_t))) {
136          *((uint16_t*) output) = (uint16_t) _mm_extract_epi16(vout${ABC[0:8]}${ABC[0:8]}, 0);
137          vout${ABC[0:8]}${ABC[0:8]} = _mm_srli_epi32(vout${ABC[0:8]}${ABC[0:8]}, 16);
138          output += 2;
139        }
140        if (n & (1 * sizeof(int8_t))) {
141          *output = (int8_t) _mm_extract_epi8(vout${ABC[0:8]}${ABC[0:8]}, 0);
142        }
143    }${" while (n != 0);" if BATCH_TILE > 8 else ""}
144  }
145}
146