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 <arm_neon.h>
12
13#include <xnnpack/vadd.h>
14
15
16void xnn_qs8_vadd_minmax_ukernel__neon_ld64_x${BATCH_TILE}(
17    size_t n,
18    const int8_t* input_x,
19    const int8_t* input_y,
20    int8_t* output,
21    const union xnn_qs8_add_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_DISABLE_TSAN
22{
23  const int8x8_t vx_zero_point = vld1_dup_s8(&params->neon.x_zero_point);
24  const int8x8_t vy_zero_point = vld1_dup_s8(&params->neon.y_zero_point);
25  const int32x4_t vx_multiplier = vld1q_dup_s32(&params->neon.x_multiplier);
26  const int32x4_t vy_multiplier = vld1q_dup_s32(&params->neon.y_multiplier);
27  const int32x4_t vright_shift = vld1q_dup_s32(&params->neon.right_shift);
28  const int32x4_t vzero_shift_mask = vreinterpretq_s32_u32(vceqq_s32(vright_shift, vmovq_n_s32(0)));
29  const int16x8_t voutput_zero_point = vld1q_dup_s16(&params->neon.output_zero_point);
30  const int8x16_t voutput_min = vld1q_dup_s8(&params->neon.output_min);
31  const int8x16_t voutput_max = vld1q_dup_s8(&params->neon.output_max);
32
33  for (; n >= ${BATCH_TILE} * sizeof(int8_t); n -= ${BATCH_TILE} * sizeof(int8_t)) {
34    $for N in range(0, BATCH_TILE, 8):
35      const int8x8_t vx${ABC[N:N+8]} = vld1_s8(input_x); input_x += 8;
36      const int8x8_t vy${ABC[N:N+8]} = vld1_s8(input_y); input_y += 8;
37
38    $for N in range(0, BATCH_TILE, 8):
39      const int16x8_t vex${ABC[N:N+8]} = vsubl_s8(vx${ABC[N:N+8]}, vx_zero_point);
40      const int16x8_t vey${ABC[N:N+8]} = vsubl_s8(vy${ABC[N:N+8]}, vy_zero_point);
41
42    $for N in range(0, BATCH_TILE, 8):
43      int32x4_t vacc${ABC[N:N+4]} = vmulq_s32(vmovl_s16(vget_low_s16(vex${ABC[N:N+8]})), vx_multiplier);
44      int32x4_t vacc${ABC[N+4:N+8]} = vmulq_s32(vmovl_s16(vget_high_s16(vex${ABC[N:N+8]})), vx_multiplier);
45
46    $for N in range(0, BATCH_TILE, 8):
47      vacc${ABC[N:N+4]} = vmlaq_s32(vacc${ABC[N:N+4]}, vmovl_s16(vget_low_s16(vey${ABC[N:N+8]})), vy_multiplier);
48      vacc${ABC[N+4:N+8]} = vmlaq_s32(vacc${ABC[N+4:N+8]}, vmovl_s16(vget_high_s16(vey${ABC[N:N+8]})), vy_multiplier);
49
50    $for N in range(0, BATCH_TILE, 4):
51      vacc${ABC[N:N+4]} = vsraq_n_s32(vacc${ABC[N:N+4]}, vbicq_s32(vacc${ABC[N:N+4]}, vzero_shift_mask), 31);
52
53    $for N in range(0, BATCH_TILE, 4):
54      vacc${ABC[N:N+4]} = vrshlq_s32(vacc${ABC[N:N+4]}, vright_shift);
55
56    $for N in range(0, BATCH_TILE, 8):
57      const int16x8_t vacc${ABC[N:N+8]} = vqaddq_s16(vcombine_s16(vqmovn_s32(vacc${ABC[N:N+4]}), vqmovn_s32(vacc${ABC[N+4:N+8]})), voutput_zero_point);
58
59    $for N in range(0, BATCH_TILE, 16):
60      $if N + 8 < BATCH_TILE:
61        int8x16_t vout${ABC[N:N+16]} = vcombine_s8(vqmovn_s16(vacc${ABC[N:N+8]}), vqmovn_s16(vacc${ABC[N+8:N+16]}));
62      $else:
63        int8x8_t vout${ABC[N:N+8]} = vqmovn_s16(vacc${ABC[N:N+8]});
64
65    $for N in range(0, BATCH_TILE, 16):
66      $if N + 8 < BATCH_TILE:
67        vout${ABC[N:N+16]} = vmaxq_s8(vout${ABC[N:N+16]}, voutput_min);
68      $else:
69        vout${ABC[N:N+8]} = vmax_s8(vout${ABC[N:N+8]}, vget_low_s8(voutput_min));
70
71    $for N in range(0, BATCH_TILE, 16):
72      $if N + 8 < BATCH_TILE:
73        vout${ABC[N:N+16]} = vminq_s8(vout${ABC[N:N+16]}, voutput_max);
74      $else:
75        vout${ABC[N:N+8]} = vmin_s8(vout${ABC[N:N+8]}, vget_low_s8(voutput_max));
76
77    $for N in range(0, BATCH_TILE, 16):
78      $if N + 8 < BATCH_TILE:
79        vst1q_s8(output, vout${ABC[N:N+16]}); output += 16;
80      $else:
81        vst1_s8(output, vout${ABC[N:N+8]}); output += 8;
82  }
83  if XNN_UNLIKELY(n != 0) {
84    ${"do " if BATCH_TILE > 8 else ""}{
85      $if BATCH_TILE > 8:
86        const int8x8_t vx${ABC[0:8]} = vld1_s8(input_x); input_x += 8;
87        const int8x8_t vy${ABC[0:8]} = vld1_s8(input_y); input_y += 8;
88      $else:
89        const int8x8_t vx${ABC[0:8]} = vld1_s8(input_x);
90        const int8x8_t vy${ABC[0:8]} = vld1_s8(input_y);
91
92      const int16x8_t vex${ABC[0:8]} = vsubl_s8(vx${ABC[0:8]}, vx_zero_point);
93      const int16x8_t vey${ABC[0:8]} = vsubl_s8(vy${ABC[0:8]}, vy_zero_point);
94
95      int32x4_t vacc${ABC[0:4]} = vmulq_s32(vmovl_s16(vget_low_s16(vex${ABC[0:8]})), vx_multiplier);
96      int32x4_t vacc${ABC[4:8]} = vmulq_s32(vmovl_s16(vget_high_s16(vex${ABC[0:8]})), vx_multiplier);
97
98      vacc${ABC[0:4]} = vmlaq_s32(vacc${ABC[0:4]}, vmovl_s16(vget_low_s16(vey${ABC[0:8]})), vy_multiplier);
99      vacc${ABC[4:8]} = vmlaq_s32(vacc${ABC[4:8]}, vmovl_s16(vget_high_s16(vey${ABC[0:8]})), vy_multiplier);
100
101      vacc${ABC[0:4]} = vsraq_n_s32(vacc${ABC[0:4]}, vbicq_s32(vacc${ABC[0:4]}, vzero_shift_mask), 31);
102      vacc${ABC[4:8]} = vsraq_n_s32(vacc${ABC[4:8]}, vbicq_s32(vacc${ABC[4:8]}, vzero_shift_mask), 31);
103
104      vacc${ABC[0:4]} = vrshlq_s32(vacc${ABC[0:4]}, vright_shift);
105      vacc${ABC[4:8]} = vrshlq_s32(vacc${ABC[4:8]}, vright_shift);
106
107      const int16x8_t vacc${ABC[0:8]} = vqaddq_s16(vcombine_s16(vqmovn_s32(vacc${ABC[0:4]}), vqmovn_s32(vacc${ABC[4:8]})), voutput_zero_point);
108
109      int8x8_t vout${ABC[0:8]} = vqmovn_s16(vacc${ABC[0:8]});
110      vout${ABC[0:8]} = vmax_s8(vout${ABC[0:8]}, vget_low_s8(voutput_min));
111      vout${ABC[0:8]} = vmin_s8(vout${ABC[0:8]}, vget_low_s8(voutput_max));
112
113      $if BATCH_TILE > 8:
114        if XNN_LIKELY(n >= (8 * sizeof(int8_t))) {
115          vst1_s8(output, vout${ABC[0:8]}); output += 8;
116          n -= 8 * sizeof(int8_t);
117        } else {
118          if (n & (4 * sizeof(int8_t))) {
119            vst1_lane_u32(__builtin_assume_aligned(output, 1), vreinterpret_u32_s8(vout${ABC[0:8]}), 0); output += 4;
120            vout${ABC[0:8]} = vext_s8(vout${ABC[0:8]}, vout${ABC[0:8]}, 4);
121          }
122          if (n & (2 * sizeof(int8_t))) {
123            vst1_lane_u16(__builtin_assume_aligned(output, 1), vreinterpret_u16_s8(vout${ABC[0:8]}), 0); output += 2;
124            vout${ABC[0:8]} = vext_s8(vout${ABC[0:8]}, vout${ABC[0:8]}, 2);
125          }
126          if (n & (1 * sizeof(int8_t))) {
127            vst1_lane_s8(output, vout${ABC[0:8]}, 0);
128          }
129          n = 0;
130        }
131      $else:
132        if (n & (4 * sizeof(int8_t))) {
133          vst1_lane_u32(__builtin_assume_aligned(output, 1), vreinterpret_u32_s8(vout${ABC[0:8]}), 0); output += 4;
134          vout${ABC[0:8]} = vext_s8(vout${ABC[0:8]}, vout${ABC[0:8]}, 4);
135        }
136        if (n & (2 * sizeof(int8_t))) {
137          vst1_lane_u16(__builtin_assume_aligned(output, 1), vreinterpret_u16_s8(vout${ABC[0:8]}), 0); output += 2;
138          vout${ABC[0:8]} = vext_s8(vout${ABC[0:8]}, vout${ABC[0:8]}, 2);
139        }
140        if (n & (1 * sizeof(int8_t))) {
141          vst1_lane_s8(output, vout${ABC[0:8]}, 0);
142        }
143    }${" while (n != 0);" if BATCH_TILE > 8 else ""}
144  }
145}
146