char-fork/llama.cpp

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

Yiwei Shaohexagon: hmx flash attention (#22347)1a03cf47f

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
14.0 KiB262 linesraw
1#ifndef HVX_COPY_H
2#define HVX_COPY_H
3
4#include <assert.h>
5#include <stddef.h>
6#include <stdint.h>
7
8#include "hvx-base.h"
9
10#define hvx_splat_pragma(x) _Pragma(#x)
11#define hvx_splat_loop_body(dst_type, vec_store, unroll_cnt)     \
12    do {                                                         \
13        dst_type * restrict vdst = (dst_type *) dst;             \
14                                                                 \
15        uint32_t nvec = n / (128 / elem_size);                   \
16        uint32_t nloe = n % (128 / elem_size);                   \
17                                                                 \
18        uint32_t i = 0;                                          \
19                                                                 \
20        hvx_splat_pragma(unroll(unroll_cnt))                     \
21        for (; i < nvec; i++) {                                  \
22            vdst[i] = src;                                       \
23        }                                                        \
24        if (nloe) {                                              \
25            vec_store((void *) &vdst[i], nloe * elem_size, src); \
26        }                                                        \
27    } while(0)
28
29static inline void hvx_splat_a(void * restrict dst, HVX_Vector src, uint32_t n, uint32_t elem_size) {
30    assert((unsigned long) dst % 128 == 0);
31    hvx_splat_loop_body(HVX_Vector, hvx_vec_store_a, 4);
32}
33
34static inline void hvx_splat_u(void * restrict dst, HVX_Vector src, uint32_t n, uint32_t elem_size) {
35    hvx_splat_loop_body(HVX_UVector, hvx_vec_store_u, 4);
36}
37
38static inline void hvx_splat_f32_a(void * restrict dst, float v, uint32_t n) {
39    hvx_splat_a(dst,  hvx_vec_splat_f32(v), n, sizeof(float));
40}
41
42static inline void hvx_splat_f32_u(void * restrict dst, float v, uint32_t n) {
43    hvx_splat_u(dst,  hvx_vec_splat_f32(v), n, sizeof(float));
44}
45
46static inline void hvx_splat_f16_a(void * restrict dst, _Float16 v, uint32_t n) {
47    hvx_splat_u(dst,  hvx_vec_splat_f16(v), n, sizeof(__fp16));
48}
49
50static inline void hvx_splat_f16_u(void * restrict dst, _Float16 v, uint32_t n) {
51    hvx_splat_u(dst,  hvx_vec_splat_f16(v), n, sizeof(__fp16));
52}
53
54static inline void hvx_splat_u16_a(void * restrict dst, uint16_t v, uint32_t n) {
55    hvx_splat_a(dst,  Q6_Vh_vsplat_R(v), n, sizeof(uint16_t));
56}
57
58static inline void hvx_splat_u16_u(void * restrict dst, uint16_t v, uint32_t n) {
59    hvx_splat_u(dst,  Q6_Vh_vsplat_R(v), n, sizeof(uint16_t));
60}
61
62static inline void hvx_splat_u8_a(void * restrict dst, uint8_t v, uint32_t n) {
63    hvx_splat_a(dst,  Q6_Vb_vsplat_R(v), n, 1);
64}
65
66static inline void hvx_splat_u8_u(void * restrict dst, uint8_t v, uint32_t n) {
67    hvx_splat_u(dst,  Q6_Vb_vsplat_R(v), n, 1);
68}
69
70#define hvx_copy_loop_body(dst_type, src_type, vec_store)            \
71    do {                                                             \
72        dst_type * restrict vdst = (dst_type *) dst;                 \
73        src_type * restrict vsrc = (src_type *) src;                 \
74                                                                     \
75        const uint32_t epv  = 128 / elem_size;                       \
76        const uint32_t nvec = n / epv;                               \
77        const uint32_t nloe = n % epv;                               \
78                                                                     \
79        uint32_t i = 0;                                              \
80                                                                     \
81        _Pragma("unroll(4)")                                         \
82        for (; i < nvec; i++) { vdst[i] = vsrc[i]; }                 \
83        if (nloe) {                                                  \
84            vec_store((void *) &vdst[i], nloe * elem_size, vsrc[i]); \
85        }                                                            \
86    } while(0)
87
88// Generic copy routines
89static inline void hvx_copy_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) {
90    assert((unsigned long) dst % 128 == 0);
91    assert((unsigned long) src % 128 == 0);
92    hvx_copy_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
93}
94
95static inline void hvx_copy_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) {
96    assert((unsigned long) dst % 128 == 0);
97    hvx_copy_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
98}
99
100static inline void hvx_copy_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) {
101    assert((unsigned long) src % 128 == 0);
102    hvx_copy_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
103}
104
105static inline void hvx_copy_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n, uint32_t elem_size) {
106    hvx_copy_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
107}
108
109// copy n fp16 elements : source and destination are aligned to HVX Vector (128)
110static inline void hvx_copy_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
111    hvx_copy_aa(dst, src, n, sizeof(__fp16));
112}
113
114// copy n fp16 elements : source is aligned, destination is potentially unaligned
115static inline void hvx_copy_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
116    hvx_copy_au(dst, src, n, sizeof(__fp16));
117}
118
119// copy n fp16 elements : source is aligned, destination is potentially unaligned
120static inline void hvx_copy_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
121    hvx_copy_ua(dst, src, n, sizeof(__fp16));
122}
123
124// copy n fp16 elements : source is aligned, destination is potentially unaligned
125static inline void hvx_copy_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
126    hvx_copy_uu(dst, src, n, sizeof(__fp16));
127}
128
129// copy n fp32 elements : source and destination are aligned to HVX Vector (128)
130static inline void hvx_copy_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
131    hvx_copy_aa(dst, src, n, sizeof(float));
132}
133
134// copy n fp32 elements : source is aligned, destination is unaligned
135static inline void hvx_copy_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
136    hvx_copy_ua(dst, src, n, sizeof(float));
137}
138
139// copy n fp32 elements : source is unaligned, destination is aligned
140static inline void hvx_copy_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
141    hvx_copy_au(dst, src, n, sizeof(float));
142}
143
144// copy n fp32 elements : source is unaligned, destination unaligned
145static inline void hvx_copy_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
146    hvx_copy_uu(dst, src, n, sizeof(float));
147}
148
149//// fp32 -> fp16
150
151#define hvx_copy_f16_f32_loop_body(dst_type, src_type, vec_store)                   \
152    do {                                                                            \
153        dst_type * restrict vdst = (dst_type *) dst;                                \
154        src_type * restrict vsrc = (src_type *) src;                                \
155                                                                                    \
156        const uint32_t elem_size = sizeof(__fp16);                                  \
157        const uint32_t epv  = 128 / elem_size;                                      \
158        const uint32_t nvec = n / epv;                                              \
159        const uint32_t nloe = n % epv;                                              \
160                                                                                    \
161        uint32_t i = 0;                                                             \
162                                                                                    \
163        _Pragma("unroll(4)")                                                        \
164        for (; i < nvec; i++) {                                                     \
165            vdst[i] = hvx_vec_f32_to_f16(vsrc[i*2+0], vsrc[i*2+1]);                 \
166        }                                                                           \
167        if (nloe) {                                                                 \
168            HVX_Vector v = hvx_vec_f32_to_f16(vsrc[i*2+0], vsrc[i*2+1]);            \
169            vec_store((void *) &vdst[i], nloe * elem_size, v);                      \
170        }                                                                           \
171    } while(0)
172
173// copy/convert n fp32 elements into n fp16 elements : source is aligned, destination is aligned
174static inline void hvx_copy_f16_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
175    assert((unsigned long) dst % 128 == 0);
176    assert((unsigned long) src % 128 == 0);
177    hvx_copy_f16_f32_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
178}
179
180// copy/convert n fp32 elements into n fp16 elements : source is unaligned, destination is aligned
181static inline void hvx_copy_f16_f32_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
182    assert((unsigned long) dst % 128 == 0);
183    hvx_copy_f16_f32_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
184}
185
186// copy/convert n fp32 elements into n fp16 elements : source is aligned, destination is unaligned
187static inline void hvx_copy_f16_f32_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
188    assert((unsigned long) src % 128 == 0);
189    hvx_copy_f16_f32_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
190}
191
192// copy/convert n fp32 elements into n fp16 elements : source is unaligned, destination is unaligned
193static inline void hvx_copy_f16_f32_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
194    hvx_copy_f16_f32_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
195}
196
197//// fp16 -> fp32
198
199#define hvx_copy_f32_f16_loop_body(dst_type, src_type, vec_store)                   \
200    do {                                                                            \
201        dst_type * restrict vdst = (dst_type *) dst;                                \
202        src_type * restrict vsrc = (src_type *) src;                                \
203                                                                                    \
204        const HVX_Vector one = hvx_vec_splat_f16(1.0);                              \
205                                                                                    \
206        const uint32_t elem_size = sizeof(__fp16);                                  \
207        const uint32_t epv  = 128 / elem_size;                                      \
208        const uint32_t nvec = n / epv;                                              \
209              uint32_t nloe = n % epv;                                              \
210                                                                                    \
211        uint32_t i = 0;                                                             \
212                                                                                    \
213        _Pragma("unroll(4)")                                                        \
214        for (i = 0; i < nvec; ++i) {                                                \
215            HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(vsrc[i]), one); \
216            vdst[i*2]   = Q6_Vsf_equals_Vqf32(Q6_V_lo_W(p));                        \
217            vdst[i*2+1] = Q6_Vsf_equals_Vqf32(Q6_V_hi_W(p));                        \
218        }                                                                           \
219                                                                                    \
220        if (nloe) {                                                                 \
221            HVX_VectorPair p = Q6_Wqf32_vmpy_VhfVhf(Q6_Vh_vshuff_Vh(vsrc[i]), one); \
222                                                                                    \
223            HVX_Vector vd = Q6_V_lo_W(p);                                           \
224            i = 2 * i;                                                              \
225                                                                                    \
226            if (nloe >= 32) {                                                       \
227                vdst[i] = Q6_Vsf_equals_Vqf32(vd);                                  \
228                nloe -= 32; ++i; vd = Q6_V_hi_W(p);                                 \
229            }                                                                       \
230                                                                                    \
231            if (nloe) {                                                             \
232                vd = Q6_Vsf_equals_Vqf32(vd);                                       \
233                hvx_vec_store_u(&vdst[i], nloe * sizeof(float), vd);                \
234            }                                                                       \
235        }                                                                           \
236    } while(0)
237
238// copy/convert n fp16 elements into n fp32 elements : source is aligned, destination is aligned
239static inline void hvx_copy_f32_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
240    assert((unsigned long) dst % 128 == 0);
241    assert((unsigned long) src % 128 == 0);
242    hvx_copy_f32_f16_loop_body(HVX_Vector, HVX_Vector, hvx_vec_store_a);
243}
244
245// copy/convert n fp16 elements into n fp32 elements : source is unaligned, destination is aligned
246static inline void hvx_copy_f32_f16_au(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
247    assert((unsigned long) dst % 128 == 0);
248    hvx_copy_f32_f16_loop_body(HVX_Vector, HVX_UVector, hvx_vec_store_a);
249}
250
251// copy/convert n fp16 elements into n fp32 elements : source is aligned, destination is unaligned
252static inline void hvx_copy_f32_f16_ua(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
253    assert((unsigned long) src % 128 == 0);
254    hvx_copy_f32_f16_loop_body(HVX_UVector, HVX_Vector, hvx_vec_store_u);
255}
256
257// copy/convert n fp16 elements into n fp32 elements : source is unaligned, destination is unaligned
258static inline void hvx_copy_f32_f16_uu(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
259    hvx_copy_f32_f16_loop_body(HVX_UVector, HVX_UVector, hvx_vec_store_u);
260}
261
262#endif // HVX_COPY_H