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 →
11.1 KiB303 linesraw
1#ifndef HVX_FA_KERNELS_H
2#define HVX_FA_KERNELS_H
3
4#include <assert.h>
5#include <math.h>
6#include "hvx-utils.h"
7
8// Little inner kernels for HVX
9
10#if __HVX_ARCH__ < 79
11#define HVX_OP_ADD_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b))
12#define HVX_OP_SUB_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b))
13#define HVX_OP_MUL_F32(a, b) Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b))
14#else
15#define HVX_OP_ADD_F32(a, b) Q6_Vsf_vadd_VsfVsf(a, b)
16#define HVX_OP_SUB_F32(a, b) Q6_Vsf_vsub_VsfVsf(a, b)
17#define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b)
18#endif
19
20// This is a bit of a hack because the compiler is struggling to properly inline
21// the default hvx_vec_f32_to_f16 with output into the local array.
22static __attribute__((unused)) __attribute__((noinline)) void hvx_vec_f32_to_f16_a(void *ptr, HVX_Vector v0, HVX_Vector v1)
23{
24    *(HVX_Vector *) ptr = hvx_vec_f32_to_f16(v0, v1);
25}
26
27// Dot product of two F16 vectors, accumulating to float
28static inline void hvx_dot_f16_f16_aa(float * restrict r, const void * restrict x, const void * restrict y, unsigned int n, float s) {
29    const HVX_Vector * restrict vx = (const HVX_Vector * restrict) x; // fp16
30    const HVX_Vector * restrict vy = (const HVX_Vector * restrict) y; // fp16
31
32    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
33    uint32_t nloe = n % VLEN_FP16; // leftover elements
34
35    HVX_VectorPair rsum_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
36
37    uint32_t i = 0;
38
39    #pragma unroll(4)
40    for (i = 0; i < nvec; i++) {
41        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, vx[i], vy[i]);
42    }
43
44    if (nloe) {
45        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
46        HVX_Vector y_hf = Q6_V_vand_QV(bmask, vy[i]);
47        HVX_Vector x_hf = Q6_V_vand_QV(bmask, vx[i]);
48
49        rsum_p = hvx_vec_mpyacc_f32_f16(rsum_p, x_hf, y_hf);
50    }
51
52    HVX_Vector rsum = HVX_OP_ADD_F32(Q6_V_lo_W(rsum_p), Q6_V_hi_W(rsum_p));
53    rsum = HVX_OP_MUL_F32(hvx_vec_splat_f32(s), hvx_vec_reduce_sum_f32(rsum));
54    hvx_vec_store_u(r, 4, rsum);
55}
56
57static inline HVX_Vector hvx_dot_f16_f16_aa_rx4(const void * restrict y,
58                                                const uint8_t * restrict x,
59                                                const size_t stride_x,
60                                                const size_t nvec,
61                                                const size_t nloe) {
62    const HVX_Vector * restrict vx0 = (const HVX_Vector * restrict) x;                   // fp16
63    const HVX_Vector * restrict vx1 = (const HVX_Vector * restrict) (x + stride_x);      // fp16
64    const HVX_Vector * restrict vx2 = (const HVX_Vector * restrict) (x + stride_x * 2);  // fp16
65    const HVX_Vector * restrict vx3 = (const HVX_Vector * restrict) (x + stride_x * 3);  // fp16
66    const HVX_Vector * restrict vy  = (const HVX_Vector * restrict) y;                   // fp16
67
68    HVX_VectorPair rsum0_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
69    HVX_VectorPair rsum1_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
70    HVX_VectorPair rsum2_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
71    HVX_VectorPair rsum3_p = Q6_W_vcombine_VV(Q6_V_vsplat_R(0), Q6_V_vsplat_R(0));
72
73    uint32_t i = 0;
74
75    for (i = 0; i < nvec; i++) {
76        HVX_Vector y_hf  = vy[i];
77        HVX_Vector x0_hf = vx0[i];
78        HVX_Vector x1_hf = vx1[i];
79        HVX_Vector x2_hf = vx2[i];
80        HVX_Vector x3_hf = vx3[i];
81
82        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
83        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
84        rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
85        rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
86    }
87
88    if (nloe) {
89        // Load x (fp16) and zero-out unused elements
90        HVX_VectorPred bmask = Q6_Q_vsetq_R(nloe * 2);
91        HVX_Vector     y_hf  = Q6_V_vand_QV(bmask, vy[i]);
92        HVX_Vector     x0_hf = Q6_V_vand_QV(bmask, vx0[i]);
93        HVX_Vector     x1_hf = Q6_V_vand_QV(bmask, vx1[i]);
94        HVX_Vector     x2_hf = Q6_V_vand_QV(bmask, vx2[i]);
95        HVX_Vector     x3_hf = Q6_V_vand_QV(bmask, vx3[i]);
96
97        rsum0_p = hvx_vec_mpyacc_f32_f16(rsum0_p, x0_hf, y_hf);
98        rsum1_p = hvx_vec_mpyacc_f32_f16(rsum1_p, x1_hf, y_hf);
99        rsum2_p = hvx_vec_mpyacc_f32_f16(rsum2_p, x2_hf, y_hf);
100        rsum3_p = hvx_vec_mpyacc_f32_f16(rsum3_p, x3_hf, y_hf);
101    }
102
103    HVX_Vector rsum0 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum0_p), Q6_V_hi_W(rsum0_p));
104    HVX_Vector rsum1 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum1_p), Q6_V_hi_W(rsum1_p));
105    HVX_Vector rsum2 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum2_p), Q6_V_hi_W(rsum2_p));
106    HVX_Vector rsum3 = HVX_OP_ADD_F32(Q6_V_lo_W(rsum3_p), Q6_V_hi_W(rsum3_p));
107
108    HVX_Vector_x4 rsum0123 = { .v = { rsum0, rsum1, rsum2, rsum3 } };
109    return hvx_vec_reduce_sum_f32x4(rsum0123);
110}
111
112static inline HVX_Vector hvx_dot_f16_f16_aa_rx32(const void * restrict y,
113                                                 const uint8_t * restrict x,
114                                                 const size_t stride_x,
115                                                 const size_t n,
116                                                 float        s) {
117
118    const size_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
119    const size_t nloe = n % VLEN_FP16; // leftover elements
120
121    HVX_Vector   sums = Q6_V_vzero();
122    const size_t stride_x_4 = stride_x * 4;
123    for (uint32_t j = 0; j < VLEN_FP32; j += 4) {
124        HVX_Vector     sums_x4 = hvx_dot_f16_f16_aa_rx4(y, x, stride_x, nvec, nloe);
125        HVX_VectorPred pred    = Q6_Q_vsetq_R(j * SIZEOF_FP32);
126        sums                   = Q6_V_vmux_QVV(pred, sums, sums_x4);
127        x += stride_x_4;
128    }
129
130    return HVX_OP_MUL_F32(hvx_vec_splat_f32(s), sums);
131}
132
133// MAD: y (F32) += x (F16) * s (F16)
134static inline void hvx_mad_f32_f16_aa(float * restrict y, const void * restrict x, const __fp16 * restrict s, uint32_t n) {
135    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
136
137    HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
138    HVX_Vector * restrict vy = (HVX_Vector *) y;
139
140    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
141    uint32_t nloe = n % VLEN_FP16; // leftover elements
142
143    HVX_Vector S0 = hvx_vec_splat_f16(*s);
144
145    uint32_t i = 0;
146
147    #pragma unroll(2)
148    for (i = 0; i < nvec; ++i) {
149        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
150    }
151
152    if (nloe) {
153        HVX_VectorPair xy_p = vy_p[i];
154        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
155
156        HVX_Vector xy = Q6_V_lo_W(xy_p);
157        i = 2 * i;  // index for vy
158
159        if (nloe >= VLEN_FP32) {
160            vy[i] = xy;
161            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
162        }
163
164        if (nloe) {
165            hvx_vec_store_a(&vy[i], nloe * 4, xy);
166        }
167    }
168}
169
170// MAD: y (F32) += x0 (F16) * s0 (F16) + x1 (F16) * s1 (F16)
171static inline void hvx_mad_f32_f16_aa_rx2(float * restrict y, const void * restrict x0, const void * restrict x1,
172                                          const __fp16 * restrict s0, const __fp16 * restrict s1, uint32_t n) {
173    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
174    const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
175
176    HVX_VectorPair * restrict vy_p  = (HVX_VectorPair *) y;
177    HVX_Vector * restrict vy        = (HVX_Vector *) y;
178
179    uint32_t nvec = n / VLEN_FP16;  // num full fp16 hvx vectors
180    uint32_t nloe = n % VLEN_FP16;  // leftover elements
181
182    HVX_Vector S0 = hvx_vec_splat_f16(*s0);
183    HVX_Vector S1 = hvx_vec_splat_f16(*s1);
184
185    uint32_t i = 0;
186
187    #pragma unroll(2)
188    for (i = 0; i < nvec; ++i) {
189        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
190        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
191    }
192
193    if (nloe) {
194        HVX_VectorPair xy_p = vy_p[i];
195        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
196        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
197
198        HVX_Vector xy = Q6_V_lo_W(xy_p);
199        i = 2 * i;  // index for vy
200
201        if (nloe >= VLEN_FP32) {
202            vy[i] = xy;
203            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
204        }
205
206        if (nloe) {
207            hvx_vec_store_a(&vy[i], nloe * 4, xy);
208        }
209    }
210}
211static inline void hvx_mad_f32_f16_aa_vec(float * restrict y, const void * restrict x, HVX_Vector S0, uint32_t n) {
212    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x;
213
214    HVX_VectorPair * restrict vy_p = (HVX_VectorPair *) y;
215    HVX_Vector * restrict vy = (HVX_Vector *) y;
216
217    uint32_t nvec = n / VLEN_FP16; // num full fp16 hvx vectors
218    uint32_t nloe = n % VLEN_FP16; // leftover elements
219
220    uint32_t i = 0;
221
222    #pragma unroll(2)
223    for (i = 0; i < nvec; ++i) {
224        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
225    }
226
227    if (nloe) {
228        HVX_VectorPair xy_p = vy_p[i];
229        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
230
231        HVX_Vector xy = Q6_V_lo_W(xy_p);
232        i = 2 * i;  // index for vy
233
234        if (nloe >= VLEN_FP32) {
235            vy[i] = xy;
236            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
237        }
238
239        if (nloe) {
240            hvx_vec_store_a(&vy[i], nloe * 4, xy);
241        }
242    }
243}
244
245static inline void hvx_mad_f32_f16_aa_rx2_vec(float * restrict y, const void * restrict x0, const void * restrict x1,
246                                          HVX_Vector S0, HVX_Vector S1, uint32_t n) {
247    const HVX_Vector * restrict vx0 = (const HVX_Vector *) x0;
248    const HVX_Vector * restrict vx1 = (const HVX_Vector *) x1;
249
250    HVX_VectorPair * restrict vy_p  = (HVX_VectorPair *) y;
251    HVX_Vector * restrict vy        = (HVX_Vector *) y;
252
253    uint32_t nvec = n / VLEN_FP16;  // num full fp16 hvx vectors
254    uint32_t nloe = n % VLEN_FP16;  // leftover elements
255
256    uint32_t i = 0;
257
258    #pragma unroll(2)
259    for (i = 0; i < nvec; ++i) {
260        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx0[i]), S0);
261        vy_p[i] = hvx_vec_mpyacc_f32_f16(vy_p[i], Q6_Vh_vshuff_Vh(vx1[i]), S1);
262    }
263
264    if (nloe) {
265        HVX_VectorPair xy_p = vy_p[i];
266        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx0[i]), S0);
267        xy_p = hvx_vec_mpyacc_f32_f16(xy_p, Q6_Vh_vshuff_Vh(vx1[i]), S1);
268
269        HVX_Vector xy = Q6_V_lo_W(xy_p);
270        i = 2 * i;  // index for vy
271
272        if (nloe >= VLEN_FP32) {
273            vy[i] = xy;
274            nloe -= VLEN_FP32; ++i; xy = Q6_V_hi_W(xy_p);
275        }
276
277        if (nloe) {
278            hvx_vec_store_a(&vy[i], nloe * 4, xy);
279        }
280    }
281}
282
283static inline void hvx_scale_vec_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, const uint32_t n, HVX_Vector vs) {
284    assert((size_t) dst % 128 == 0);
285    assert((size_t) src % 128 == 0);
286
287    const HVX_Vector * restrict vsrc = (const HVX_Vector * restrict) src;
288    HVX_Vector * restrict vdst       = (HVX_Vector * restrict) dst;
289
290    const uint32_t nvec = n / VLEN_FP32;
291    const uint32_t nloe = n % VLEN_FP32;
292
293    uint32_t i = 0;
294    #pragma unroll(4)
295    for (; i < nvec; ++i) {
296        vdst[i] = HVX_OP_MUL_F32(vsrc[i], vs);
297    }
298    if (nloe) {
299        hvx_vec_store_a(&vdst[i], nloe * sizeof(float), HVX_OP_MUL_F32(vsrc[i], vs));
300    }
301}
302
303#endif /* HVX_FA_KERNELS_H */