char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Max Krasnyanskyhexagon: support for multi-device model split (aka row-split) (#28589)eafe15a5e

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
41.2 KiB725 linesraw
1#ifndef HVX_ARITH_H
2#define HVX_ARITH_H
3
4#include <assert.h>
5#include <stddef.h>
6#include <stdint.h>
7#include <math.h>
8
9#include "hvx-base.h"
10#include "hex-utils.h"
11
12//
13// Binary operations (add, mul, sub)
14//
15
16#define UNUSED(x) (void)(x)
17
18#define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \
19    do {                                                                                  \
20        dst_type * vdst  = (dst_type *) dst;                                              \
21        src0_type * vsrc0 = (src0_type *) src0;                                           \
22        src1_type * vsrc1 = (src1_type *) src1;                                           \
23                                                                                          \
24        const uint32_t epv  = 128 / (elem_size);                                          \
25        const uint32_t nvec = n / epv;                                                    \
26        const uint32_t nloe = n % epv;                                                    \
27                                                                                          \
28        uint32_t i = 0;                                                                   \
29                                                                                          \
30        _Pragma("unroll(4)")                                                              \
31        for (; i < nvec; i++) {                                                           \
32            vdst[i] = vec_op(vsrc0[i], vsrc1[i]);                                         \
33        }                                                                                 \
34        if (nloe) {                                                                       \
35            HVX_Vector v = vec_op(vsrc0[i], vsrc1[i]);                                    \
36            vec_store((void *) &vdst[i], nloe * (elem_size), v);                          \
37        }                                                                                 \
38    } while(0)
39
40#if __HVX_ARCH__ < 79
41
42#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
43#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b))
44#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
45
46#else
47
48#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
49#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b)
50#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
51
52#endif
53
54#define HVX_OP_ADD_F16(a, b) hvx_vec_add_f16_f16(a, b)
55#define HVX_OP_SUB_F16(a, b) hvx_vec_sub_f16_f16(a, b)
56#define HVX_OP_MUL_F16(a, b) hvx_vec_mul_f16_f16(a, b)
57
58// Generic macro to define alignment permutations for an op
59#define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE)                                           \
60static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
61    assert((uintptr_t) dst % 128 == 0);                                                                       \
62    assert((uintptr_t) src0 % 128 == 0);                                                                      \
63    assert((uintptr_t) src1 % 128 == 0);                                                                      \
64    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);    \
65}                                                                                                             \
66static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
67    assert((uintptr_t) dst % 128 == 0);                                                                       \
68    assert((uintptr_t) src0 % 128 == 0);                                                                      \
69    hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);   \
70}                                                                                                             \
71static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
72    assert((uintptr_t) dst % 128 == 0);                                                                       \
73    assert((uintptr_t) src1 % 128 == 0);                                                                      \
74    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);   \
75}                                                                                                             \
76static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
77    assert((uintptr_t) dst % 128 == 0);                                                                       \
78    hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);  \
79}                                                                                                             \
80static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
81    assert((uintptr_t) src0 % 128 == 0);                                                                      \
82    assert((uintptr_t) src1 % 128 == 0);                                                                      \
83    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);   \
84}                                                                                                             \
85static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
86    assert((uintptr_t) src0 % 128 == 0);                                                                      \
87    hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);  \
88}                                                                                                             \
89static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
90    assert((uintptr_t) src1 % 128 == 0);                                                                      \
91    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);  \
92}                                                                                                             \
93static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) {     \
94    hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \
95}                                                                                                             \
96
97DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f32, HVX_OP_ADD_F32, float)
98DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f32, HVX_OP_SUB_F32, float)
99DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f32, HVX_OP_MUL_F32, float)
100
101DEFINE_HVX_BINARY_OP_VARIANTS(hvx_add_f16, HVX_OP_ADD_F16, _Float16)
102DEFINE_HVX_BINARY_OP_VARIANTS(hvx_sub_f16, HVX_OP_SUB_F16, _Float16)
103DEFINE_HVX_BINARY_OP_VARIANTS(hvx_mul_f16, HVX_OP_MUL_F16, _Float16)
104
105// Dispatcher logic
106#define HVX_BINARY_DISPATCHER(OP_NAME)                                                                                                       \
107static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint32_t num_elems) { \
108    if (hex_is_aligned((void *) dst, 128)) {                                                                                                 \
109        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
110            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aaa(dst, src0, src1, num_elems);                                               \
111            else                                    OP_NAME##_aau(dst, src0, src1, num_elems);                                               \
112        } else {                                                                                                                             \
113            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_aua(dst, src0, src1, num_elems);                                               \
114            else                                    OP_NAME##_auu(dst, src0, src1, num_elems);                                               \
115        }                                                                                                                                    \
116    } else {                                                                                                                                 \
117        if (hex_is_aligned((void *) src0, 128)) {                                                                                            \
118            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uaa(dst, src0, src1, num_elems);                                               \
119            else                                    OP_NAME##_uau(dst, src0, src1, num_elems);                                               \
120        } else {                                                                                                                             \
121            if (hex_is_aligned((void *) src1, 128)) OP_NAME##_uua(dst, src0, src1, num_elems);                                               \
122            else                                    OP_NAME##_uuu(dst, src0, src1, num_elems);                                               \
123        }                                                                                                                                    \
124    }                                                                                                                                        \
125}
126
127HVX_BINARY_DISPATCHER(hvx_add_f32)
128HVX_BINARY_DISPATCHER(hvx_sub_f32)
129HVX_BINARY_DISPATCHER(hvx_mul_f32)
130
131HVX_BINARY_DISPATCHER(hvx_add_f16)
132HVX_BINARY_DISPATCHER(hvx_sub_f16)
133HVX_BINARY_DISPATCHER(hvx_mul_f16)
134
135// Mul-Mul Optimized
136static inline void hvx_mul_mul_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, const uint8_t * restrict src2, const uint32_t num_elems) {
137    assert((unsigned long) dst % 128 == 0);
138    assert((unsigned long) src0 % 128 == 0);
139    assert((unsigned long) src1 % 128 == 0);
140    assert((unsigned long) src2 % 128 == 0);
141
142    HVX_Vector * restrict vdst  = (HVX_Vector *) dst;
143    HVX_Vector * restrict vsrc0 = (HVX_Vector *) src0;
144    HVX_Vector * restrict vsrc1 = (HVX_Vector *) src1;
145    HVX_Vector * restrict vsrc2 = (HVX_Vector *) src2;
146
147    const uint32_t elem_size = sizeof(float);
148    const uint32_t epv  = 128 / elem_size;
149    const uint32_t nvec = num_elems / epv;
150    const uint32_t nloe = num_elems % epv;
151
152    uint32_t i = 0;
153
154    _Pragma("unroll(4)")
155    for (; i < nvec; i++) {
156        HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]);
157        vdst[i] = HVX_OP_MUL(v1, vsrc2[i]);
158    }
159
160    if (nloe) {
161        HVX_Vector v1 = HVX_OP_MUL_F32(vsrc0[i], vsrc1[i]);
162        HVX_Vector v2 = HVX_OP_MUL_F32(v1, vsrc2[i]);
163        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v2);
164    }
165}
166
167// Scalar Operations
168
169#define hvx_scalar_loop_body(dst_type, src_type, elem_size, vec_store, scalar_op_macro) \
170    do {                                                                                \
171        dst_type * restrict vdst = (dst_type *) dst;                                    \
172        src_type * restrict vsrc = (src_type *) src;                                    \
173                                                                                        \
174        const uint32_t epv  = 128 / (elem_size);                                        \
175        const uint32_t nvec = n / epv;                                                  \
176        const uint32_t nloe = n % epv;                                                  \
177                                                                                        \
178        uint32_t i = 0;                                                                 \
179                                                                                        \
180        _Pragma("unroll(4)")                                                            \
181        for (; i < nvec; i++) {                                                         \
182            HVX_Vector v = vsrc[i];                                                     \
183            vdst[i] = scalar_op_macro(v);                                               \
184        }                                                                               \
185        if (nloe) {                                                                     \
186            HVX_Vector v = vsrc[i];                                                     \
187            v = scalar_op_macro(v);                                                     \
188            vec_store((void *) &vdst[i], nloe * (elem_size), v);                        \
189        }                                                                               \
190    } while(0)
191
192#define HVX_OP_ADD_SCALAR_F32(v)                                   \
193    ({                                                             \
194        const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VwVw(inf, v); \
195        HVX_Vector out = HVX_OP_ADD_F32(v, val_vec);               \
196        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
197    })
198
199#define HVX_OP_MUL_SCALAR_F32(v) HVX_OP_MUL_F32(v, val_vec)
200#define HVX_OP_SUB_SCALAR_F32(v) HVX_OP_SUB_F32(v, val_vec)
201
202#define HVX_OP_ADD_SCALAR_F16(v)                                   \
203    ({                                                             \
204        const HVX_VectorPred pred_inf = Q6_Q_vcmp_eq_VhVh(inf, v); \
205        HVX_Vector out = HVX_OP_ADD_F16(v, val_vec);               \
206        Q6_V_vmux_QVV(pred_inf, inf, out);                         \
207    })
208
209#define HVX_OP_MUL_SCALAR_F16(v) HVX_OP_MUL_F16(v, val_vec)
210#define HVX_OP_SUB_SCALAR_F16(v) HVX_OP_SUB_F16(v, val_vec)
211
212// Scalar Variants
213
214// Generic macro to define alignment permutations for an op
215#define DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(OP_NAME, OP_MACRO, SPLAT_MACRO, ELEM_TYPE)                                  \
216static inline void OP_NAME##_aa(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
217    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
218    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
219    assert((uintptr_t) dst % 128 == 0);                                                                                  \
220    assert((uintptr_t) src % 128 == 0);                                                                                  \
221    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                          \
222}                                                                                                                        \
223static inline void OP_NAME##_au(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
224    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
225    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
226    assert((uintptr_t) dst % 128 == 0);                                                                                  \
227    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO);                         \
228}                                                                                                                        \
229static inline void OP_NAME##_ua(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
230    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
231    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
232    assert((uintptr_t) src % 128 == 0);                                                                                  \
233    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                         \
234}                                                                                                                        \
235static inline void OP_NAME##_uu(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, uint32_t n) { \
236    const HVX_Vector val_vec = SPLAT_MACRO(val);                                                                         \
237    const HVX_Vector inf = SPLAT_MACRO((ELEM_TYPE)INFINITY); UNUSED(inf);                                                \
238    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO);                        \
239}                                                                                                                        \
240
241DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f32, HVX_OP_ADD_SCALAR_F32, hvx_vec_splat_f32, float)
242DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f32, HVX_OP_SUB_SCALAR_F32, hvx_vec_splat_f32, float)
243DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f32, HVX_OP_MUL_SCALAR_F32, hvx_vec_splat_f32, float)
244
245DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_add_scalar_f16, HVX_OP_ADD_SCALAR_F16, hvx_vec_splat_f16, _Float16)
246DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_sub_scalar_f16, HVX_OP_SUB_SCALAR_F16, hvx_vec_splat_f16, _Float16)
247DEFINE_HVX_BINARY_SCALAR_OP_VARIANTS(hvx_mul_scalar_f16, HVX_OP_MUL_SCALAR_F16, hvx_vec_splat_f16, _Float16)
248
249// Dispatcher logic
250#define HVX_BINARY_SCALAR_DISPATCHER(OP_NAME, ELEM_TYPE)                                                                          \
251static inline void OP_NAME(uint8_t * restrict dst, const uint8_t * restrict src, const ELEM_TYPE val, const uint32_t num_elems) { \
252    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {                                                 \
253        OP_NAME##_aa(dst, src, val, num_elems);                                                                                   \
254    } else if (hex_is_aligned((void *) dst, 128)) {                                                                               \
255        OP_NAME##_au(dst, src, val, num_elems);                                                                                   \
256    } else if (hex_is_aligned((void *) src, 128)) {                                                                               \
257        OP_NAME##_ua(dst, src, val, num_elems);                                                                                   \
258    } else {                                                                                                                      \
259        OP_NAME##_uu(dst, src, val, num_elems);                                                                                   \
260    }                                                                                                                             \
261}
262
263HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f32, float)
264HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f32, float)
265HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f32, float)
266
267HVX_BINARY_SCALAR_DISPATCHER(hvx_add_scalar_f16, _Float16)
268HVX_BINARY_SCALAR_DISPATCHER(hvx_sub_scalar_f16, _Float16)
269HVX_BINARY_SCALAR_DISPATCHER(hvx_mul_scalar_f16, _Float16)
270
271// MIN Scalar variants
272
273#define HVX_OP_MIN_SCALAR(v) Q6_Vsf_vmin_VsfVsf(val_vec, v)
274
275static inline void hvx_min_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
276    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
277    assert((unsigned long) dst % 128 == 0);
278    assert((unsigned long) src % 128 == 0);
279    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR);
280}
281
282static inline void hvx_min_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
283    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
284    assert((unsigned long) dst % 128 == 0);
285    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_MIN_SCALAR);
286}
287
288static inline void hvx_min_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
289    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
290    assert((unsigned long) src % 128 == 0);
291    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR);
292}
293
294static inline void hvx_min_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
295    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
296    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_MIN_SCALAR);
297}
298
299static inline void hvx_min_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const int num_elems) {
300    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
301        hvx_min_scalar_f32_aa(dst, src, val, num_elems);
302    } else if (hex_is_aligned((void *) dst, 128)) {
303        hvx_min_scalar_f32_au(dst, src, val, num_elems);
304    } else if (hex_is_aligned((void *) src, 128)) {
305        hvx_min_scalar_f32_ua(dst, src, val, num_elems);
306    } else {
307        hvx_min_scalar_f32_uu(dst, src, val, num_elems);
308    }
309}
310
311// MAX Scalar variants
312
313#define HVX_OP_MAX_SCALAR(v) Q6_Vsf_vmax_VsfVsf(val_vec, v)
314
315static inline void hvx_max_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
316    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
317    assert((unsigned long) dst % 128 == 0);
318    assert((unsigned long) src % 128 == 0);
319    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_MAX_SCALAR);
320}
321
322static inline void hvx_max_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
323    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
324    assert((unsigned long) dst % 128 == 0);
325    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_MAX_SCALAR);
326}
327
328static inline void hvx_max_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
329    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
330    assert((unsigned long) src % 128 == 0);
331    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_MAX_SCALAR);
332}
333
334static inline void hvx_max_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float val, uint32_t n) {
335    const HVX_Vector val_vec = hvx_vec_splat_f32(val);
336    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_MAX_SCALAR);
337}
338
339static inline void hvx_max_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const int num_elems) {
340    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
341        hvx_max_scalar_f32_aa(dst, src, val, num_elems);
342    } else if (hex_is_aligned((void *) dst, 128)) {
343        hvx_max_scalar_f32_au(dst, src, val, num_elems);
344    } else if (hex_is_aligned((void *) src, 128)) {
345        hvx_max_scalar_f32_ua(dst, src, val, num_elems);
346    } else {
347        hvx_max_scalar_f32_uu(dst, src, val, num_elems);
348    }
349}
350
351// CLAMP Scalar variants
352
353#define HVX_OP_CLAMP_SCALAR(v)                                           \
354    ({                                                                   \
355        HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VsfVsf(v, max_vec); \
356        HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VsfVsf(min_vec, v); \
357        HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v);      \
358        Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp);                      \
359    })
360
361static inline void hvx_clamp_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {
362    const HVX_Vector min_vec = hvx_vec_splat_f32(min);
363    const HVX_Vector max_vec = hvx_vec_splat_f32(max);
364    assert((unsigned long) dst % 128 == 0);
365    assert((unsigned long) src % 128 == 0);
366    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR);
367}
368
369static inline void hvx_clamp_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {
370    const HVX_Vector min_vec = hvx_vec_splat_f32(min);
371    const HVX_Vector max_vec = hvx_vec_splat_f32(max);
372    assert((unsigned long) dst % 128 == 0);
373    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR);
374}
375
376static inline void hvx_clamp_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {
377    const HVX_Vector min_vec = hvx_vec_splat_f32(min);
378    const HVX_Vector max_vec = hvx_vec_splat_f32(max);
379    assert((unsigned long) src % 128 == 0);
380    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR);
381}
382
383static inline void hvx_clamp_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, uint32_t n) {
384    const HVX_Vector min_vec = hvx_vec_splat_f32(min);
385    const HVX_Vector max_vec = hvx_vec_splat_f32(max);
386    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR);
387}
388
389static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float min, const float max, const int num_elems) {
390    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
391        hvx_clamp_scalar_f32_aa(dst, src, min, max, num_elems);
392    } else if (hex_is_aligned((void *) dst, 128)) {
393        hvx_clamp_scalar_f32_au(dst, src, min, max, num_elems);
394    } else if (hex_is_aligned((void *) src, 128)) {
395        hvx_clamp_scalar_f32_ua(dst, src, min, max, num_elems);
396    } else {
397        hvx_clamp_scalar_f32_uu(dst, src, min, max, num_elems);
398    }
399}
400
401#define HVX_OP_CLAMP_SCALAR_F16(v)                                       \
402    ({                                                                   \
403        HVX_VectorPred pred_cap_right = Q6_Q_vcmp_gt_VhfVhf(v, max_vec); \
404        HVX_VectorPred pred_cap_left  = Q6_Q_vcmp_gt_VhfVhf(min_vec, v); \
405        HVX_Vector tmp = Q6_V_vmux_QVV(pred_cap_right, max_vec, v);      \
406        Q6_V_vmux_QVV(pred_cap_left, min_vec, tmp);                      \
407    })
408
409static inline void hvx_clamp_scalar_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
410    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
411    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
412    assert((unsigned long) dst % 128 == 0);
413    assert((unsigned long) src % 128 == 0);
414    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(_Float16), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR_F16);
415}
416
417static inline void hvx_clamp_scalar_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
418    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
419    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
420    assert((unsigned long) dst % 128 == 0);
421    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(_Float16), hvx_vec_store_a, HVX_OP_CLAMP_SCALAR_F16);
422}
423
424static inline void hvx_clamp_scalar_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
425    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
426    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
427    assert((unsigned long) src % 128 == 0);
428    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(_Float16), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR_F16);
429}
430
431static inline void hvx_clamp_scalar_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, uint32_t n) {
432    const HVX_Vector min_vec = hvx_vec_splat_f16(min);
433    const HVX_Vector max_vec = hvx_vec_splat_f16(max);
434    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(_Float16), hvx_vec_store_u, HVX_OP_CLAMP_SCALAR_F16);
435}
436
437static inline void hvx_clamp_scalar_f16(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 min, const _Float16 max, const int num_elems) {
438    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
439        hvx_clamp_scalar_f16_aa(dst, src, min, max, num_elems);
440    } else if (hex_is_aligned((void *) dst, 128)) {
441        hvx_clamp_scalar_f16_au(dst, src, min, max, num_elems);
442    } else if (hex_is_aligned((void *) src, 128)) {
443        hvx_clamp_scalar_f16_ua(dst, src, min, max, num_elems);
444    } else {
445        hvx_clamp_scalar_f16_uu(dst, src, min, max, num_elems);
446    }
447}
448
449#define HVX_OP_LEAKY_RELU_SCALAR(v)                                 \
450    ({                                                              \
451        HVX_VectorPred pred_neg = Q6_Q_vcmp_gt_VsfVsf(zero_vec, v); \
452        HVX_Vector     scaled   = HVX_OP_MUL_F32(v, ns_vec);        \
453        Q6_V_vmux_QVV(pred_neg, scaled, v);                         \
454    })
455
456static inline void hvx_leaky_relu_scalar_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const float ns, uint32_t n) {
457    const HVX_Vector zero_vec = hvx_vec_splat_f32(0.0f);
458    const HVX_Vector ns_vec   = hvx_vec_splat_f32(ns);
459    assert((unsigned long) dst % 128 == 0);
460    assert((unsigned long) src % 128 == 0);
461    hvx_scalar_loop_body(HVX_Vector, HVX_Vector, sizeof(float), hvx_vec_store_a, HVX_OP_LEAKY_RELU_SCALAR);
462}
463
464static inline void hvx_leaky_relu_scalar_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, const float ns, uint32_t n) {
465    const HVX_Vector zero_vec = hvx_vec_splat_f32(0.0f);
466    const HVX_Vector ns_vec   = hvx_vec_splat_f32(ns);
467    assert((unsigned long) dst % 128 == 0);
468    hvx_scalar_loop_body(HVX_Vector, HVX_UVector, sizeof(float), hvx_vec_store_a, HVX_OP_LEAKY_RELU_SCALAR);
469}
470
471static inline void hvx_leaky_relu_scalar_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, const float ns, uint32_t n) {
472    const HVX_Vector zero_vec = hvx_vec_splat_f32(0.0f);
473    const HVX_Vector ns_vec   = hvx_vec_splat_f32(ns);
474    assert((unsigned long) src % 128 == 0);
475    hvx_scalar_loop_body(HVX_UVector, HVX_Vector, sizeof(float), hvx_vec_store_u, HVX_OP_LEAKY_RELU_SCALAR);
476}
477
478static inline void hvx_leaky_relu_scalar_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, const float ns, uint32_t n) {
479    const HVX_Vector zero_vec = hvx_vec_splat_f32(0.0f);
480    const HVX_Vector ns_vec   = hvx_vec_splat_f32(ns);
481    hvx_scalar_loop_body(HVX_UVector, HVX_UVector, sizeof(float), hvx_vec_store_u, HVX_OP_LEAKY_RELU_SCALAR);
482}
483
484static inline void hvx_leaky_relu_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float ns, const int num_elems) {
485    if (hex_is_aligned((void *) dst, 128) && hex_is_aligned((void *) src, 128)) {
486        hvx_leaky_relu_scalar_f32_aa(dst, src, ns, num_elems);
487    } else if (hex_is_aligned((void *) dst, 128)) {
488        hvx_leaky_relu_scalar_f32_au(dst, src, ns, num_elems);
489    } else if (hex_is_aligned((void *) src, 128)) {
490        hvx_leaky_relu_scalar_f32_ua(dst, src, ns, num_elems);
491    } else {
492        hvx_leaky_relu_scalar_f32_uu(dst, src, ns, num_elems);
493    }
494}
495
496//
497// Abs
498//
499
500static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
501    assert((unsigned long) dst % 128 == 0);
502    assert((unsigned long) src % 128 == 0);
503
504    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
505    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
506
507    const uint32_t elem_size = sizeof(float);
508    const uint32_t epv       = 128 / elem_size;
509    const uint32_t nvec      = n / epv;
510    const uint32_t nloe      = n % epv;
511
512    uint32_t i = 0;
513
514    _Pragma("unroll(4)")
515    for (; i < nvec; i++) {
516        vdst[i] = hvx_vec_abs_f32(vsrc[i]);
517    }
518    if (nloe) {
519        HVX_Vector v = hvx_vec_abs_f32(vsrc[i]);
520        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
521    }
522}
523
524#define hvx_abs_f16_loop_body(dst_type, src_type, vec_store)               \
525    do {                                                                   \
526        dst_type * restrict vdst  = (dst_type *) dst;                      \
527        src_type * restrict vsrc = (src_type *) src;                       \
528                                                                           \
529        const uint32_t elem_size = sizeof(_Float16);                       \
530        const uint32_t epv  = 128 / elem_size;                             \
531        const uint32_t nvec = n / epv;                                     \
532        const uint32_t nloe = n % epv;                                     \
533                                                                           \
534        uint32_t i = 0;                                                    \
535                                                                           \
536        _Pragma("unroll(4)")                                               \
537        for (; i < nvec; i++) {                                            \
538            vdst[i] = hvx_vec_abs_f16(vsrc[i]);                            \
539        }                                                                  \
540        if (nloe) {                                                        \
541            HVX_Vector v = hvx_vec_abs_f16(vsrc[i]);                       \
542            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
543        }                                                                  \
544    } while(0)
545
546static inline void hvx_abs_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
547    assert((unsigned long) dst % 128 == 0);
548    assert((unsigned long) src % 128 == 0);
549    hvx_abs_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
550}
551
552static inline void hvx_abs_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
553    assert((unsigned long) dst % 128 == 0);
554    hvx_abs_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
555}
556
557static inline void hvx_abs_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
558    assert((unsigned long) src % 128 == 0);
559    hvx_abs_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
560}
561
562static inline void hvx_abs_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
563    hvx_abs_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
564}
565
566static inline void hvx_abs_f16(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) {
567    if (hex_is_aligned((void *) dst, 128)) {
568        if (hex_is_aligned((void *) src, 128)) {
569            hvx_abs_f16_aa(dst, src, num_elems);
570        } else {
571            hvx_abs_f16_au(dst, src, num_elems);
572        }
573    } else {
574        if (hex_is_aligned((void *) src, 128)) {
575            hvx_abs_f16_ua(dst, src, num_elems);
576        } else {
577            hvx_abs_f16_uu(dst, src, num_elems);
578        }
579    }
580}
581
582//
583// Square
584//
585
586#define hvx_sqr_f32_loop_body(dst_type, src_type, vec_store)               \
587    do {                                                                   \
588        dst_type * restrict vdst  = (dst_type *) dst;                      \
589        src_type * restrict vsrc = (src_type *) src;                       \
590                                                                           \
591        const uint32_t elem_size = sizeof(float);                          \
592        const uint32_t epv  = 128 / elem_size;                             \
593        const uint32_t nvec = n / epv;                                     \
594        const uint32_t nloe = n % epv;                                     \
595                                                                           \
596        uint32_t i = 0;                                                    \
597                                                                           \
598        _Pragma("unroll(4)")                                               \
599        for (; i < nvec; i++) {                                            \
600            vdst[i] = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);                    \
601        }                                                                  \
602        if (nloe) {                                                        \
603            HVX_Vector v = HVX_OP_MUL_F32(vsrc[i], vsrc[i]);               \
604            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
605        }                                                                  \
606    } while(0)
607
608static inline void hvx_sqr_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
609    assert((unsigned long) dst % 128 == 0);
610    assert((unsigned long) src % 128 == 0);
611    hvx_sqr_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
612}
613
614static inline void hvx_sqr_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
615    assert((unsigned long) dst % 128 == 0);
616    hvx_sqr_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
617}
618
619static inline void hvx_sqr_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
620    assert((unsigned long) src % 128 == 0);
621    hvx_sqr_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
622}
623
624static inline void hvx_sqr_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
625    hvx_sqr_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
626}
627
628static inline void hvx_sqr_f32(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) {
629    if (hex_is_aligned((void *) dst, 128)) {
630        if (hex_is_aligned((void *) src, 128)) {
631            hvx_sqr_f32_aa(dst, src, num_elems);
632        } else {
633            hvx_sqr_f32_au(dst, src, num_elems);
634        }
635    } else {
636        if (hex_is_aligned((void *) src, 128)) {
637            hvx_sqr_f32_ua(dst, src, num_elems);
638        } else {
639            hvx_sqr_f32_uu(dst, src, num_elems);
640        }
641    }
642}
643
644#define hvx_sqr_f16_loop_body(dst_type, src_type, vec_store)               \
645    do {                                                                   \
646        dst_type * restrict vdst  = (dst_type *) dst;                      \
647        src_type * restrict vsrc = (src_type *) src;                       \
648                                                                           \
649        const uint32_t elem_size = sizeof(_Float16);                       \
650        const uint32_t epv  = 128 / elem_size;                             \
651        const uint32_t nvec = n / epv;                                     \
652        const uint32_t nloe = n % epv;                                     \
653                                                                           \
654        uint32_t i = 0;                                                    \
655                                                                           \
656        _Pragma("unroll(4)")                                               \
657        for (; i < nvec; i++) {                                            \
658            vdst[i] = HVX_OP_MUL_F16(vsrc[i], vsrc[i]);                    \
659        }                                                                  \
660        if (nloe) {                                                        \
661            HVX_Vector v = HVX_OP_MUL_F16(vsrc[i], vsrc[i]);               \
662            vec_store((void *) &vdst[i], nloe * elem_size, v);             \
663        }                                                                  \
664    } while(0)
665
666static inline void hvx_sqr_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
667    assert((unsigned long) dst % 128 == 0);
668    assert((unsigned long) src % 128 == 0);
669    hvx_sqr_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
670}
671
672static inline void hvx_sqr_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
673    assert((unsigned long) dst % 128 == 0);
674    hvx_sqr_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
675}
676
677static inline void hvx_sqr_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
678    assert((unsigned long) src % 128 == 0);
679    hvx_sqr_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
680}
681
682static inline void hvx_sqr_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
683    hvx_sqr_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
684}
685
686static inline void hvx_sqr_f16(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t num_elems) {
687    if (hex_is_aligned((void *) dst, 128)) {
688        if (hex_is_aligned((void *) src, 128)) {
689            hvx_sqr_f16_aa(dst, src, num_elems);
690        } else {
691            hvx_sqr_f16_au(dst, src, num_elems);
692        }
693    } else {
694        if (hex_is_aligned((void *) src, 128)) {
695            hvx_sqr_f16_ua(dst, src, num_elems);
696        } else {
697            hvx_sqr_f16_uu(dst, src, num_elems);
698        }
699    }
700}
701
702#undef HVX_OP_ADD_F32
703#undef HVX_OP_SUB_F32
704#undef HVX_OP_MUL_F32
705#undef HVX_OP_ADD_F16
706#undef HVX_OP_SUB_F16
707#undef HVX_OP_MUL_F16
708#undef hvx_arith_loop_body
709#undef HVX_OP_ADD_SCALAR_F32
710#undef HVX_OP_SUB_SCALAR_F32
711#undef HVX_OP_MUL_SCALAR_F32
712#undef HVX_OP_ADD_SCALAR_F16
713#undef HVX_OP_SUB_SCALAR_F16
714#undef HVX_OP_MUL_SCALAR_F16
715#undef hvx_scalar_loop_body
716#undef HVX_OP_MIN_SCALAR
717#undef HVX_OP_MAX_SCALAR
718#undef HVX_OP_CLAMP_SCALAR
719#undef HVX_OP_CLAMP_SCALAR_F16
720#undef HVX_OP_LEAKY_RELU_SCALAR
721#undef DEFINE_HVX_BINARY_OP_VARIANTS
722#undef HVX_BINARY_DISPATCHER
723#undef UNUSED
724
725#endif // HVX_ARITH_H