1 // Copyright (c) Facebook, Inc. and its affiliates.
2 // All rights reserved.
3 //
4 // Copyright 2019 Google LLC
5 //
6 // This source code is licensed under the BSD-style license found in the
7 // LICENSE file in the root directory of this source tree.
8
9 #include <assert.h>
10 #include <stdint.h>
11
12 #include <arm_neon.h>
13
14 #include <fp16/bitcasts.h>
15
16 #include <xnnpack/requantization-stubs.h>
17
18
xnn_requantize_q31__neon(size_t n,const int32_t * input,float scale,uint8_t zero_point,uint8_t qmin,uint8_t qmax,uint8_t * output)19 void xnn_requantize_q31__neon(
20 size_t n,
21 const int32_t* input,
22 float scale,
23 uint8_t zero_point,
24 uint8_t qmin,
25 uint8_t qmax,
26 uint8_t* output)
27 {
28 assert(n % 16 == 0);
29 assert(scale < 1.0f);
30 assert(scale >= 0x1.0p-32f);
31
32 // Compute requantization parameters.
33 const uint32_t scale_bits = fp32_to_bits(scale);
34
35 // Multiplier is in [0x40000000, 0x7FFFFF80] range.
36 const int32_t multiplier = (int32_t)(((scale_bits & UINT32_C(0x007FFFFF)) | UINT32_C(0x00800000)) << 7);
37 assert(multiplier >= INT32_C(0x40000000));
38 assert(multiplier <= INT32_C(0x7FFFFF80));
39
40 // Shift is in [0, 31] range.
41 const int32_t shift = 127 + 31 - 32 - (fp32_to_bits(scale) >> 23);
42 assert(shift >= 0);
43 assert(shift < 32);
44
45 const int32x4_t vmultiplier = vdupq_n_s32(multiplier);
46 const int16x8_t vzero_point = vdupq_n_s16((int16_t)(uint16_t) zero_point);
47 const int32x4_t vshift = vdupq_n_s32(-shift);
48 const int32x4_t vshift_eq_0_mask = vreinterpretq_s32_u32(vceqq_s32(vshift, vmovq_n_s32(0)));
49 const uint8x16_t vqmin = vdupq_n_u8(qmin);
50 const uint8x16_t vqmax = vdupq_n_u8(qmax);
51 for (; n != 0; n -= 16) {
52 const int32x4_t x = vld1q_s32(input);
53 const int32x4_t y = vld1q_s32(input + 4);
54 const int32x4_t z = vld1q_s32(input + 8);
55 const int32x4_t w = vld1q_s32(input + 12);
56 input += 16;
57
58 // Directly use VQRDMULH/SQRDMULH instruction for Q31 multiplication with rounding.
59 // Although these instruction saturate out-of-range outputs, we never hit this case in requantization.
60 const int32x4_t x_product = vqrdmulhq_s32(x, vmultiplier);
61 const int32x4_t y_product = vqrdmulhq_s32(y, vmultiplier);
62 const int32x4_t z_product = vqrdmulhq_s32(z, vmultiplier);
63 const int32x4_t w_product = vqrdmulhq_s32(w, vmultiplier);
64
65 // Shift the 32-bit product right with rounding.
66 // Rounding is performed towards closest integer, with midpoints rounded up (same as away from zero).
67 //
68 // We leverage the "right shift with rounding" instruction (VRSHL.S32 on ARM NEON, SRSHL in ARM64 Advanced SIMD) to
69 // do the shift. However, as this instruction rounds midpoints up, rather than away from zero, we adjust the input
70 // by subtracting 1 from negative values, but only if shift is non-zero.
71 const int32x4_t x_adjusted_product = vsraq_n_s32(x_product, vbicq_s32(x, vshift_eq_0_mask), 31);
72 const int32x4_t y_adjusted_product = vsraq_n_s32(y_product, vbicq_s32(y, vshift_eq_0_mask), 31);
73 const int32x4_t z_adjusted_product = vsraq_n_s32(z_product, vbicq_s32(z, vshift_eq_0_mask), 31);
74 const int32x4_t w_adjusted_product = vsraq_n_s32(w_product, vbicq_s32(w, vshift_eq_0_mask), 31);
75
76 const int32x4_t x_scaled = vrshlq_s32(x_adjusted_product, vshift);
77 const int32x4_t y_scaled = vrshlq_s32(y_adjusted_product, vshift);
78 const int32x4_t z_scaled = vrshlq_s32(z_adjusted_product, vshift);
79 const int32x4_t w_scaled = vrshlq_s32(w_adjusted_product, vshift);
80
81 #ifdef __aarch64__
82 const int16x8_t xy_packed = vqaddq_s16(vqmovn_high_s32(vqmovn_s32(x_scaled), y_scaled), vzero_point);
83 const int16x8_t zw_packed = vqaddq_s16(vqmovn_high_s32(vqmovn_s32(z_scaled), w_scaled), vzero_point);
84 const uint8x16_t xyzw_packed = vqmovun_high_s16(vqmovun_s16(xy_packed), zw_packed);
85 #else
86 const int16x8_t xy_packed = vqaddq_s16(vcombine_s16(vqmovn_s32(x_scaled), vqmovn_s32(y_scaled)), vzero_point);
87 const int16x8_t zw_packed = vqaddq_s16(vcombine_s16(vqmovn_s32(z_scaled), vqmovn_s32(w_scaled)), vzero_point);
88 const uint8x16_t xyzw_packed = vcombine_u8(vqmovun_s16(xy_packed), vqmovun_s16(zw_packed));
89 #endif
90
91 const uint8x16_t xyzw_clamped = vmaxq_u8(vminq_u8(xyzw_packed, vqmax), vqmin);
92
93 // AArch32 version:
94 // 4x VQRDMULH.S32 Qd, Qm, Qn
95 // 4x VAND Qd, Qm, Dn
96 // 4x VSRA.S32 Qd, Qm, #31
97 // 4x VRSHL.S32 Qd, Qm, Qn
98 // 4x VQMOVN.S32 Dd, Qm
99 // 2x VADD.S16 Qd, Qm, Qn
100 // 2x VQMOVUN.S16 Dd, Qm
101 // 1x VMAX.U8 Qd, Qm, Qn
102 // 1x VMIN.U8 Qd, Qm, Qn
103 // ---------------------
104 // 26 instructions total
105 //
106 // AArch64 version:
107 // 4x SQRDMULH Vd.4S, Vn.4S, Vm.4S
108 // 4x AND Vd.16B, Vn.16B, Vm.16B
109 // 4x SSRA Vd.4S, Vn.4S, #31
110 // 4x SRSHL Vd.4S, Vn.4S, Vm.4S
111 // 2x SQXTN Vd.4H, Vn.4S
112 // 2x SQXTN2 Vd.8H, Vn.4S
113 // 2x ADD Vd.8H, Vn.8H, Vm.8H
114 // 1x SQXTUN Vd.8B, Vn.8H
115 // 1x SQXTUN2 Vd.16B, Vn.8H
116 // 1x UMIN Vd.16B, Vn.16B, Vm.16B
117 // 1x UMAX Vd.16B, Vn.16B, Vm.16B
118 // ---------------------
119 // 26 instructions total
120
121 vst1q_u8(output, xyzw_clamped);
122 output += 16;
123 }
124 }
125