1 /*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include <executorch/backends/cadence/hifi/kernels/kernels.h>
10 #include <xa_nnlib_common.h>
11 #include <xa_nnlib_common_macros.h>
12
13 namespace cadence {
14 namespace impl {
15 namespace HiFi {
16 namespace kernels {
17
18 __attribute__((always_inline)) void
memcpy(void * dst,const void * src,size_t num_bytes)19 memcpy(void* dst, const void* src, size_t num_bytes) {
20 MEMCPY_8b(dst, src, num_bytes);
21 }
22
23 // Quantize a fp32 value to an int8_t/uint8_t value
24 template <typename T>
25 __attribute__((always_inline)) T
quantize(const float x,float scale,int32_t zero_point)26 quantize(const float x, float scale, int32_t zero_point) {
27 constexpr float min_val = std::numeric_limits<T>::min();
28 constexpr float max_val = std::numeric_limits<T>::max();
29 float tmp = roundf(x * scale + zero_point);
30 return std::max(std::min(tmp, max_val), min_val);
31 }
32
33 // Quantize a fp32 array to an int8_t/uint8_t/int16_t array
34 template <typename T>
quantize(T * __restrict__ y,const float * __restrict__ x,float scale,int32_t zero_point,size_t size)35 void quantize(
36 T* __restrict__ y,
37 const float* __restrict__ x,
38 float scale,
39 int32_t zero_point,
40 size_t size) {
41 xtfloatx2 scale_vec = (xtfloatx2)scale;
42 xtfloatx2 zero_vec = XT_FLOAT_SX2(zero_point, 0);
43
44 constexpr float min_val = std::numeric_limits<T>::min();
45 constexpr float max_val = std::numeric_limits<T>::max();
46
47 const xtfloatx2* __restrict__ p0 = (const xtfloatx2* __restrict__)x;
48 ae_valign va0 = XT_LASX2PP(p0);
49
50 size_t i = 0;
51 // Vectorize by 2
52 for (; i < (size & ~1); i += 2) {
53 xtfloatx2 in_vec;
54 XT_LASX2IP(in_vec, va0, p0);
55 xtfloatx2 acc = zero_vec;
56 XT_MADD_SX2(acc, scale_vec, in_vec);
57 xtfloatx2 t0 = XT_FIROUND_SX2(acc);
58 ae_int32x2 t1 =
59 XT_UTRUNC_SX2(XT_MAX_SX2(XT_MIN_SX2(t0, max_val), min_val), 0);
60 y[i] = AE_MOVAD32_H(t1);
61 y[i + 1] = AE_MOVAD32_L(t1);
62 }
63 // Handle residual iteration
64 if (i < size) {
65 y[i] = quantize<T>(x[i], scale, zero_point);
66 }
67 }
68
69 // Dequantize an int8_t/uint8_t value to a fp32 value
70 template <typename T>
71 __attribute__((always_inline)) float
dequantize(const T x,float scale,int32_t zero_point)72 dequantize(const T x, float scale, int32_t zero_point) {
73 return scale * (x - zero_point);
74 }
75
76 // Deuantize an int8_t/uint8_t/int16_t array to an fp32 array
77 template <typename T>
dequantize(float * __restrict__ y,const T * __restrict__ x,float scale,int32_t zero_point,size_t size)78 void dequantize(
79 float* __restrict__ y,
80 const T* __restrict__ x,
81 float scale,
82 int32_t zero_point,
83 size_t size) {
84 xtfloatx2 scale_vec = (xtfloatx2)scale;
85 xtfloatx2 zero_vec = XT_FLOAT_SX2(zero_point, 0);
86
87 xtfloatx2* __restrict__ p0 = (xtfloatx2*)y;
88 ae_valign va0 = AE_ZALIGN64();
89 size_t i = 0;
90 // Vectorize by 2
91 for (size_t e = (size >> 1) << 1; i < e; i += 2) {
92 xtfloatx2 in_vec = {(float)x[i], (float)x[i + 1]};
93 xtfloatx2 t0 = XT_SUB_SX2(in_vec, zero_vec);
94 xtfloatx2 t1 = XT_MUL_SX2(t0, scale_vec);
95 XT_SASX2IP(t1, va0, p0);
96 }
97 // Flush the output stream
98 XT_SASX2POSFP(va0, p0);
99
100 // Handle residual iteration
101 if (i < size) {
102 y[i] = dequantize<T>(x[i], scale, zero_point);
103 }
104 }
105
106 // Requantize the int8_t/uint8_t in value to a uint8_t/int8_t out value.
107 // The scale and zero_point for requantization are in the args.
108 template <typename IT, typename OT>
requantize(const IT in,float in_scale,int32_t in_zero_point,float inv_out_scale,int32_t out_zero_point)109 __attribute__((always_inline)) OT requantize(
110 const IT in,
111 float in_scale,
112 int32_t in_zero_point,
113 float inv_out_scale,
114 int32_t out_zero_point) {
115 float dequant = dequantize<IT>(in, in_scale, in_zero_point);
116 return quantize<OT>(dequant, inv_out_scale, out_zero_point);
117 }
118
119 // Requantize the int8_t/uint8_t in array to a uint8_t/int8_t out array.
120 // The scale and zero_point for requantization are in the args.
121 template <typename IT, typename OT>
requantize(OT * __restrict__ out,const IT * __restrict__ in,float in_scale,int32_t in_zero_point,float inv_out_scale,int32_t out_zero_point,size_t size)122 void requantize(
123 OT* __restrict__ out,
124 const IT* __restrict__ in,
125 float in_scale,
126 int32_t in_zero_point,
127 float inv_out_scale,
128 int32_t out_zero_point,
129 size_t size) {
130 xtfloatx2 in_scale_vec = (xtfloatx2)in_scale;
131 xtfloatx2 in_zero_vec = XT_FLOAT_SX2(in_zero_point, 0);
132 xtfloatx2 inv_out_scale_vec = (xtfloatx2)inv_out_scale;
133 xtfloatx2 out_zero_vec = XT_FLOAT_SX2(out_zero_point, 0);
134
135 float min_val = std::numeric_limits<OT>::min();
136 float max_val = std::numeric_limits<OT>::max();
137
138 size_t i = 0;
139 // Vectorize by 2
140 for (; i < (size & ~1); i += 2) {
141 xtfloatx2 in_vec = {(float)in[i], (float)in[i + 1]};
142 xtfloatx2 t0 = XT_SUB_SX2(in_vec, in_zero_vec);
143 xtfloatx2 t1 = XT_MUL_SX2(t0, in_scale_vec);
144
145 xtfloatx2 acc = out_zero_vec;
146 XT_MADD_SX2(acc, inv_out_scale_vec, t1);
147 xtfloatx2 t2 = XT_FIROUND_SX2(acc);
148 ae_int32x2 t3 =
149 XT_UTRUNC_SX2(XT_MAX_SX2(XT_MIN_SX2(t2, max_val), min_val), 0);
150 out[i] = AE_MOVAD32_H(t3);
151 out[i + 1] = AE_MOVAD32_L(t3);
152 }
153 // Handle residual iteration
154 if (i < size) {
155 out[i] = requantize<IT, OT>(
156 in[i], in_scale, in_zero_point, inv_out_scale, out_zero_point);
157 }
158 }
159
160 // explicit template instantiation
161
162 #define typed_quantize_val(dtype) \
163 template __attribute__((always_inline)) dtype quantize( \
164 const float x, float inv_scale, int32_t zero_point);
165 typed_quantize_val(int8_t);
166 typed_quantize_val(uint8_t);
167 typed_quantize_val(int16_t);
168 typed_quantize_val(uint16_t);
169 #undef typed_quantize_val
170
171 #define typed_quantize_vec(dtype) \
172 template void quantize( \
173 dtype* __restrict__ y, \
174 const float* __restrict__ x, \
175 float inv_scale, \
176 int32_t zero_point, \
177 size_t size);
178 typed_quantize_vec(int8_t);
179 typed_quantize_vec(uint8_t);
180 typed_quantize_vec(int16_t);
181 typed_quantize_vec(uint16_t);
182 typed_quantize_vec(int32_t);
183 #undef typed_quantize_vec
184
185 #define typed_dequantize_val(dtype) \
186 template __attribute__((always_inline)) float dequantize( \
187 const dtype x, float scale, int32_t zero_point);
188 typed_dequantize_val(int8_t);
189 typed_dequantize_val(uint8_t);
190 typed_dequantize_val(int16_t);
191 typed_dequantize_val(uint16_t);
192 #undef typed_dequantize_val
193
194 #define typed_dequantize_vec(dtype) \
195 template void dequantize( \
196 float* __restrict__ y, \
197 const dtype* __restrict__ x, \
198 float scale, \
199 int32_t zero_point, \
200 size_t size);
201 typed_dequantize_vec(int8_t);
202 typed_dequantize_vec(uint8_t);
203 typed_dequantize_vec(int16_t);
204 typed_dequantize_vec(uint16_t);
205 typed_dequantize_vec(int32_t);
206 #undef typed_dequantize_vec
207
208 #define typed_requantize_val(itype, otype) \
209 template __attribute__((always_inline)) otype requantize( \
210 const itype in, \
211 float in_scale, \
212 int32_t in_zero_point, \
213 float inv_out_scale, \
214 int32_t out_zero_point);
215 typed_requantize_val(int8_t, int8_t);
216 typed_requantize_val(uint8_t, uint8_t);
217 typed_requantize_val(int8_t, uint8_t);
218 typed_requantize_val(uint8_t, int8_t);
219 #undef typed_requantize_val
220
221 #define typed_requantize_vec(itype, otype) \
222 template void requantize( \
223 otype* __restrict__ out, \
224 const itype* __restrict__ in, \
225 float in_scale, \
226 int32_t in_zero_point, \
227 float inv_out_scale, \
228 int32_t out_zero_point, \
229 size_t size);
230 typed_requantize_vec(int8_t, int8_t);
231 typed_requantize_vec(uint8_t, uint8_t);
232 typed_requantize_vec(int8_t, uint8_t);
233 typed_requantize_vec(uint8_t, int8_t);
234 #undef typed_requantize_vec
235
236 }; // namespace kernels
237 }; // namespace HiFi
238 }; // namespace impl
239 }; // namespace cadence
240