xref: /aosp_15_r20/external/XNNPACK/src/qs8-vlrelu/neon.c.in (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)
1// Copyright 2022 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
7$assert BATCH_TILE == 8 or BATCH_TILE % 16 == 0
8$SIMD_TILE = BATCH_TILE // 16
9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
10#include <assert.h>
11
12#include <arm_neon.h>
13
14#include <xnnpack/common.h>
15#include <xnnpack/vlrelu.h>
16
17
18$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
19$XINT8X8_T = {"QS8": "int8x8_t", "QU8": "uint8x8_t"}[DATATYPE]
20$XINT8X16_T = {"QS8": "int8x16_t", "QU8": "uint8x16_t"}[DATATYPE]
21$XINT16X8_T = {"QS8": "int16x8_t", "QU8": "uint16x8_t"}[DATATYPE]
22$VLD1_X8 = {"QS8": "vld1_s8", "QU8": "vld1_u8"}[DATATYPE]
23$VLD1Q_X8 = {"QS8": "vld1q_s8", "QU8": "vld1q_u8"}[DATATYPE]
24$VLD1Q_DUP_X16 = {"QS8": "vld1q_dup_s16", "QU8": "vld1q_dup_u16"}[DATATYPE]
25$VQMOVXN_S16 = {"QS8": "vqmovn_s16", "QU8": "vqmovun_s16"}[DATATYPE]
26$VGET_LOW_X8 = {"QS8": "vget_low_s8", "QU8": "vget_low_u8"}[DATATYPE]
27$VGET_HIGH_X8 = {"QS8": "vget_high_s8", "QU8": "vget_high_u8"}[DATATYPE]
28$VCOMBINE_X8 = {"QS8": "vcombine_s8", "QU8": "vcombine_u8"}[DATATYPE]
29$VREINTERPRET_U32_X8 = {"QS8": "vreinterpret_u32_s8", "QU8": "vreinterpret_u32_u8"}[DATATYPE]
30$VREINTERPRET_U16_X8 = {"QS8": "vreinterpret_u16_s8", "QU8": "vreinterpret_u16_u8"}[DATATYPE]
31$VEXT_X8 = {"QS8": "vext_s8", "QU8": "vext_u8"}[DATATYPE]
32$VST1_X8 = {"QS8": "vst1_s8", "QU8": "vst1_u8"}[DATATYPE]
33$VST1Q_X8 = {"QS8": "vst1q_s8", "QU8": "vst1q_u8"}[DATATYPE]
34$VST1_LANE_X8 = {"QS8": "vst1_lane_s8", "QU8": "vst1_lane_u8"}[DATATYPE]
35void xnn_${DATATYPE.lower()}_vlrelu_ukernel__neon_x${BATCH_TILE}(
36    size_t n,
37    const ${XINT8_T}* x,
38    ${XINT8_T}* y,
39    const union xnn_${DATATYPE.lower()}_lrelu_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
40{
41  assert(n != 0);
42  assert(n % sizeof(${XINT8_T}) == 0);
43  assert(x != NULL);
44  assert(y != NULL);
45
46  const ${XINT16X8_T} vinput_zero_point = ${VLD1Q_DUP_X16}(&params->neon.input_zero_point);
47  const int16x8_t vpositive_multiplier = vld1q_dup_s16(&params->neon.positive_multiplier);
48  const int16x8_t vnegative_multiplier = vld1q_dup_s16(&params->neon.negative_multiplier);
49  const int16x8_t voutput_zero_point = vld1q_dup_s16(&params->neon.output_zero_point);
50  $if BATCH_TILE > 8:
51    for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
52      $for N in range(SIMD_TILE):
53        const ${XINT8X16_T} vx${ABC[N]} = ${VLD1Q_X8}(x); x += 16;
54
55      $for N in range(SIMD_TILE):
56        $if DATATYPE == "QU8":
57          int16x8_t vacc${ABC[2*N]} = vreinterpretq_s16_u16(vsubw_u8(vinput_zero_point, ${VGET_LOW_X8}(vx${ABC[N]})));
58          int16x8_t vacc${ABC[2*N+1]} = vreinterpretq_s16_u16(vsubw_u8(vinput_zero_point, ${VGET_HIGH_X8}(vx${ABC[N]})));
59        $else:
60          int16x8_t vacc${ABC[2*N]} = vsubw_s8(vinput_zero_point, ${VGET_LOW_X8}(vx${ABC[N]}));
61          int16x8_t vacc${ABC[2*N+1]} = vsubw_s8(vinput_zero_point, ${VGET_HIGH_X8}(vx${ABC[N]}));
62
63      $for N in range(2*SIMD_TILE):
64        const uint16x8_t vmask${ABC[N]} = vcltq_s16(vacc${ABC[N]}, vmovq_n_s16(0));
65
66      $for N in range(2*SIMD_TILE):
67        vacc${ABC[N]} = vshlq_n_s16(vacc${ABC[N]}, 7);
68
69      $for N in range(2*SIMD_TILE):
70        const int16x8_t vmultiplier${ABC[N]} = vbslq_s16(vmask${ABC[N]}, vpositive_multiplier, vnegative_multiplier);
71
72      $for N in range(2*SIMD_TILE):
73        vacc${ABC[N]} = vqrdmulhq_s16(vacc${ABC[N]}, vmultiplier${ABC[N]});
74
75      $for N in range(2*SIMD_TILE):
76        vacc${ABC[N]} = vqaddq_s16(vacc${ABC[N]}, voutput_zero_point);
77
78      $for N in range(SIMD_TILE):
79        const ${XINT8X16_T} vy${ABC[N]} = ${VCOMBINE_X8}(${VQMOVXN_S16}(vacc${ABC[2*N]}), ${VQMOVXN_S16}(vacc${ABC[2*N+1]}));
80
81      $for N in range(SIMD_TILE):
82        ${VST1Q_X8}(y, vy${ABC[N]}); y += 16;
83    }
84  for (; n >= 8 * sizeof(${XINT8_T}); n -= 8 * sizeof(${XINT8_T})) {
85    const ${XINT8X8_T} vx = ${VLD1_X8}(x); x += 8;
86    $if DATATYPE == "QU8":
87      int16x8_t vacc = vreinterpretq_s16_u16(vsubw_u8(vinput_zero_point, vx));
88    $else:
89      int16x8_t vacc = vsubw_s8(vinput_zero_point, vx);
90    const uint16x8_t vmask = vcltq_s16(vacc, vmovq_n_s16(0));
91    vacc = vshlq_n_s16(vacc, 7);
92    const int16x8_t vmultiplier = vbslq_s16(vmask, vpositive_multiplier, vnegative_multiplier);
93    vacc = vqrdmulhq_s16(vacc, vmultiplier);
94    vacc = vqaddq_s16(vacc, voutput_zero_point);
95    const ${XINT8X8_T} vy = ${VQMOVXN_S16}(vacc);
96    ${VST1_X8}(y, vy); y += 8;
97  }
98  if XNN_UNLIKELY(n != 0) {
99    assert(n >= 1 * sizeof(${XINT8_T}));
100    assert(n <= 7 * sizeof(${XINT8_T}));
101
102    const ${XINT8X8_T} vx = ${VLD1_X8}(x);
103    $if DATATYPE == "QU8":
104      int16x8_t vacc = vreinterpretq_s16_u16(vsubw_u8(vinput_zero_point, vx));
105    $else:
106      int16x8_t vacc = vsubw_s8(vinput_zero_point, vx);
107    const uint16x8_t vmask = vcltq_s16(vacc, vmovq_n_s16(0));
108    vacc = vshlq_n_s16(vacc, 7);
109    const int16x8_t vmultiplier = vbslq_s16(vmask, vpositive_multiplier, vnegative_multiplier);
110    vacc = vqrdmulhq_s16(vacc, vmultiplier);
111    vacc = vqaddq_s16(vacc, voutput_zero_point);
112    ${XINT8X8_T} vy = ${VQMOVXN_S16}(vacc);
113
114    if (n & (4 * sizeof(${XINT8_T}))) {
115      vst1_lane_u32((void*) y, ${VREINTERPRET_U32_X8}(vy), 0); y += 4;
116      vy = ${VEXT_X8}(vy, vy, 4);
117    }
118    if (n & (2 * sizeof(${XINT8_T}))) {
119      vst1_lane_u16((void*) y, ${VREINTERPRET_U16_X8}(vy), 0); y += 2;
120      vy = ${VEXT_X8}(vy, vy, 2);
121    }
122    if (n & (1 * sizeof(${XINT8_T}))) {
123      ${VST1_LANE_X8}(y, vy, 0);
124    }
125  }
126}
127