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 →
32.9 KiB602 linesraw
1#pragma clang diagnostic ignored "-Wunused-variable"
2#pragma clang diagnostic ignored "-Wunused-function"
3#pragma clang diagnostic ignored "-Wunused-but-set-variable"
4
5#include <HAP_farf.h>
6
7#include <math.h>
8#include <string.h>
9
10#include "hex-dma.h"
11#include "hvx-utils.h"
12
13#define GGML_COMMON_DECL_C
14#include "ggml-common.h"
15#include "htp-ctx.h"
16#include "htp-ops.h"
17#include "hex-common.h"
18#include "htp-tensor.h"
19#include "htp-vtcm.h"
20
21#define htp_act_preamble                                 \
22    const struct htp_tensor * src0 = actx->octx->src[0]; \
23    const struct htp_tensor * src1 = actx->octx->src[1]; \
24    const struct htp_tensor * dst  = actx->octx->dst;    \
25                                                         \
26    const uint32_t ne00 = src0->ne[0];                   \
27    const uint32_t ne01 = src0->ne[1];                   \
28    const uint32_t ne02 = src0->ne[2];                   \
29    const uint32_t ne03 = src0->ne[3];                   \
30                                                         \
31    const uint32_t nb00 = src0->nb[0];                   \
32    const uint32_t nb01 = src0->nb[1];                   \
33    const uint32_t nb02 = src0->nb[2];                   \
34    const uint32_t nb03 = src0->nb[3];                   \
35                                                         \
36    const uint32_t ne10 = src1 ? src1->ne[0] : 0;        \
37    const uint32_t ne11 = src1 ? src1->ne[1] : 0;        \
38    const uint32_t ne12 = src1 ? src1->ne[2] : 0;        \
39    const uint32_t ne13 = src1 ? src1->ne[3] : 0;        \
40                                                         \
41    const uint32_t nb10 = src1 ? src1->nb[0] : 0;        \
42    const uint32_t nb11 = src1 ? src1->nb[1] : 0;        \
43    const uint32_t nb12 = src1 ? src1->nb[2] : 0;        \
44    const uint32_t nb13 = src1 ? src1->nb[3] : 0;        \
45                                                         \
46    const uint32_t ne0 = dst->ne[0];                     \
47    const uint32_t ne1 = dst->ne[1];                     \
48    const uint32_t ne2 = dst->ne[2];                     \
49    const uint32_t ne3 = dst->ne[3];                     \
50                                                         \
51    const uint32_t nb0 = dst->nb[0];                     \
52    const uint32_t nb1 = dst->nb[1];                     \
53    const uint32_t nb2 = dst->nb[2];                     \
54    const uint32_t nb3 = dst->nb[3];
55
56struct htp_act_context {
57    struct htp_ops_context * octx;
58
59    // Precomputed values
60    const uint8_t *          data_src0;
61    const uint8_t *          data_src1;
62    uint8_t *                data_dst;
63
64    size_t                   src0_row_size;
65    size_t                   src1_row_size;
66    size_t                   dst_row_size;
67
68    size_t                   src0_row_stride;
69    size_t                   src1_row_stride;
70
71    size_t                   src0_row_size_aligned;
72    size_t                   src1_row_size_aligned;
73    size_t                   dst_row_size_aligned;
74
75    size_t                   src0_spad_half_size;
76    size_t                   src1_spad_half_size;
77    size_t                   dst_spad_half_size;
78
79    uint32_t                 block;
80    uint32_t                 src0_nrows;
81    uint32_t                 src0_nrows_per_thread;
82    uint32_t                 row_start;
83    int                      nc;
84
85    uint8_t *                vtcm_src0;
86    uint8_t *                vtcm_src1;
87    uint8_t *                vtcm_dst;
88
89    size_t                   vtcm_src0_size_per_thread;
90    size_t                   vtcm_src1_size_per_thread;
91    size_t                   vtcm_dst_size_per_thread;
92};
93
94struct htp_act_vtcm_layout {
95    size_t total_bytes;
96    size_t off_src0;
97    size_t off_src1;
98    size_t off_dst;
99
100    size_t src0_bytes_per_thread;
101    size_t src1_bytes_per_thread;
102    size_t dst_bytes_per_thread;
103
104    uint32_t vtcm_row_per_thread;
105};
106
107static inline void htp_act_vtcm_layout_build(struct htp_act_vtcm_layout * L,
108                                             size_t                       src0_row_size_aligned,
109                                             size_t                       src1_row_size_aligned,
110                                             size_t                       dst_row_size_aligned,
111                                             uint32_t                     n_threads,
112                                             size_t                       vtcm_size) {
113    const size_t   spad_size_per_row   = src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned;
114    const uint32_t vtcm_row_per_thread = (uint32_t) (vtcm_size / (n_threads * spad_size_per_row));
115
116    L->vtcm_row_per_thread = vtcm_row_per_thread;
117
118    L->src0_bytes_per_thread = src0_row_size_aligned * vtcm_row_per_thread;
119    L->src1_bytes_per_thread = src1_row_size_aligned * vtcm_row_per_thread;
120    L->dst_bytes_per_thread  = dst_row_size_aligned * vtcm_row_per_thread;
121
122    L->off_src0 = 0;
123    L->off_src1 = L->off_src0 + L->src0_bytes_per_thread * n_threads;
124    L->off_dst  = L->off_src1 + L->src1_bytes_per_thread * n_threads;
125
126    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
127}
128
129#define htp_glu_op_preamble                                            \
130    const size_t src0_row_size_aligned = actx->src0_row_size_aligned;  \
131    const size_t src1_row_size_aligned = actx->src1_row_size_aligned;  \
132    const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;   \
133    const int    nc                    = actx->nc;
134
135// swiglu(x) = x1 * sigmoid(x0)
136static void swiglu_f32(const float * restrict src0,
137                       const float * restrict src1,
138                       float * restrict dst,
139                       const uint32_t num_rows,
140                       const struct htp_act_context * actx) {
141    htp_glu_op_preamble;
142
143    for (uint32_t ib = 0; ib < num_rows; ib++) {
144        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
145        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
146        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
147
148        hvx_sigmoid_f32_aa(dst_ptr, src0_ptr, nc);
149        hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
150    }
151}
152
153// out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
154static void swiglu_oai_f32(const float * restrict src0,
155                           const float * restrict src1,
156                           float * restrict dst,
157                           const uint32_t num_rows,
158                           const struct htp_act_context * actx) {
159    htp_glu_op_preamble;
160    const float alpha = ((const float *) (actx->octx->op_params))[2];
161    const float limit = ((const float *) (actx->octx->op_params))[3];
162
163    for (uint32_t ib = 0; ib < num_rows; ib++) {
164        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
165        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
166        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
167
168        // x (src0_ptr) = std::min(src0_p[k], limit);
169        hvx_min_scalar_f32((uint8_t *) src0_ptr, src0_ptr, limit, nc);
170        // y1 (src1_ptr) = std::clamp(src1_p[k], -limit, limit);
171        hvx_clamp_scalar_f32((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
172        // y (src1_ptr) = y1 + 1.f
173        hvx_add_scalar_f32((uint8_t *) src1_ptr, src1_ptr, 1.0, nc);
174        // x1 (dst_ptr) = alpha * x
175        hvx_mul_scalar_f32(dst_ptr, src0_ptr, alpha, nc);
176        // x2 (dst_ptr) = sigmoid(x1) = 1/(1+exp(-x1))
177        hvx_sigmoid_f32_aa(dst_ptr, dst_ptr, nc);
178        // out = x * sigmoid(alpha * x) * (y + 1.f)
179        hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
180    }
181}
182
183static void swiglu_clamp_f32(const float * restrict src0,
184                             const float * restrict src1,
185                             float * restrict dst,
186                             const uint32_t                 num_rows,
187                             const struct htp_act_context * actx) {
188    htp_glu_op_preamble;
189    const float limit = ((const float *) (actx->octx->op_params))[3];
190
191    for (uint32_t ib = 0; ib < num_rows; ib++) {
192        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
193        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
194        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
195
196        hvx_min_scalar_f32((uint8_t *) src0_ptr, src0_ptr, limit, nc);
197        hvx_clamp_scalar_f32((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
198        hvx_sigmoid_f32_aa(dst_ptr, src0_ptr, nc);
199        hvx_mul_mul_f32_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
200    }
201}
202
203static const float GELU_COEF_A     = 0.044715f;
204static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
205
206static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
207    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
208    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
209
210    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
211    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
212    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
213
214    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
215    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
216
217    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
218    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
219    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
220
221    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
222    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
223
224    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
225    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
226
227    // Newton-Raphson with 2 iterations
228    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
229    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
230    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
231        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
232    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
233        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
234    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
235
236    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
237
238    return Q6_Vsf_equals_Vqf32(res);
239}
240
241static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
242                                                            HVX_Vector one,
243                                                            HVX_Vector max_exp,
244                                                            HVX_Vector min_exp) {
245    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
246    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
247
248    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
249    out            = Q6_V_vmux_QVV(pred_max, out, one);
250    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
251}
252
253static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
254    assert((unsigned long) dst  % 128 == 0);
255    assert((unsigned long) src0 % 128 == 0);
256    assert((unsigned long) src1 % 128 == 0);
257
258    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
259    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
260    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
261
262    const uint32_t epv  = 128 / sizeof(float);
263    const uint32_t nvec = n / epv;
264    const uint32_t nloe = n % epv;
265
266    const float GELU_COEF_A_TIMES_SQRT = GELU_COEF_A * SQRT_2_OVER_PI;
267
268    const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT);
269    const HVX_Vector v_sqrt_2_pi         = hvx_vec_splat_f32(SQRT_2_OVER_PI);
270    const HVX_Vector v_one               = hvx_vec_splat_f32(1.0f);
271    const HVX_Vector v_max_exp           = hvx_vec_splat_f32(87.0f);
272    const HVX_Vector v_min_exp           = hvx_vec_splat_f32(-87.0f);
273
274    uint32_t i = 0;
275
276    _Pragma("unroll(4)")
277    for (; i < nvec; i++) {
278        HVX_Vector x = vsrc0[i];
279        HVX_Vector g = vsrc1[i];
280
281        HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
282        HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
283        coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
284        HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
285
286        // y2 = 2 * inner = inner + inner
287        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
288
289        // Fast sigmoid approximation (2 iterations)
290        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
291
292        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
293        vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g);
294    }
295
296    if (nloe) {
297        HVX_Vector x = vsrc0[i];
298        HVX_Vector g = vsrc1[i];
299
300        HVX_Vector x2 = hvx_vec_mul_f32_f32(x, x);
301        HVX_Vector coef = hvx_vec_mul_f32_f32(x2, v_coef_a_times_sqrt);
302        coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi);
303        HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef);
304
305        HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner);
306
307        HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp);
308
309        HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y);
310        HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g);
311        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res);
312    }
313}
314
315// geglu(x, g) = gelu(x) * g
316static void geglu_f32(const float * restrict src0,
317                      const float * restrict src1,
318                      float * restrict dst,
319                      const uint32_t num_rows,
320                      const struct htp_act_context * actx) {
321    htp_glu_op_preamble;
322
323    for (uint32_t ib = 0; ib < num_rows; ib++) {
324        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
325        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
326        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
327
328        hvx_geglu_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
329    }
330}
331
332#define DEFINE_GLU_PER_THREAD(NAME, OP_STR, CORE_EXPR)                                                                   \
333    static void glu_##NAME##_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {                           \
334        struct htp_act_context * actx = (struct htp_act_context *) data;                                                 \
335        htp_act_preamble;                                                                                                \
336                                                                                                                         \
337        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                            \
338                                                                                                                         \
339        size_t src0_row_size = actx->src0_row_size;                                                                      \
340        size_t src1_row_size = actx->src1_row_size;                                                                      \
341        size_t dst_row_size  = actx->dst_row_size;                                                                       \
342                                                                                                                         \
343        size_t src0_row_stride = actx->src0_row_stride;                                                                  \
344        size_t src1_row_stride = actx->src1_row_stride;                                                                  \
345                                                                                                                         \
346        const uint32_t src0_nrows            = actx->src0_nrows;                                                         \
347        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                              \
348                                                                                                                         \
349        const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;                                   \
350        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);       \
351                                                                                                                         \
352        /* no work for this thread */                                                                                    \
353        if (src0_start_row >= src0_end_row) {                                                                            \
354            return;                                                                                                      \
355        }                                                                                                                \
356                                                                                                                         \
357        const uint8_t * restrict data_src0 = actx->data_src0;                                                            \
358        const uint8_t * restrict data_src1 = actx->data_src1;                                                            \
359        uint8_t * restrict data_dst        = actx->data_dst;                                                             \
360                                                                                                                         \
361        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                                \
362        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                                \
363        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                                 \
364                                                                                                                         \
365        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                   \
366        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                   \
367        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                    \
368                                                                                                                         \
369        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                          \
370        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                          \
371        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                           \
372                                                                                                                         \
373        const int BLOCK = actx->block;                                                                                   \
374        if (BLOCK == 0) {                                                                                                \
375            FARF(ERROR,                                                                                                  \
376                 OP_STR                                                                                                  \
377                 " : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n",        \
378                 actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                                \
379            return;                                                                                                      \
380        }                                                                                                                \
381                                                                                                                         \
382        dma_queue * dma_queue = actx->octx->ctx->dma[ith];                                                               \
383                                                                                                                         \
384        /* See discussion: https://github.com/ggml-org/llama.cpp/pull/18151#issuecomment-3678235379 */                   \
385        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {   \
386            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
387                                                                                                                         \
388            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                     \
389            dma_queue_push_vtcm_to_ddr(dma_queue,                                                                        \
390                                       dma_make_ptr(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),          \
391                                       dst_row_size, dst_row_size_aligned, 0);                                           \
392                                                                                                                         \
393            dma_queue_push(                                                                                              \
394                dma_queue,                                                                                               \
395                dma_make_ptr(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),     \
396                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                      \
397            dma_queue_push(                                                                                              \
398                dma_queue,                                                                                               \
399                dma_make_ptr(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),     \
400                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                      \
401        }                                                                                                                \
402                                                                                                                         \
403        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                             \
404            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                   \
405                                                                                                                         \
406            float * dst_spad  = (float *) dma_queue_pop(dma_queue).src;                                                  \
407            float * src0_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
408            float * src1_spad = (float *) dma_queue_pop(dma_queue).dst;                                                  \
409                                                                                                                         \
410            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                       \
411            CORE_EXPR;                                                                                                   \
412            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                        \
413                                                                                                                         \
414            dma_queue_push_vtcm_to_ddr(dma_queue, dma_make_ptr(data_dst + (ir * dst_row_size), dst_spad),                \
415                                       dst_row_size, dst_row_size_aligned, block_size);                                  \
416                                                                                                                         \
417            /* prefetch N+2 loop iteration if any */                                                                     \
418            const uint32_t pref_block = (ir + BLOCK * 2);                                                                \
419            if (pref_block < src0_end_row) {                                                                             \
420                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                  \
421                dma_queue_push(dma_queue, dma_make_ptr(src0_spad, data_src0 + (pref_block * src0_row_stride)),           \
422                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                  \
423                dma_queue_push(dma_queue, dma_make_ptr(src1_spad, data_src1 + (pref_block * src1_row_stride)),           \
424                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                  \
425            }                                                                                                            \
426        }                                                                                                                \
427                                                                                                                         \
428        dma_queue_flush(dma_queue);                                                                                      \
429                                                                                                                         \
430    }
431
432DEFINE_GLU_PER_THREAD(swiglu, "swiglu-f32", swiglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
433DEFINE_GLU_PER_THREAD(swiglu_oai, "swiglu-oai-f32", swiglu_oai_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
434DEFINE_GLU_PER_THREAD(swiglu_clamp, "swiglu-clamp-f32", swiglu_clamp_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
435DEFINE_GLU_PER_THREAD(geglu, "geglu-f32", geglu_f32(src0_spad, src1_spad, dst_spad, block_size, actx))
436
437static int execute_op_activations_f32(struct htp_ops_context * octx) {
438    const struct htp_tensor * src0 = octx->src[0];
439    const struct htp_tensor * src1 = octx->src[1];
440    const struct htp_tensor * dst  = octx->dst;
441
442    if ((dst->ne[0] * SIZEOF_FP32) != dst->nb[1]) {
443        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
444        return HTP_STATUS_NO_SUPPORT;
445    }
446
447    worker_callback_t act_op_func;
448    const char *      op_type = NULL;
449
450    switch (octx->op) {
451        case HTP_OP_GLU_SWIGLU:
452            act_op_func = (worker_callback_t)glu_swiglu_f32_per_thread;
453            op_type     = "swiglu-f32";
454            break;
455
456        case HTP_OP_GLU_SWIGLU_OAI:
457            act_op_func = (worker_callback_t)glu_swiglu_oai_f32_per_thread;
458            op_type     = "swiglu-oai-f32";
459            break;
460
461        case HTP_OP_GLU_SWIGLU_CLAMP:
462            act_op_func = (worker_callback_t) glu_swiglu_clamp_f32_per_thread;
463            op_type     = "swiglu-clamp-f32";
464            break;
465
466        case HTP_OP_GLU_GEGLU:
467            act_op_func = (worker_callback_t)glu_geglu_f32_per_thread;
468            op_type     = "geglu-f32";
469            break;
470        default:
471            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
472            return HTP_STATUS_NO_SUPPORT;
473    }
474
475    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
476    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP32;
477
478    uint32_t row_start = 0;
479    uint32_t nrows     = src0_nrows;
480
481    if (octx->ctx->mdev.count > 1) {
482        uint32_t rows_per_chunk = 0;
483        htp_tensor_mdev_rows_per_chunk(dst, sizeof(float), (uint32_t) dst_row_size, &rows_per_chunk);
484        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
485        row_start = range.start;
486        nrows     = range.count;
487    }
488
489    if (nrows == 0) {
490        return HTP_STATUS_OK;
491    }
492
493    const uint32_t n_threads = octx->n_threads;
494
495    // row_size   = bytes of useful data per row (what the kernel touches / what DMA copies).
496    // row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
497    const size_t nc_bytes        = dst_row_size;
498    const size_t src0_row_size   = nc_bytes;
499    const size_t src1_row_size   = nc_bytes;
500    const size_t src0_row_stride = src0->nb[1];
501    const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
502
503    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
504    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
505    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
506
507    struct htp_act_vtcm_layout L;
508    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
509                              octx->ctx->vtcm_size);
510
511    // Make sure the reserved vtcm size is sufficient
512    if (L.vtcm_row_per_thread == 0) {
513        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
514             (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
515        return HTP_STATUS_VTCM_TOO_SMALL;
516    }
517
518    if (src1) {
519        FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n",
520             op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
521             src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], L.src0_bytes_per_thread * n_threads,
522             L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
523    } else {
524        FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n", op_type,
525             src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
526             L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
527    }
528
529    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
530        return HTP_STATUS_OK;
531    }
532
533    // Prepare context
534    struct htp_act_context actx;
535    actx.octx = octx;
536
537    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
538
539    actx.src0_row_size = src0_row_size;
540    actx.src1_row_size = src1_row_size;
541    actx.dst_row_size  = dst_row_size;
542
543    actx.src0_row_size_aligned = src0_row_size_aligned;
544    actx.src1_row_size_aligned = src1_row_size_aligned;
545    actx.dst_row_size_aligned  = dst_row_size_aligned;
546
547    actx.src0_row_stride = src0_row_stride;
548    actx.src1_row_stride = src1_row_stride;
549
550    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
551    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
552    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
553    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
554
555    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
556    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
557    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
558
559    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
560    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
561    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
562
563    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
564    actx.src0_nrows = nrows;
565    actx.row_start  = row_start;
566
567    actx.nc = dst->ne[0];
568
569    // Pointers and GLU logic
570    const uint8_t * data_src0 = (const uint8_t *) src0->data;
571    const uint8_t * data_src1 = src1 ? (const uint8_t *) src1->data : NULL;
572
573    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU || octx->op == HTP_OP_GLU_SWIGLU_OAI || octx->op == HTP_OP_GLU_SWIGLU_CLAMP || octx->op == HTP_OP_GLU_GEGLU)) {
574         const int32_t swapped = octx->op_params[1];
575         data_src1 = data_src0;
576         actx.src1_row_size = actx.src0_row_size;
577
578         size_t nc_in_bytes = actx.nc * SIZEOF_FP32;
579         if (swapped) {
580             data_src0 += nc_in_bytes;
581         } else {
582             data_src1 += nc_in_bytes;
583         }
584    }
585
586    actx.data_src0 = data_src0;
587    actx.data_src1 = data_src1;
588    actx.data_dst  = (uint8_t *) dst->data;
589
590    work_queue_run(octx->ctx->work_queue, act_op_func, &actx, n_threads);
591    return HTP_STATUS_OK;
592}
593
594int op_activations(struct htp_ops_context * octx) {
595    switch (octx->src[0]->type) {
596        case HTP_TYPE_F32:
597            return execute_op_activations_f32(octx);
598
599        default:
600            return HTP_STATUS_NO_SUPPORT;
601    }
602}