char-fork/llama.cpp

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

Max Krasnyanskyhexagon: flash attention rework (optimizations, accuracy improvements, etc) (#25085)13e673863

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
10.1 KiB292 linesraw
1#ifndef HVX_BASE_H
2#define HVX_BASE_H
3
4#include <stdbool.h>
5#include <stdint.h>
6#include <math.h>
7#include <assert.h>
8
9#include "hex-utils.h"
10#include "hvx-types.h"
11
12#define hvx_vmem(A)   *((HVX_Vector *)(A))
13#define hvx_vmemu(A)  *((HVX_UVector *)(A))
14
15static inline void hvx_vec_store_u(void * restrict dst, uint32_t n, HVX_Vector v) {
16    // Rotate as needed.
17    v = Q6_V_vlalign_VVR(v, v, (size_t) dst);
18
19    uint32_t left_off  = (size_t) dst & 127;
20    uint32_t right_off = left_off + n;
21
22    HVX_VectorPred ql_not = Q6_Q_vsetq_R((size_t) dst);
23    HVX_VectorPred qr     = Q6_Q_vsetq2_R(right_off);
24
25    if (right_off > 128) {
26        Q6_vmem_QRIV(qr, (HVX_Vector *) dst + 1, v);
27        // all 1's
28        qr = Q6_Q_vcmp_eq_VbVb(v, v);
29    }
30
31    ql_not = Q6_Q_or_QQn(ql_not, qr);
32    Q6_vmem_QnRIV(ql_not, (HVX_Vector *) dst, v);
33}
34
35static inline void hvx_vec_store_a(void * restrict dst, uint32_t n, HVX_Vector v) {
36    assert((unsigned long) dst % 128 == 0);
37    HVX_VectorPred m = Q6_Q_or_QQn(Q6_Q_vsetq_R((unsigned long) dst), Q6_Q_vsetq2_R(n));
38    Q6_vmem_QnRIV(m, (HVX_Vector *) dst, v);
39}
40
41static inline HVX_Vector hvx_vec_splat_f32(float v) {
42    union { float  f; uint32_t i; } u = { .f = v };
43    return Q6_V_vsplat_R(u.i);
44}
45
46static inline HVX_Vector hvx_vec_splat_f16(_Float16 v) {
47    union { __fp16 f; uint16_t i; } u = { .f = v };
48    return Q6_Vh_vsplat_R(u.i);
49}
50
51static inline HVX_Vector hvx_vec_repl4(HVX_Vector v) {
52    // vdelta control to replicate first 4 bytes across all elements
53    static const uint8_t __attribute__((aligned(128))) repl[128] = {
54        0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
55        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
56        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
57        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
58        0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
59        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
60        0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
61        0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04,
62    };
63
64    HVX_Vector ctrl = *(HVX_Vector *) repl;
65    return Q6_V_vdelta_VV(v, ctrl);
66}
67
68static inline float hvx_vec_get_f32(HVX_Vector v) {
69    float __attribute__((aligned(128))) x;
70    hvx_vec_store_a(&x, 4, v);
71    return x;
72}
73
74static inline int32_t hvx_vec_get_i32(HVX_Vector v) {
75    int32_t __attribute__((aligned(128))) x;
76    hvx_vec_store_a(&x, 4, v);
77    return x;
78}
79
80static inline _Float16 hvx_vec_get_f16(HVX_Vector v) {
81    _Float16 __attribute__((aligned(128))) x;
82    hvx_vec_store_a(&x, 2, v);
83    return x;
84}
85
86static inline HVX_Vector hvx_vec_abs_f16(HVX_Vector v) {
87    // abs by clearing the fp16 sign bit
88    HVX_Vector mask = Q6_Vh_vsplat_R(0x7fff);
89    return Q6_V_vand_VV(v, mask);
90}
91
92static inline HVX_Vector hvx_vec_neg_f16(HVX_Vector v) {
93    // neg by setting the fp16 sign bit
94    HVX_Vector mask = Q6_Vh_vsplat_R(0x8000);
95    return Q6_V_vxor_VV(v, mask);
96}
97
98static inline HVX_Vector hvx_vec_abs_f32(HVX_Vector v) {
99    // abs by clearing the fp32 sign bit
100    HVX_Vector mask = Q6_V_vsplat_R(0x7fffffff);
101    return Q6_V_vand_VV(v, mask);
102}
103
104static inline HVX_Vector hvx_vec_neg_f32(HVX_Vector v) {
105#if __HVX_ARCH__ > 75
106    return Q6_Vsf_vfneg_Vsf(v);
107#else
108    // neg by setting the fp32 sign bit
109    HVX_Vector mask = Q6_V_vsplat_R(0x80000000);
110    return Q6_V_vxor_VV(v, mask);
111#endif  // __HVX_ARCH__ > 75
112}
113
114static inline HVX_VectorPred hvx_vec_is_nan_f16(HVX_Vector v) {
115    const HVX_Vector vnan_exp  = Q6_Vh_vsplat_R(0x7C00);
116    const HVX_Vector vnan_frac = Q6_Vh_vsplat_R(0x7FFF);
117
118    // get pred of which are NaN, i.e., exponent bits all 1s and fraction bits non 0s
119    HVX_VectorPred p_exp  = Q6_Q_vcmp_eq_VhVh(Q6_V_vand_VV(v, vnan_exp), vnan_exp);
120    HVX_VectorPred p_frac = Q6_Q_not_Q(Q6_Q_vcmp_eq_VhVh(Q6_V_vand_VV(v, vnan_frac), vnan_exp));
121    return Q6_Q_and_QQ(p_exp, p_frac);
122}
123
124static inline HVX_Vector hvx_vec_f32_to_f16_shuff(HVX_Vector v0, HVX_Vector v1) {
125#if __HVX_ARCH__ >= 81
126    HVX_Vector q0 = Q6_Vqf32_equals_Vsf(v0);
127    HVX_Vector q1 = Q6_Vqf32_equals_Vsf(v1);
128#else
129    const HVX_Vector zero = Q6_V_vzero();
130    HVX_Vector q0 = Q6_Vqf32_vadd_VsfVsf(v0, zero);
131    HVX_Vector q1 = Q6_Vqf32_vadd_VsfVsf(v1, zero);
132#endif
133    return Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(q1, q0));
134}
135
136static inline HVX_Vector hvx_vec_f32_to_f16(HVX_Vector v0, HVX_Vector v1) {
137    return Q6_Vh_vdeal_Vh(hvx_vec_f32_to_f16_shuff(v0, v1));
138}
139
140#if __HVX_ARCH__ >= 79
141static inline HVX_VectorPair hvx_vec_f16_to_f32_shuff(HVX_Vector v) {
142    const HVX_Vector one = hvx_vec_splat_f16(1.0);
143    HVX_VectorPair p = Q6_Wsf_vmpy_VhfVhf(v, one);
144    return Q6_W_vcombine_VV(Q6_V_hi_W(p), Q6_V_lo_W(p));
145}
146static inline HVX_VectorPair hvx_vec_f16_to_f32(HVX_Vector v) {
147    const HVX_Vector one = hvx_vec_splat_f16(1.0);
148    HVX_VectorPair p = Q6_Wsf_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(v), one);
149    return Q6_W_vcombine_VV(Q6_V_hi_W(p), Q6_V_lo_W(p));
150}
151#else
152static inline HVX_VectorPair hvx_vec_f16_to_f32_shuff(HVX_Vector v) {
153    const HVX_Vector one = hvx_vec_splat_f16(1.0);
154    HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(v, one);
155    return Q6_W_vcombine_VV(Q6_Vsf_equals_Vqf32(Q6_V_hi_W(p)), Q6_Vsf_equals_Vqf32(Q6_V_lo_W(p)));
156}
157static inline HVX_VectorPair hvx_vec_f16_to_f32(HVX_Vector v) {
158    const HVX_Vector one = hvx_vec_splat_f16(1.0);
159    HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(v), one);
160    return Q6_W_vcombine_VV(Q6_Vsf_equals_Vqf32(Q6_V_hi_W(p)), Q6_Vsf_equals_Vqf32(Q6_V_lo_W(p)));
161}
162#endif
163
164static inline HVX_Vector hvx_vec_i16_from_hf_rnd_sat(HVX_Vector vin) {
165    // This looks complicated.
166    // Ideally should just be Q6_Vh_equals_Vhf(vin)
167    // but that instruction does not do proper rounding.
168
169    // convert to qf32, multiplying by 1.0 in the process.
170    HVX_VectorPair v32 = Q6_Wqf32_vmpy_VhfVhf(vin, Q6_Vh_vsplat_R(0x3C00));
171
172    // 'in-range' values are +/32752.
173    // add 192K to it, convert to sf
174    HVX_Vector v192K = Q6_V_vsplat_R(0x48400000);
175    HVX_Vector vsf_0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_V_lo_W(v32), v192K));
176    HVX_Vector vsf_1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_V_hi_W(v32), v192K));
177
178    // for in-range cases, result is {163858... 229360} so the exponent is always 144.
179    // if we extract bits 21..0 as a signed quantity, and round 6 bits off, that will be the answer.
180    // Start by <<10 to get the final 'sign' bit in bit 15...
181    vsf_0 = Q6_Vw_vasl_VwR(vsf_0, 10);
182    vsf_1 = Q6_Vw_vasl_VwR(vsf_1, 10);
183
184    // now round down to 16
185    return Q6_Vh_vround_VwVw_sat(vsf_1, vsf_0);
186}
187
188#if __HVX_ARCH__ < 79
189
190static inline HVX_VectorPair hvx_vec_mpyacc_f32_f16(HVX_VectorPair acc, HVX_Vector x, HVX_Vector y)
191{
192    HVX_VectorPair m = Q6_Wqf32_vmpy_VhfVhf(x, y);
193    HVX_Vector a0 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_V_lo_W(m), Q6_V_lo_W(acc)));
194    HVX_Vector a1 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vsf(Q6_V_hi_W(m), Q6_V_hi_W(acc)));
195    return Q6_W_vcombine_VV(a1, a0);
196}
197
198#else
199
200static inline HVX_VectorPair hvx_vec_mpyacc_f32_f16(HVX_VectorPair acc, HVX_Vector x, HVX_Vector y)
201{
202    return Q6_Wsf_vmpyacc_WsfVhfVhf(acc, x, y);
203}
204
205#endif
206
207#if __HVX_ARCH__ < 79
208
209static inline HVX_Vector hvx_vec_add_f16_f16(HVX_Vector a, HVX_Vector b)
210{
211    const HVX_Vector negone = Q6_Vh_vsplat_R(0xBC00); // -1.0 in IEEE FP16
212    const HVX_Vector one    = Q6_Vh_vsplat_R(0x3C00); //  1.0 in IEEE FP16
213    HVX_VectorPair a_p = Q6_Wqf32_vmpy_VhfVhf(a, one);
214    HVX_VectorPair b_p = Q6_Wqf32_vmpy_VhfVhf(b, negone);
215    HVX_Vector a0 = Q6_Vqf32_vsub_Vqf32Vqf32(Q6_V_lo_W(a_p), Q6_V_lo_W(b_p));
216    HVX_Vector a1 = Q6_Vqf32_vsub_Vqf32Vqf32(Q6_V_hi_W(a_p), Q6_V_hi_W(b_p));
217    return Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(a1, a0));
218}
219
220static inline HVX_Vector hvx_vec_sub_f16_f16(HVX_Vector a, HVX_Vector b)
221{
222    const HVX_Vector negone = Q6_Vh_vsplat_R(0xBC00); // -1.0 in IEEE FP16
223    const HVX_Vector one    = Q6_Vh_vsplat_R(0x3C00); //  1.0 in IEEE FP16
224    HVX_VectorPair a_p = Q6_Wqf32_vmpy_VhfVhf(a, one);
225    HVX_VectorPair b_p = Q6_Wqf32_vmpy_VhfVhf(b, negone);
226    HVX_Vector a0 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_lo_W(a_p), Q6_V_lo_W(b_p));
227    HVX_Vector a1 = Q6_Vqf32_vadd_Vqf32Vqf32(Q6_V_hi_W(a_p), Q6_V_hi_W(b_p));
228    return Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(a1, a0));
229}
230
231static inline HVX_Vector hvx_vec_mul_f16_f16(HVX_Vector a, HVX_Vector b)
232{
233    return Q6_Vhf_equals_Wqf32(Q6_Wqf32_vmpy_VhfVhf(a, b));
234}
235
236static inline HVX_Vector hvx_vec_add_f32_f32(HVX_Vector a, HVX_Vector b) {
237    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_VsfVsf(a, b));
238}
239
240static inline HVX_Vector hvx_vec_sub_f32_f32(HVX_Vector a, HVX_Vector b) {
241    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(a, b));
242}
243
244static inline HVX_Vector hvx_vec_mul_f32_f32(HVX_Vector a, HVX_Vector b) {
245    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(a, b));
246}
247
248#else
249
250static inline HVX_Vector hvx_vec_add_f16_f16(HVX_Vector a, HVX_Vector b)
251{
252    return Q6_Vhf_vadd_VhfVhf(a, b);
253}
254
255static inline HVX_Vector hvx_vec_sub_f16_f16(HVX_Vector a, HVX_Vector b)
256{
257    return Q6_Vhf_vsub_VhfVhf(a, b);
258}
259
260static inline HVX_Vector hvx_vec_mul_f16_f16(HVX_Vector a, HVX_Vector b)
261{
262    return Q6_Vhf_vmpy_VhfVhf(a, b);
263}
264
265static inline HVX_Vector hvx_vec_add_f32_f32(HVX_Vector a, HVX_Vector b) {
266    return Q6_Vsf_vadd_VsfVsf(a, b);
267}
268
269static inline HVX_Vector hvx_vec_sub_f32_f32(HVX_Vector a, HVX_Vector b) {
270    return Q6_Vsf_vsub_VsfVsf(a, b);
271}
272
273static inline HVX_Vector hvx_vec_mul_f32_f32(HVX_Vector a, HVX_Vector b) {
274    return Q6_Vsf_vmpy_VsfVsf(a, b);
275}
276
277#endif // __HVX_ARCH__ < 79
278
279static inline HVX_Vector hvx_vec_load_act_tile(const uint8_t * y_q, uint32_t kt, HVX_Vector * v_act_all) {
280    if (kt % 4 == 0) {
281        *v_act_all = hvx_vmem(y_q + kt * 32);
282        return *v_act_all;
283    } else if (kt % 4 == 1) {
284        return Q6_V_vror_VR(*v_act_all, 32);
285    } else if (kt % 4 == 2) {
286        return Q6_V_vror_VR(*v_act_all, 64);
287    } else {
288        return Q6_V_vror_VR(*v_act_all, 96);
289    }
290}
291
292#endif /* HVX_BASE_H */