xref: /aosp_15_r20/external/XNNPACK/src/qs8-vlrelu/armsimd32.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 DATATYPE in ["QS8", "QU8"]
7$assert BATCH_TILE % 4 == 0
8$SIMD_TILE = BATCH_TILE // 4
9$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
10#include <assert.h>
11
12#include <arm_acle.h>
13
14#include <xnnpack/intrinsics-polyfill.h>
15#include <xnnpack/math.h>
16#include <xnnpack/unaligned.h>
17#include <xnnpack/vlrelu.h>
18
19
20$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
21$XINT8X4_T = {"QS8": "int8x4_t", "QU8": "uint8x4_t"}[DATATYPE]
22$XINT16X2_T = {"QS8": "int16x2_t", "QU8": "uint16x2_t"}[DATATYPE]
23$__XXTB16 = {"QS8": "__sxtb16", "QU8": "__uxtb16"}[DATATYPE]
24$__XSUB16 = {"QS8": "__ssub16", "QU8": "__usub16"}[DATATYPE]
25$__XSAT = {"QS8": "__ssat", "QU8": "__usat"}[DATATYPE]
26void xnn_${DATATYPE.lower()}_vlrelu_ukernel__armsimd32_x${BATCH_TILE}(
27    size_t n,
28    const ${XINT8_T}* x,
29    ${XINT8_T}* y,
30    const union xnn_${DATATYPE.lower()}_lrelu_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
31{
32  const ${XINT16X2_T} vinput_zero_point = (${XINT16X2_T}) params->armsimd32.input_zero_point;
33  const int16x2_t vpositive_multiplier = (int16x2_t) params->armsimd32.positive_multiplier;
34  const int16x2_t vnegative_multiplier = (int16x2_t) params->armsimd32.negative_multiplier;
35  const int32_t vbias = params->armsimd32.bias;
36  $if BATCH_TILE > 4:
37    for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
38      $for N in range(SIMD_TILE):
39        const ${XINT8X4_T} vx${ABC[4*N:4*N+4]} = (${XINT8X4_T}) unaligned_indexed_load_u32(x, ${N});
40      x += ${BATCH_TILE};
41
42      $for N in range(0, BATCH_TILE, 4):
43        ${XINT16X2_T} vx${ABC[N]}${ABC[N+2]} = ${__XXTB16}(vx${ABC[N:N+4]});
44        ${XINT16X2_T} vx${ABC[N+1]}${ABC[N+3]} = ${__XXTB16}(__ror(vx${ABC[N:N+4]}, 8));
45
46      $for N in range(0, BATCH_TILE, 4):
47        vx${ABC[N]}${ABC[N+2]} = ${__XSUB16}(vinput_zero_point, vx${ABC[N]}${ABC[N+2]});
48        const int16x2_t vmultiplier${ABC[N]}${ABC[N+2]} = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
49        vx${ABC[N+1]}${ABC[N+3]} = ${__XSUB16}(vinput_zero_point, vx${ABC[N+1]}${ABC[N+3]});
50        const int16x2_t vmultiplier${ABC[N+1]}${ABC[N+3]} = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
51
52      $for N in range(0, BATCH_TILE, 4):
53        int32_t vacc${ABC[N]} = __smlabb(vmultiplier${ABC[N]}${ABC[N+2]}, vx${ABC[N]}${ABC[N+2]}, vbias);
54        int32_t vacc${ABC[N+1]} = __smlabb(vmultiplier${ABC[N+1]}${ABC[N+3]}, vx${ABC[N+1]}${ABC[N+3]}, vbias);
55        int32_t vacc${ABC[N+2]} = __smlatt(vmultiplier${ABC[N]}${ABC[N+2]}, vx${ABC[N]}${ABC[N+2]}, vbias);
56        int32_t vacc${ABC[N+3]} = __smlatt(vmultiplier${ABC[N+1]}${ABC[N+3]}, vx${ABC[N+1]}${ABC[N+3]}, vbias);
57
58      $for N in range(BATCH_TILE):
59        vacc${ABC[N]} = ${__XSAT}(math_asr_s32(vacc${ABC[N]}, 8), 8);
60
61      $for N in range(BATCH_TILE):
62        y[${N}] = (${XINT8_T}) vacc${ABC[N]};
63      y += ${BATCH_TILE};
64    }
65  for (; n >= 4 * sizeof(${XINT8_T}); n -= 4 * sizeof(${XINT8_T})) {
66    const ${XINT8X4_T} vx0123 = (${XINT8X4_T}) unaligned_load_u32(x);
67    x += 4;
68
69    ${XINT16X2_T} vx02 = ${__XXTB16}(vx0123);
70    ${XINT16X2_T} vx13 = ${__XXTB16}(__ror(vx0123, 8));
71
72    vx02 = ${__XSUB16}(vinput_zero_point, vx02);
73    const int16x2_t vmultiplier02 = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
74    vx13 = ${__XSUB16}(vinput_zero_point, vx13);
75    const int16x2_t vmultiplier13 = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
76
77    int32_t vacc0 = __smlabb(vmultiplier02, vx02, vbias);
78    int32_t vacc1 = __smlabb(vmultiplier13, vx13, vbias);
79    int32_t vacc2 = __smlatt(vmultiplier02, vx02, vbias);
80    int32_t vacc3 = __smlatt(vmultiplier13, vx13, vbias);
81
82    vacc0 = ${__XSAT}(math_asr_s32(vacc0, 8), 8);
83    vacc1 = ${__XSAT}(math_asr_s32(vacc1, 8), 8);
84    vacc2 = ${__XSAT}(math_asr_s32(vacc2, 8), 8);
85    vacc3 = ${__XSAT}(math_asr_s32(vacc3, 8), 8);
86
87    y[0] = (${XINT8_T}) vacc0;
88    y[1] = (${XINT8_T}) vacc1;
89    y[2] = (${XINT8_T}) vacc2;
90    y[3] = (${XINT8_T}) vacc3;
91    y += 4;
92  }
93  if XNN_UNLIKELY(n != 0) {
94    const ${XINT8X4_T} vx0123 = (${XINT8X4_T}) unaligned_load_u32(x);
95
96    ${XINT16X2_T} vx02 = ${__XXTB16}(vx0123);
97    ${XINT16X2_T} vx13 = ${__XXTB16}(__ror(vx0123, 8));
98
99    vx02 = ${__XSUB16}(vinput_zero_point, vx02);
100    const int16x2_t vmultiplier02 = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
101    vx13 = ${__XSUB16}(vinput_zero_point, vx13);
102    const int16x2_t vmultiplier13 = (int16x2_t) __sel((uint8x4_t) vnegative_multiplier, (uint8x4_t) vpositive_multiplier);
103
104    int32_t vacc0 = __smlabb(vmultiplier02, vx02, vbias);
105    int32_t vacc1 = __smlabb(vmultiplier13, vx13, vbias);
106    const int32_t vacc2 = __smlatt(vmultiplier02, vx02, vbias);
107
108    vacc0 = ${__XSAT}(math_asr_s32(vacc0, 8), 8);
109    vacc1 = ${__XSAT}(math_asr_s32(vacc1, 8), 8);
110
111    if (n & (2 * sizeof(${XINT8_T}))) {
112      y[0] = (${XINT8_T}) vacc0;
113      y[1] = (${XINT8_T}) vacc1;
114      vacc0 = ${__XSAT}(math_asr_s32(vacc2, 8), 8);
115      y += 2;
116    }
117    if (n & (1 * sizeof(${XINT8_T}))) {
118      y[0] = (${XINT8_T}) vacc0;
119    }
120  }
121}
122