xref: /aosp_15_r20/external/XNNPACK/src/qs8-vmul/wasmsimd-mul32-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 BATCH_TILE % 8 == 0
9$assert BATCH_TILE >= 8
10$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
11#include <assert.h>
12
13#include <wasm_simd128.h>
14
15#include <xnnpack/vmul.h>
16
17
18$PARAMS_STRUCT = REQUANTIZATION.lower() + "_wasmsimd"
19$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
20$WASM_X16X8_LOAD8X8 = {"QS8": "wasm_i16x8_load8x8", "QU8": "wasm_u16x8_load8x8"}[DATATYPE]
21$WASM_X32X4_EXTEND_LOW_X16X8 = {"QS8": "wasm_i32x4_extend_low_i16x8", "QU8": "wasm_u32x4_extend_low_u16x8"}[DATATYPE]
22$WASM_X32X4_EXTEND_HIGH_X16X8 = {"QS8": "wasm_i32x4_extend_high_i16x8", "QU8": "wasm_u32x4_extend_high_u16x8"}[DATATYPE]
23$WASM_X8X16_NARROW_I16X8 = {"QS8": "wasm_i8x16_narrow_i16x8", "QU8": "wasm_u8x16_narrow_i16x8"}[DATATYPE]
24$WASM_X8X16_MIN = {"QS8": "wasm_i8x16_min", "QU8": "wasm_u8x16_min"}[DATATYPE]
25void xnn_${DATATYPE.lower()}_vmul_minmax_${REQUANTIZATION.lower()}_ukernel__wasmsimd_mul32_ld64_x${BATCH_TILE}(
26    size_t n,
27    const ${XINT8_T}* input_a,
28    const ${XINT8_T}* input_b,
29    ${XINT8_T}* output,
30    const union xnn_${DATATYPE.lower()}_mul_minmax_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
31
32{
33  const v128_t va_zero_point = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.a_zero_point);
34  const v128_t vb_zero_point = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.b_zero_point);
35  const v128_t vscale = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.scale);
36  const v128_t vmagic_bias = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.magic_bias);
37  const v128_t vmagic_min = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.magic_min);
38  const v128_t vmagic_bias_less_output_zero_point = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.magic_bias_less_output_zero_point);
39  const v128_t voutput_max = wasm_v128_load64_splat(params->${PARAMS_STRUCT}.output_max);
40
41  for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
42    const v128_t va${ABC[0:8]} = ${WASM_X16X8_LOAD8X8}(input_a);
43    const v128_t vb${ABC[0:8]} = ${WASM_X16X8_LOAD8X8}(input_b);
44    $for N in range(8, BATCH_TILE, 8):
45      const v128_t va${ABC[N:N+8]} = ${WASM_X16X8_LOAD8X8}(input_a + ${N});
46      const v128_t vb${ABC[N:N+8]} = ${WASM_X16X8_LOAD8X8}(input_b + ${N});
47    input_a += ${BATCH_TILE};
48    input_b += ${BATCH_TILE};
49
50    $for N in range(0, BATCH_TILE, 8):
51      const v128_t vxa${ABC[N:N+8]} = wasm_i16x8_sub(va${ABC[N:N+8]}, va_zero_point);
52      const v128_t vxb${ABC[N:N+8]} = wasm_i16x8_sub(vb${ABC[N:N+8]}, vb_zero_point);
53
54    $for N in range(0, BATCH_TILE, 8):
55      v128_t vacc${ABC[N:N+4]} = wasm_i32x4_mul(wasm_i32x4_extend_low_i16x8(vxa${ABC[N:N+8]}), wasm_i32x4_extend_low_i16x8(vxb${ABC[N:N+8]}));
56      v128_t vacc${ABC[N+4:N+8]} = wasm_i32x4_mul(wasm_i32x4_extend_high_i16x8(vxa${ABC[N:N+8]}), wasm_i32x4_extend_high_i16x8(vxb${ABC[N:N+8]}));
57
58    $for N in range(0, BATCH_TILE, 4):
59      vacc${ABC[N:N+4]} = wasm_f32x4_convert_i32x4(vacc${ABC[N:N+4]});
60
61    $for N in range(0, BATCH_TILE, 4):
62      vacc${ABC[N:N+4]} = wasm_f32x4_mul(vacc${ABC[N:N+4]}, vscale);
63
64    $for N in range(0, BATCH_TILE, 4):
65      vacc${ABC[N:N+4]} = wasm_f32x4_add(vacc${ABC[N:N+4]}, vmagic_bias);
66
67    $for N in range(0, BATCH_TILE, 4):
68      vacc${ABC[N:N+4]} = wasm_i32x4_max(vacc${ABC[N:N+4]}, vmagic_min);
69
70    $for N in range(0, BATCH_TILE, 4):
71      vacc${ABC[N:N+4]} = wasm_i32x4_sub(vacc${ABC[N:N+4]}, vmagic_bias_less_output_zero_point);
72
73    $for N in range(0, BATCH_TILE, 8):
74      v128_t vout${ABC[N:N+8]} = wasm_i16x8_narrow_i32x4(vacc${ABC[N:N+4]}, vacc${ABC[N+4:N+8]});
75
76    $for N in range(0, BATCH_TILE, 16):
77      $if N + 8 < BATCH_TILE:
78        v128_t vout${ABC[N:N+16]} = ${WASM_X8X16_NARROW_I16X8}(vout${ABC[N:N+8]}, vout${ABC[N+8:N+16]});
79      $else:
80        v128_t vout${ABC[N:N+8]}${ABC[N:N+8]} = ${WASM_X8X16_NARROW_I16X8}(vout${ABC[N:N+8]}, vout${ABC[N:N+8]});
81
82    $for N in range(0, BATCH_TILE, 16):
83      $if N + 8 < BATCH_TILE:
84        vout${ABC[N:N+16]} = ${WASM_X8X16_MIN}(vout${ABC[N:N+16]}, voutput_max);
85      $else:
86        vout${ABC[N:N+8]}${ABC[N:N+8]} = ${WASM_X8X16_MIN}(vout${ABC[N:N+8]}${ABC[N:N+8]}, voutput_max);
87
88    $if BATCH_TILE >= 16:
89      wasm_v128_store(output, vout${ABC[0:16]});
90    $else:
91      *((double*) output) = wasm_f64x2_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
92    $for N in range(16, BATCH_TILE, 16):
93      $if N + 8 < BATCH_TILE:
94        wasm_v128_store(output + ${N}, vout${ABC[N:N+16]});
95      $else:
96        *((double*) output) = wasm_f64x2_extract_lane(output + ${N}, vout${ABC[N:N+8]}${ABC[N:N+8]});
97    output += ${BATCH_TILE};
98  }
99  if XNN_UNLIKELY(n != 0) {
100    ${"do " if BATCH_TILE > 8 else ""}{
101      const v128_t va${ABC[0:8]} = ${WASM_X16X8_LOAD8X8}(input_a);
102      const v128_t vb${ABC[0:8]} = ${WASM_X16X8_LOAD8X8}(input_b);
103      $if BATCH_TILE > 8:
104        input_a += 8;
105        input_b += 8;
106
107      const v128_t vxa${ABC[0:8]} = wasm_i16x8_sub(va${ABC[0:8]}, va_zero_point);
108      const v128_t vxb${ABC[0:8]} = wasm_i16x8_sub(vb${ABC[0:8]}, vb_zero_point);
109
110      v128_t vacc${ABC[0:4]} = wasm_i32x4_mul(wasm_i32x4_extend_low_i16x8(vxa${ABC[0:8]}), wasm_i32x4_extend_low_i16x8(vxb${ABC[0:8]}));
111      v128_t vacc${ABC[4:8]} = wasm_i32x4_mul(wasm_i32x4_extend_high_i16x8(vxa${ABC[0:8]}), wasm_i32x4_extend_high_i16x8(vxb${ABC[0:8]}));
112
113      vacc${ABC[0:4]} = wasm_f32x4_convert_i32x4(vacc${ABC[0:4]});
114      vacc${ABC[4:8]} = wasm_f32x4_convert_i32x4(vacc${ABC[4:8]});
115
116      vacc${ABC[0:4]} = wasm_f32x4_mul(vacc${ABC[0:4]}, vscale);
117      vacc${ABC[4:8]} = wasm_f32x4_mul(vacc${ABC[4:8]}, vscale);
118
119      vacc${ABC[0:4]} = wasm_f32x4_add(vacc${ABC[0:4]}, vmagic_bias);
120      vacc${ABC[4:8]} = wasm_f32x4_add(vacc${ABC[4:8]}, vmagic_bias);
121
122      vacc${ABC[0:4]} = wasm_i32x4_max(vacc${ABC[0:4]}, vmagic_min);
123      vacc${ABC[4:8]} = wasm_i32x4_max(vacc${ABC[4:8]}, vmagic_min);
124
125      vacc${ABC[0:4]} = wasm_i32x4_sub(vacc${ABC[0:4]}, vmagic_bias_less_output_zero_point);
126      vacc${ABC[4:8]} = wasm_i32x4_sub(vacc${ABC[4:8]}, vmagic_bias_less_output_zero_point);
127
128      v128_t vout${ABC[0:8]} = wasm_i16x8_narrow_i32x4(vacc${ABC[0:4]}, vacc${ABC[4:8]});
129      v128_t vout${ABC[0:8]}${ABC[0:8]} = ${WASM_X8X16_NARROW_I16X8}(vout${ABC[0:8]}, vout${ABC[0:8]});
130      vout${ABC[0:8]}${ABC[0:8]} = ${WASM_X8X16_MIN}(vout${ABC[0:8]}${ABC[0:8]}, voutput_max);
131
132      $if BATCH_TILE > 8:
133        if XNN_LIKELY(n >= (8 * sizeof(${XINT8_T}))) {
134          *((double*) output) = wasm_f64x2_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
135          output += 8;
136          n -= 8 * sizeof(${XINT8_T});
137        } else {
138          if (n & (4 * sizeof(${XINT8_T}))) {
139            *((float*) output) = wasm_f32x4_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
140            vout${ABC[0:8]}${ABC[0:8]} = wasm_u64x2_shr(vout${ABC[0:8]}${ABC[0:8]}, 32);
141            output += 4;
142          }
143          uint32_t vout${ABC[0:4]} = wasm_i32x4_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
144          if (n & (2 * sizeof(${XINT8_T}))) {
145            *((uint16_t*) output) = (uint16_t) vout${ABC[0:4]};
146            vout${ABC[0:4]} >>= 16;
147            output += 2;
148          }
149          if (n & (1 * sizeof(${XINT8_T}))) {
150            *output = (${XINT8_T}) vout${ABC[0:4]};
151          }
152          n = 0;
153        }
154      $else:
155        if (n & (4 * sizeof(${XINT8_T}))) {
156          *((float*) output) = wasm_f32x4_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
157          vout${ABC[0:8]}${ABC[0:8]} = wasm_u64x2_shr(vout${ABC[0:8]}${ABC[0:8]}, 32);
158          output += 4;
159        }
160        uint32_t vout${ABC[0:4]} = wasm_i32x4_extract_lane(vout${ABC[0:8]}${ABC[0:8]}, 0);
161        if (n & (2 * sizeof(${XINT8_T}))) {
162          *((uint16_t*) output) = (uint16_t) vout${ABC[0:4]};
163          vout${ABC[0:4]} >>= 16;
164          output += 2;
165        }
166        if (n & (1 * sizeof(${XINT8_T}))) {
167          *output = (${XINT8_T}) vout${ABC[0:4]};
168        }
169    }${" while (n != 0);" if BATCH_TILE > 8 else ""}
170  }
171}
172