char-fork/llama.cpp

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

Max Krasnyanskyhexagon: further improved pipeline of the core bits (L2, DMA, MM, FA) (#26049)0a50d9909

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
13.7 KiB336 linesraw
1#ifndef HVX_REDUCE_H
2#define HVX_REDUCE_H
3
4#include <math.h>
5#include <stdbool.h>
6#include <stdint.h>
7#include <assert.h>
8
9#include "hex-utils.h"
10#include "hvx-base.h"
11#include "hvx-types.h"
12
13static inline HVX_Vector hvx_vec_reduce_sum_n_i32(HVX_Vector in, unsigned int n) {
14    unsigned int total = n * 4;  // total vec nbytes
15    unsigned int width = 4;      // int32
16
17    HVX_Vector sum = in, sum_t;
18    while (width < total) {
19        sum_t = Q6_V_vror_VR(sum, width);     // rotate right
20        sum   = Q6_Vw_vadd_VwVw(sum_t, sum);  // elementwise sum
21        width = width << 1;
22    }
23    return sum;
24}
25
26static inline HVX_Vector hvx_vec_reduce_sum_i32(HVX_Vector in) {
27    return hvx_vec_reduce_sum_n_i32(in, 32);
28}
29
30static inline HVX_Vector hvx_vec_reduce_sum_n_qf32(HVX_Vector in, unsigned int n) {
31    unsigned int total = n * 4;  // total vec nbytes
32    unsigned int width = 4;      // fp32 nbytes
33
34    HVX_Vector sum = in, sum_t;
35    while (width < total) {
36        sum_t = Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum), width);  // rotate right
37        sum   = Q6_Vqf32_vadd_Vqf32Vsf(sum, sum_t);             // elementwise sum
38        width = width << 1;
39    }
40    return sum;
41}
42
43static inline HVX_Vector hvx_vec_reduce_sum_qf32(HVX_Vector in) {
44    return hvx_vec_reduce_sum_n_qf32(in, 32);
45}
46
47#if __HVX_ARCH__ > 75
48
49static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) {
50    HVX_VectorPair sum_p01 = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4);
51    HVX_VectorPair sum_p23 = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4);
52    HVX_Vector  sum_sf01  = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01));
53    HVX_Vector  sum_sf23  = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23));
54
55    HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(sum_sf23, sum_sf01, 8);
56    HVX_Vector  sum_sf       = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123));
57
58    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2));
59    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4));
60    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8));
61    return sum_sf;
62}
63
64static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) {
65    HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4);
66    HVX_Vector  sum_sf  = Q6_Vsf_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump));
67
68    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 2));
69    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 4));
70    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 8));
71    sum_sf = Q6_Vsf_vadd_VsfVsf(sum_sf, Q6_V_vror_VR(sum_sf, VLEN / 16));
72    return sum_sf;
73}
74
75static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) {
76    unsigned int total = n * 4;  // total vec nbytes
77    unsigned int width = 4;      // fp32 nbytes
78
79    HVX_Vector sum = in, sum_t;
80    while (width < total) {
81        sum_t = Q6_V_vror_VR(sum, width);       // rotate right
82        sum   = Q6_Vsf_vadd_VsfVsf(sum, sum_t); // elementwise sum
83        width = width << 1;
84    }
85    return sum;
86}
87
88#else
89
90static inline HVX_Vector hvx_vec_reduce_sum_f32x4(HVX_Vector_x4 in) {
91    HVX_VectorPair sum_p01  = Q6_W_vshuff_VVR(in.v[1], in.v[0], 4);
92    HVX_VectorPair sum_p23  = Q6_W_vshuff_VVR(in.v[3], in.v[2], 4);
93    HVX_Vector     sum_qf01 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p01), Q6_V_hi_W(sum_p01));
94    HVX_Vector     sum_qf23 = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p23), Q6_V_hi_W(sum_p23));
95
96    HVX_VectorPair sum_p0123 = Q6_W_vshuff_VVR(Q6_Vsf_equals_Vqf32(sum_qf23), Q6_Vsf_equals_Vqf32(sum_qf01), 8);
97    HVX_Vector     sum_qf    = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sum_p0123), Q6_V_hi_W(sum_p0123));
98
99    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2));
100    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4));
101    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8));
102    return Q6_Vsf_equals_Vqf32(sum_qf);
103}
104
105static inline HVX_Vector hvx_vec_reduce_sum_f32x2(HVX_Vector in0, HVX_Vector in1) {
106    HVX_VectorPair sump = Q6_W_vshuff_VVR(in1, in0, 4);
107    HVX_Vector  sum_qf  = Q6_Vqf32_vadd_VsfVsf(Q6_V_lo_W(sump), Q6_V_hi_W(sump));
108
109    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 2));
110    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 4));
111    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 8));
112    sum_qf = Q6_Vqf32_vadd_Vqf32Vsf(sum_qf, Q6_V_vror_VR(Q6_Vsf_equals_Vqf32(sum_qf), VLEN / 16));
113    return Q6_Vsf_equals_Vqf32(sum_qf);
114}
115
116static inline HVX_Vector hvx_vec_reduce_sum_n_f32(HVX_Vector in, unsigned int n) {
117    unsigned int total = n * 4;  // total vec nbytes
118    unsigned int width = 4;      // fp32 nbytes
119
120    HVX_Vector sum = in, sum_t;
121    while (width < total) {
122        sum_t = Q6_V_vror_VR(sum, width);                               // rotate right
123        sum   = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(sum, sum_t));  // elementwise sum
124        width = width << 1;
125    }
126    return sum;
127}
128
129#endif
130
131static inline HVX_Vector hvx_vec_reduce_sum_f32(HVX_Vector in) {
132    return hvx_vec_reduce_sum_n_f32(in, 32);
133}
134
135static inline HVX_Vector hvx_vec_reduce_max_f16(HVX_Vector in) {
136    unsigned total = 128;  // total vec nbytes
137    unsigned width = 2;    // fp16 nbytes
138
139    HVX_Vector _max = in, _max_t;
140    while (width < total) {
141        _max_t = Q6_V_vror_VR(_max, width);         // rotate right
142        _max   = Q6_Vhf_vmax_VhfVhf(_max_t, _max);  // elementwise max
143        width  = width << 1;
144    }
145
146    return _max;
147}
148
149static inline HVX_Vector hvx_vec_reduce_max2_f16(HVX_Vector in, HVX_Vector _max) {
150    unsigned total = 128;  // total vec nbytes
151    unsigned width = 2;    // fp32 nbytes
152
153    HVX_Vector _max_t;
154
155    _max = Q6_Vhf_vmax_VhfVhf(in, _max);
156    while (width < total) {
157        _max_t = Q6_V_vror_VR(_max, width);         // rotate right
158        _max   = Q6_Vhf_vmax_VhfVhf(_max_t, _max);  // elementwise max
159        width  = width << 1;
160    }
161
162    return _max;
163}
164
165static inline HVX_Vector hvx_vec_reduce_max_f32(HVX_Vector in) {
166    unsigned total = 128;  // total vec nbytes
167    unsigned width = 4;    // fp32 nbytes
168
169    HVX_Vector _max = in, _max_t;
170    while (width < total) {
171        _max_t = Q6_V_vror_VR(_max, width);         // rotate right
172        _max   = Q6_Vsf_vmax_VsfVsf(_max_t, _max);  // elementwise max
173        width  = width << 1;
174    }
175
176    return _max;
177}
178
179static inline HVX_Vector hvx_vec_reduce_max2_f32(HVX_Vector in, HVX_Vector _max) {
180    unsigned total = 128;  // total vec nbytes
181    unsigned width = 4;    // fp32 nbytes
182
183    HVX_Vector _max_t;
184
185    _max = Q6_Vsf_vmax_VsfVsf(in, _max);
186    while (width < total) {
187        _max_t = Q6_V_vror_VR(_max, width);         // rotate right
188        _max   = Q6_Vsf_vmax_VsfVsf(_max_t, _max);  // elementwise max
189        width  = width << 1;
190    }
191
192    return _max;
193}
194
195#define hvx_reduce_loop_body(src_type, init_vec, pad_vec, vec_op, reduce_op, scalar_reduce) \
196    do {                                                                                    \
197        src_type * restrict vsrc = (src_type *) src;                                        \
198        HVX_Vector acc = init_vec;                                                          \
199                                                                                            \
200        const uint32_t elem_size = sizeof(float);                                           \
201        const uint32_t epv  = 128 / elem_size;                                              \
202        const uint32_t nvec = num_elems / epv;                                              \
203        const uint32_t nloe = num_elems % epv;                                              \
204                                                                                            \
205        uint32_t i = 0;                                                                     \
206        _Pragma("unroll(4)")                                                                \
207        for (; i < nvec; i++) {                                                             \
208            acc = vec_op(acc, vsrc[i]);                                                     \
209        }                                                                                   \
210        if (nloe) {                                                                         \
211            const float * srcf = (const float *) src + i * epv;                             \
212            HVX_Vector in = *(HVX_UVector *) srcf;                                          \
213            HVX_Vector temp = Q6_V_valign_VVR(in, pad_vec, nloe * elem_size);               \
214            acc = vec_op(acc, temp);                                                        \
215        }                                                                                   \
216        HVX_Vector v = reduce_op(acc);                                                      \
217        return scalar_reduce(v);                                                            \
218    } while(0)
219
220#define HVX_REDUCE_MAX_OP(acc, val) Q6_Vsf_vmax_VsfVsf(acc, val)
221#define HVX_REDUCE_SUM_OP(acc, val) Q6_Vqf32_vadd_VsfVsf(Q6_Vsf_equals_Vqf32(acc), val)
222#define HVX_SUM_SQ_OP(acc, val) Q6_Vqf32_vadd_Vqf32Vqf32(acc, Q6_Vqf32_vmpy_VsfVsf(val, val))
223#define HVX_REDUCE_MAX_SCALAR(v) hvx_vec_get_f32(v)
224#define HVX_REDUCE_SUM_SCALAR(v) hvx_vec_get_f32(Q6_Vsf_equals_Vqf32(v))
225
226// Max variants
227
228static inline float hvx_reduce_max_f32_a(const uint8_t * restrict src, const int num_elems) {
229    HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]);
230    assert((unsigned long) src % 128 == 0);
231    hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR);
232}
233
234static inline float hvx_reduce_max_f32_u(const uint8_t * restrict src, const int num_elems) {
235    HVX_Vector init_vec = hvx_vec_splat_f32(((const float *) src)[0]);
236    hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_MAX_OP, hvx_vec_reduce_max_f32, HVX_REDUCE_MAX_SCALAR);
237}
238
239static inline float hvx_reduce_max_f32(const uint8_t * restrict src, const int num_elems) {
240    if (hex_is_aligned((void *) src, 128)) {
241        return hvx_reduce_max_f32_a(src, num_elems);
242    } else {
243        return hvx_reduce_max_f32_u(src, num_elems);
244    }
245}
246
247// Sum variants
248
249static inline float hvx_reduce_sum_f32_a(const uint8_t * restrict src, const int num_elems) {
250    HVX_Vector init_vec = Q6_V_vsplat_R(0);
251    assert((unsigned long) src % 128 == 0);
252    hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
253}
254
255static inline float hvx_reduce_sum_f32_u(const uint8_t * restrict src, const int num_elems) {
256    HVX_Vector init_vec = Q6_V_vsplat_R(0);
257    hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_REDUCE_SUM_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
258}
259
260static inline float hvx_reduce_sum_f32(const uint8_t * restrict src, const int num_elems) {
261    if (hex_is_aligned((void *) src, 128)) {
262        return hvx_reduce_sum_f32_a(src, num_elems);
263    } else {
264        return hvx_reduce_sum_f32_u(src, num_elems);
265    }
266}
267
268// Sum of squares variants
269
270static inline float hvx_sum_of_squares_f32_a(const uint8_t * restrict src, const int num_elems) {
271    HVX_Vector init_vec = Q6_V_vsplat_R(0);
272    assert((uintptr_t) src % 128 == 0);
273    hvx_reduce_loop_body(HVX_Vector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
274}
275
276static inline float hvx_sum_of_squares_f32_u(const uint8_t * restrict src, const int num_elems) {
277    HVX_Vector init_vec = Q6_V_vsplat_R(0);
278    hvx_reduce_loop_body(HVX_UVector, init_vec, init_vec, HVX_SUM_SQ_OP, hvx_vec_reduce_sum_qf32, HVX_REDUCE_SUM_SCALAR);
279}
280
281static inline float hvx_sum_of_squares_f32(const uint8_t * restrict src, const int num_elems) {
282    if (hex_is_aligned((void *) src, 128)) {
283        return hvx_sum_of_squares_f32_a(src, num_elems);
284    } else {
285        return hvx_sum_of_squares_f32_u(src, num_elems);
286    }
287}
288
289// Signed 32-bit Integer Max variants
290
291static inline HVX_Vector hvx_vec_reduce_max_n_i32(HVX_Vector in, unsigned int n) {
292    unsigned int total = n * 4;  // total vec nbytes
293    unsigned int width = 4;      // int32 nbytes
294
295    HVX_Vector max_val = in, max_t;
296    while (width < total) {
297        max_t = Q6_V_vror_VR(max_val, width);          // rotate right
298        max_val = Q6_Vw_vmax_VwVw(max_t, max_val);      // elementwise signed max
299        width = width << 1;
300    }
301    return max_val;
302}
303
304static inline HVX_Vector hvx_vec_reduce_max_i32(HVX_Vector in) {
305    return hvx_vec_reduce_max_n_i32(in, 32);
306}
307
308static inline int32_t hvx_reduce_max_i32_a(const uint8_t * restrict src, const int num_elems) {
309    HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
310    HVX_Vector pad_vec  = Q6_V_vsplat_R(0x80000000);
311    assert((uintptr_t) src % 128 == 0);
312    hvx_reduce_loop_body(HVX_Vector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
313}
314
315static inline int32_t hvx_reduce_max_i32_u(const uint8_t * restrict src, const int num_elems) {
316    HVX_Vector init_vec = Q6_V_vsplat_R(((const int32_t *) src)[0]);
317    HVX_Vector pad_vec  = Q6_V_vsplat_R(0x80000000);
318    hvx_reduce_loop_body(HVX_UVector, init_vec, pad_vec, Q6_Vw_vmax_VwVw, hvx_vec_reduce_max_i32, hvx_vec_get_i32);
319}
320
321static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int num_elems) {
322    if (hex_is_aligned((void *) src, 128)) {
323        return hvx_reduce_max_i32_a(src, num_elems);
324    } else {
325        return hvx_reduce_max_i32_u(src, num_elems);
326    }
327}
328
329#undef hvx_reduce_loop_body
330#undef HVX_REDUCE_MAX_OP
331#undef HVX_REDUCE_SUM_OP
332#undef HVX_REDUCE_MAX_SCALAR
333#undef HVX_REDUCE_SUM_SCALAR
334#undef HVX_SUM_SQ_OP
335
336#endif /* HVX_REDUCE_H */