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 →
43.9 KiB917 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#include <HAP_perf.h>
7
8#include <math.h>
9#include <string.h>
10
11#include "hex-dma.h"
12#include "hvx-utils.h"
13
14#define GGML_COMMON_DECL_C
15#include "ggml-common.h"
16#include "hex-common.h"
17#include "hex-profile.h"
18#include "htp-ctx.h"
19#include "htp-ops.h"
20#include "htp-tensor.h"
21
22#ifndef MIN
23#define MIN(a, b) ((a) < (b) ? (a) : (b))
24#endif
25
26// Context for binary operations
27struct htp_binary_context {
28    struct htp_ops_context * octx;
29
30    struct fastdiv_values src0_dim1_div; // ne01
31    struct fastdiv_values src0_dim2_div; // ne02
32    struct fastdiv_values src0_dim12_div;// ne03
33
34    struct fastdiv_values src1_dim1_div; // ne11
35    struct fastdiv_values src1_dim2_div; // ne12
36    struct fastdiv_values src1_dim3_div; // ne13
37
38    uint32_t block_max;
39    uint32_t nrows_per_thread;
40    uint32_t total_rows;
41    uint32_t row_start;
42    size_t   src0_row_size_aligned;
43    size_t   src1_row_size_aligned;
44    size_t   dst_row_size_aligned;
45
46    bool split_at_ne01;
47    bool split_at_ne02;
48};
49
50#define htp_binary_preamble                        \
51    const struct htp_tensor * src0 = octx->src[0]; \
52    const struct htp_tensor * src1 = octx->src[1]; \
53    const struct htp_tensor * dst  = octx->dst;    \
54                                                   \
55    const uint32_t ne00 = src0->ne[0];             \
56    const uint32_t ne01 = src0->ne[1];             \
57    const uint32_t ne02 = src0->ne[2];             \
58    const uint32_t ne03 = src0->ne[3];             \
59                                                   \
60    const uint32_t ne10 = src1->ne[0];             \
61    const uint32_t ne11 = src1->ne[1];             \
62    const uint32_t ne12 = src1->ne[2];             \
63    const uint32_t ne13 = src1->ne[3];             \
64                                                   \
65    const uint32_t nb01 = src0->nb[1];             \
66    const uint32_t nb02 = src0->nb[2];             \
67    const uint32_t nb03 = src0->nb[3];             \
68                                                   \
69    const uint32_t nb11 = src1->nb[1];             \
70    const uint32_t nb12 = src1->nb[2];             \
71    const uint32_t nb13 = src1->nb[3];             \
72                                                   \
73    const uint32_t nb1 = dst->nb[1];               \
74    const uint32_t nb2 = dst->nb[2];               \
75    const uint32_t nb3 = dst->nb[3];
76
77static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_t ir, uint32_t end_row, uint32_t ne01, uint32_t ne02) {
78    uint32_t i03, i02, i01, rem;
79    i03 = fastdiv(ir, &bctx->src0_dim12_div);
80    rem = ir - i03 * (ne02 * ne01);
81    i02 = fastdiv(rem, &bctx->src0_dim1_div);
82    i01 = rem - i02 * ne01;
83
84    uint32_t rows_left = end_row - ir;
85    uint32_t block_limit = rows_left;
86
87    if (bctx->split_at_ne01) {
88        block_limit = MIN(block_limit, ne01 - i01);
89    }
90    if (bctx->split_at_ne02) {
91         uint32_t rows_in_plane = (ne02 * ne01) - rem;
92         block_limit = MIN(block_limit, rows_in_plane);
93    }
94
95    return MIN(bctx->block_max, block_limit);
96}
97
98// Macro for scalar op switch
99#define COMPUTE_SCALAR_OP(DST, SRC, VAL, TYPE, N)                                               \
100    if(TYPE == HTP_TYPE_F32) {                                                                  \
101        switch (octx->op) {                                                                     \
102            case HTP_OP_ADD: hvx_add_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
103            case HTP_OP_SUB: hvx_sub_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
104            case HTP_OP_MUL: hvx_mul_scalar_f32_aa(DST, SRC, *(float *)VAL, N); break;          \
105            case HTP_OP_DIV: hvx_mul_scalar_f32_aa(DST, SRC, 1.0f / (*(float *)VAL), N); break; \
106            default: break;                                                                     \
107        }                                                                                       \
108    }                                                                                           \
109    else {                                                                                      \
110        switch (octx->op) {                                                                     \
111            case HTP_OP_ADD: hvx_add_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
112            case HTP_OP_SUB: hvx_sub_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
113            case HTP_OP_MUL: hvx_mul_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
114            case HTP_OP_DIV: hvx_div_scalar_f16_aa(DST, SRC, *(_Float16 *)VAL, N); break;       \
115            default: break;                                                                     \
116        }                                                                                       \
117    }
118
119// Macro for vector op switch (All Aligned)
120#define COMPUTE_VECTOR_OP_AAA(DST, SRC0, SRC1, TYPE, N)                  \
121    if(TYPE == HTP_TYPE_F32) {                                           \
122        switch (octx->op) {                                              \
123            case HTP_OP_ADD: hvx_add_f32_aaa(DST, SRC0, SRC1, N); break; \
124            case HTP_OP_SUB: hvx_sub_f32_aaa(DST, SRC0, SRC1, N); break; \
125            case HTP_OP_MUL: hvx_mul_f32_aaa(DST, SRC0, SRC1, N); break; \
126            case HTP_OP_DIV: hvx_div_f32_aaa(DST, SRC0, SRC1, N); break; \
127            default: break;                                              \
128        }                                                                \
129    }                                                                    \
130    else {                                                               \
131        switch (octx->op) {                                              \
132            case HTP_OP_ADD: hvx_add_f16_aaa(DST, SRC0, SRC1, N); break; \
133            case HTP_OP_SUB: hvx_sub_f16_aaa(DST, SRC0, SRC1, N); break; \
134            case HTP_OP_MUL: hvx_mul_f16_aaa(DST, SRC0, SRC1, N); break; \
135            case HTP_OP_DIV: hvx_div_f16_aaa(DST, SRC0, SRC1, N); break; \
136            default: break;                                              \
137        }                                                                \
138    }
139
140// Macro for vector op switch (Dst Aligned, Src0 Aligned, Src1 Unaligned)
141#define COMPUTE_VECTOR_OP_AAU(DST, SRC0, SRC1, TYPE, N)                  \
142    if(TYPE == HTP_TYPE_F32) {                                           \
143        switch (octx->op) {                                              \
144            case HTP_OP_ADD: hvx_add_f32_aau(DST, SRC0, SRC1, N); break; \
145            case HTP_OP_SUB: hvx_sub_f32_aau(DST, SRC0, SRC1, N); break; \
146            case HTP_OP_MUL: hvx_mul_f32_aau(DST, SRC0, SRC1, N); break; \
147            case HTP_OP_DIV: hvx_div_f32_aau(DST, SRC0, SRC1, N); break; \
148            default: break;                                              \
149        }                                                                \
150    }                                                                    \
151    else {                                                               \
152        switch (octx->op) {                                              \
153            case HTP_OP_ADD: hvx_add_f16_aau(DST, SRC0, SRC1, N); break; \
154            case HTP_OP_SUB: hvx_sub_f16_aau(DST, SRC0, SRC1, N); break; \
155            case HTP_OP_MUL: hvx_mul_f16_aau(DST, SRC0, SRC1, N); break; \
156            case HTP_OP_DIV: hvx_div_f16_aau(DST, SRC0, SRC1, N); break; \
157            default: break;                                              \
158        }                                                                \
159    }
160
161// Macro for vector op switch (All Unaligned - generic loop used in element repeat)
162#define COMPUTE_VECTOR_OP_UUU(DST, SRC0, SRC1, TYPE, N)                  \
163    if(TYPE == HTP_TYPE_F32) {                                           \
164        switch (octx->op) {                                              \
165            case HTP_OP_ADD: hvx_add_f32_uuu(DST, SRC0, SRC1, N); break; \
166            case HTP_OP_SUB: hvx_sub_f32_uuu(DST, SRC0, SRC1, N); break; \
167            case HTP_OP_MUL: hvx_mul_f32_uuu(DST, SRC0, SRC1, N); break; \
168            case HTP_OP_DIV: hvx_div_f32_uuu(DST, SRC0, SRC1, N); break; \
169            default: break;                                              \
170        }                                                                \
171    }                                                                    \
172    else {                                                               \
173        switch (octx->op) {                                              \
174            case HTP_OP_ADD: hvx_add_f16_uuu(DST, SRC0, SRC1, N); break; \
175            case HTP_OP_SUB: hvx_sub_f16_uuu(DST, SRC0, SRC1, N); break; \
176            case HTP_OP_MUL: hvx_mul_f16_uuu(DST, SRC0, SRC1, N); break; \
177            case HTP_OP_DIV: hvx_div_f16_uuu(DST, SRC0, SRC1, N); break; \
178            default: break;                                              \
179        }                                                                \
180    }
181
182// 1. Scalar src1 (ne10 == 1)
183static void binary_job_scalar(unsigned int nth, unsigned int ith, void * data) {
184    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
185    struct htp_ops_context * octx = bctx->octx;
186    htp_binary_preamble;
187
188    const uint32_t src0_type = octx->src[0]->type;
189    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
190    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
191    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
192    if (start_row >= end_row) return;
193
194    FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
195
196    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
197    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
198    size_t src0_spad_half    = octx->src0_spad.size_per_thread / 2;
199    size_t dst_spad_half     = octx->dst_spad.size_per_thread  / 2;
200
201    dma_queue * q = octx->ctx->dma[ith];
202    uint32_t ir_prefetch = start_row;
203    int spad_idx = 0;
204
205    // Preamble
206    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
207        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
208        uint32_t i03, i02, i01, rem;
209        i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
210        rem = ir_prefetch - i03 * (ne02 * ne01);
211        i02 = fastdiv(rem, &bctx->src0_dim1_div);
212        i01 = rem - i02 * ne01;
213
214        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
215        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
216
217        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
218        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
219
220        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
221        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
222        ir_prefetch += current_block_size;
223        spad_idx ^= 1;
224    }
225
226    // Main loop
227    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
228
229    for (uint32_t ir = start_row; ir < end_row; ) {
230        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
231
232        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
233        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
234
235        uint32_t i03, i02, i01, rem;
236        i03 = fastdiv(ir, &bctx->src0_dim12_div);
237        rem = ir - i03 * (ne02 * ne01);
238        i02 = fastdiv(rem, &bctx->src0_dim1_div);
239        i01 = rem - i02 * ne01;
240
241        // src1 indices (broadcast/repeat)
242        uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
243        uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
244        uint32_t i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div);
245
246        uint8_t * src1_ptr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
247        uint32_t s1_stride = (ne11 == 1) ? 0 : nb11;
248
249        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
250        for (uint32_t r = 0; r < current_block_size; r++) {
251            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
252            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
253            COMPUTE_SCALAR_OP(r_dst, r_src0, src1_ptr, src0_type, ne00);
254            src1_ptr += s1_stride;
255        }
256        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
257
258        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
259        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
260
261        if (ir_prefetch < end_row) {
262             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
263             uint32_t p03, p02, p01, prem;
264             p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
265             prem = ir_prefetch - p03 * (ne02 * ne01);
266             p02 = fastdiv(prem, &bctx->src0_dim1_div);
267             p01 = prem - p02 * ne01;
268             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
269
270             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
271             ir_prefetch += next_block_size;
272        }
273        ir += current_block_size;
274    }
275
276    dma_queue_flush(q);
277}
278
279// 2. Vector Same Shape (ne1x == ne0x) or Simple Broadcast
280static void binary_job_vector_same_shape(unsigned int nth, unsigned int ith, void * data) {
281    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
282    struct htp_ops_context * octx = bctx->octx;
283    htp_binary_preamble;
284
285    const uint32_t src0_type = octx->src[0]->type;
286    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
287    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
288    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
289    if (start_row >= end_row) return;
290
291    FARF(HIGH, "binary-same-shape: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
292
293    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
294    uint8_t * src1_spad_base = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread);
295    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
296
297    size_t src0_spad_half = octx->src0_spad.size_per_thread / 2;
298    size_t src1_spad_half = octx->src1_spad.size_per_thread / 2;
299    size_t dst_spad_half  = octx->dst_spad.size_per_thread  / 2;
300
301    dma_queue * q = octx->ctx->dma[ith];
302    uint32_t ir_prefetch = start_row;
303    int spad_idx = 0;
304
305    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
306        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
307        uint32_t i03, i02, i01, rem;
308        i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
309        rem = ir_prefetch - i03 * (ne02 * ne01);
310        i02 = fastdiv(rem, &bctx->src0_dim1_div);
311        i01 = rem - i02 * ne01;
312
313        uint32_t i13 = (ne13 == 1) ? 0 : i03;
314        uint32_t i12 = (ne12 == 1) ? 0 : i02;
315        uint32_t i11 = (ne11 == 1) ? 0 : i01;
316
317        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
318        uint8_t * src1_curr = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
319        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
320
321        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
322        uint8_t * s1_spad = src1_spad_base + spad_idx * src1_spad_half;
323        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
324
325        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
326        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
327        dma_queue_push(q, dma_make_ptr(s1_spad, src1_curr), bctx->src1_row_size_aligned, nb11, row_size_bytes, current_block_size);
328        ir_prefetch += current_block_size;
329        spad_idx ^= 1;
330    }
331
332    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
333
334    for (uint32_t ir = start_row; ir < end_row; ) {
335        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
336        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
337        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
338        uint8_t * s1_spad = (uint8_t *) dma_queue_pop(q).dst;
339
340        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
341        for (uint32_t r = 0; r < current_block_size; r++) {
342            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
343            uint8_t * r_src1 = s1_spad + r * bctx->src1_row_size_aligned;
344            uint8_t * r_dst  = d_spad  + r * bctx->dst_row_size_aligned;
345            COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
346        }
347        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
348
349        uint32_t i03, i02, i01, rem;
350        i03 = fastdiv(ir, &bctx->src0_dim12_div);
351        rem = ir - i03 * (ne02 * ne01);
352        i02 = fastdiv(rem, &bctx->src0_dim1_div);
353        i01 = rem - i02 * ne01;
354        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
355        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
356
357        if (ir_prefetch < end_row) {
358             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
359             uint32_t p03, p02, p01, prem;
360             p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
361             prem = ir_prefetch - p03 * (ne02 * ne01);
362             p02 = fastdiv(prem, &bctx->src0_dim1_div);
363             p01 = prem - p02 * ne01;
364
365             uint32_t p13 = (ne13 == 1) ? 0 : p03;
366             uint32_t p12 = (ne12 == 1) ? 0 : p02;
367             uint32_t p11 = (ne11 == 1) ? 0 : p01;
368
369             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
370             uint8_t * s1_next = (uint8_t *)src1->data + p13 * nb13 + p12 * nb12 + p11 * nb11;
371
372             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
373             dma_queue_push(q, dma_make_ptr(s1_spad, s1_next), bctx->src1_row_size_aligned, nb11, row_size_bytes, next_block_size);
374
375             ir_prefetch += next_block_size;
376        }
377        ir += current_block_size;
378    }
379
380    dma_queue_flush(q);
381}
382
383// 3. Row Broadcast (ne11 == 1, ne12 == 1, single row src1)
384static void binary_job_vector_row_broadcast(unsigned int nth, unsigned int ith, void * data) {
385    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
386    struct htp_ops_context * octx = bctx->octx;
387    htp_binary_preamble;
388
389    const uint32_t src0_type  = octx->src[0]->type;
390    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
391    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
392    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
393    if (start_row >= end_row) return;
394
395    FARF(HIGH, "binary-row-bcast: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
396
397    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
398    uint8_t * src1_spad_base = octx->src1_spad.data + (ith * octx->src1_spad.size_per_thread);
399    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
400
401    size_t src0_spad_half = octx->src0_spad.size_per_thread / 2;
402    size_t dst_spad_half  = octx->dst_spad.size_per_thread  / 2;
403
404    dma_queue * q = octx->ctx->dma[ith];
405    uint32_t ir_prefetch = start_row;
406    int spad_idx = 0;
407
408    void * s1_ptr = (void *) src1_spad_base;
409
410    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
411        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
412        uint32_t i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
413        uint32_t rem = ir_prefetch - i03 * (ne02 * ne01);
414        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
415        uint32_t i01 = rem - i02 * ne01;
416
417        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
418        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
419
420        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
421        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
422
423        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
424        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
425        ir_prefetch += current_block_size;
426        spad_idx ^= 1;
427    }
428
429    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
430
431    for (uint32_t ir = start_row; ir < end_row; ) {
432        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
433        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(q).src;
434        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
435
436        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
437        for (uint32_t r = 0; r < current_block_size; r++) {
438            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
439            uint8_t * r_src1 = (uint8_t *)s1_ptr; // Constant
440            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
441            COMPUTE_VECTOR_OP_AAA(r_dst, r_src0, r_src1, src0_type, ne00);
442        }
443        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
444
445        uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
446        uint32_t rem = ir - i03 * (ne02 * ne01);
447        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
448        uint32_t i01 = rem - i02 * ne01;
449        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
450        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
451
452        if (ir_prefetch < end_row) {
453             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
454             uint32_t p03  = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
455             uint32_t prem = ir_prefetch - p03 * (ne02 * ne01);
456             uint32_t p02  = fastdiv(prem, &bctx->src0_dim1_div);
457             uint32_t p01  = prem - p02 * ne01;
458             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
459             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
460             ir_prefetch += next_block_size;
461        }
462        ir += current_block_size;
463    }
464
465    dma_queue_flush(q);
466}
467
468// 4. Vector Complex (ne10 == ne00, complex broadcast)
469static void binary_job_vector_complex(unsigned int nth, unsigned int ith, void * data) {
470    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
471    struct htp_ops_context * octx = bctx->octx;
472    htp_binary_preamble;
473
474    const uint32_t src0_type = octx->src[0]->type;
475    const uint32_t row_size_bytes = (src0_type == HTP_TYPE_F32) ? ne00 * sizeof(float) : ne00 * sizeof(_Float16);
476    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
477    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
478    if (start_row >= end_row) return;
479
480    FARF(HIGH, "binary-complex: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
481
482    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
483    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
484    size_t src0_spad_half    = octx->src0_spad.size_per_thread / 2;
485    size_t dst_spad_half     = octx->dst_spad.size_per_thread  / 2;
486
487    dma_queue * q = octx->ctx->dma[ith];
488    uint32_t ir_prefetch = start_row;
489    int spad_idx = 0;
490
491    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
492        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
493        uint32_t i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
494        uint32_t rem = ir_prefetch - i03 * (ne02 * ne01);
495        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
496        uint32_t i01 = rem - i02 * ne01;
497
498        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
499        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
500
501        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
502        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
503
504        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
505        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
506        ir_prefetch += current_block_size;
507        spad_idx ^= 1;
508    }
509
510    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
511
512    for (uint32_t ir = start_row; ir < end_row; ) {
513        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
514        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
515        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
516
517        uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
518        uint32_t rem = ir - i03 * (ne02 * ne01);
519        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
520        uint32_t i01 = rem - i02 * ne01;
521
522        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
523        for (uint32_t r = 0; r < current_block_size; r++) {
524            uint32_t r_i01 = i01 + r;
525            uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
526            uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
527            uint32_t i11 = fastmodulo(r_i01, ne11, &bctx->src1_dim1_div);
528
529            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
530            uint8_t * r_src1 = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
531            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
532
533            // Read src1 from DDR (unaligned)
534            COMPUTE_VECTOR_OP_AAU(r_dst, r_src0, r_src1, src0_type, ne00);
535        }
536        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
537
538        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
539        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
540
541        if (ir_prefetch < end_row) {
542             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
543             uint32_t p03  = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
544             uint32_t prem = ir_prefetch - p03 * (ne02 * ne01);
545             uint32_t p02  = fastdiv(prem, &bctx->src0_dim1_div);
546             uint32_t p01  = prem - p02 * ne01;
547             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
548             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
549             ir_prefetch += next_block_size;
550        }
551        ir += current_block_size;
552    }
553
554    dma_queue_flush(q);
555}
556
557// 5. Element Repeat (ne10 != ne00)
558static void binary_job_element_repeat(unsigned int nth, unsigned int ith, void * data) {
559    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
560    struct htp_ops_context * octx = bctx->octx;
561    htp_binary_preamble;
562
563    const uint32_t src0_type = octx->src[0]->type;
564    const uint32_t elem_size_bytes = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
565    const uint32_t row_size_bytes = ne00 * elem_size_bytes;;
566    const uint32_t start_row  = bctx->row_start + bctx->nrows_per_thread * ith;
567    const uint32_t end_row    = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
568    if (start_row >= end_row) return;
569
570    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
571    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
572    size_t src0_spad_half    = octx->src0_spad.size_per_thread / 2;
573    size_t dst_spad_half     = octx->dst_spad.size_per_thread  / 2;
574
575    FARF(HIGH, "binary-repeat: %d/%d (%u:%u) row-size %u (%u)", ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
576
577    dma_queue * q = octx->ctx->dma[ith];
578    uint32_t ir_prefetch = start_row;
579    int spad_idx = 0;
580
581    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
582        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
583        uint32_t i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
584        uint32_t rem = ir_prefetch - i03 * (ne02 * ne01);
585        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
586        uint32_t i01 = rem - i02 * ne01;
587
588        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
589        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
590
591        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
592        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
593
594        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
595        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
596        ir_prefetch += current_block_size;
597        spad_idx ^= 1;
598    }
599
600    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
601
602    for (uint32_t ir = start_row; ir < end_row; ) {
603        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
604        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
605        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
606
607        uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
608        uint32_t rem = ir - i03 * (ne02 * ne01);
609        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
610        uint32_t i01 = rem - i02 * ne01;
611
612        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
613        for (uint32_t r = 0; r < current_block_size; r++) {
614            uint32_t r_i01 = i01 + r;
615            uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
616            uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
617            uint32_t i11 = fastmodulo(r_i01, ne11, &bctx->src1_dim1_div);
618
619            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
620            uint8_t * r_src1_row = (uint8_t *)src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
621            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
622
623            // Repeat src1 row
624            for (uint32_t c = 0; c < ne00; c += ne10) {
625                uint32_t len = MIN(ne10, ne00 - c);
626                // Use UUU for speed and simplicity
627                COMPUTE_VECTOR_OP_UUU(r_dst + c * elem_size_bytes, r_src0 + c * elem_size_bytes, r_src1_row, src0_type, len);
628            }
629        }
630        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
631
632        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
633        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
634
635        if (ir_prefetch < end_row) {
636             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
637             uint32_t p03  = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
638             uint32_t prem = ir_prefetch - p03 * (ne02 * ne01);
639             uint32_t p02  = fastdiv(prem, &bctx->src0_dim1_div);
640             uint32_t p01  = prem - p02 * ne01;
641             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
642             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
643             ir_prefetch += next_block_size;
644        }
645        ir += current_block_size;
646    }
647
648    dma_queue_flush(q);
649}
650
651// 6. ADD_ID (src1 gathered via src2 indices)
652static void binary_job_add_id(unsigned int nth, unsigned int ith, void * data) {
653    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
654    struct htp_ops_context * octx = bctx->octx;
655
656    const struct htp_tensor * src0 = octx->src[0];
657    const struct htp_tensor * src1 = octx->src[1];
658    const struct htp_tensor * src2 = octx->src[2];
659    const struct htp_tensor * dst  = octx->dst;
660
661    const uint32_t ne00 = src0->ne[0];
662    const uint32_t ne01 = src0->ne[1];
663    const uint32_t ne02 = src0->ne[2];
664    const uint32_t ne03 = src0->ne[3];
665    const uint32_t ne11 = src1->ne[1]; // for bounds check
666
667    const uint32_t nb01 = src0->nb[1];
668    const uint32_t nb02 = src0->nb[2];
669    const uint32_t nb03 = src0->nb[3];
670    const uint32_t nb11 = src1->nb[1]; // src1 row stride
671
672    const uint32_t nb1 = dst->nb[1];
673    const uint32_t nb2 = dst->nb[2];
674    const uint32_t nb3 = dst->nb[3];
675
676    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
677    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
678    if (start_row >= end_row) return;
679
680    uint8_t * src0_spad_base = octx->src0_spad.data + (ith * octx->src0_spad.size_per_thread);
681    uint8_t * dst_spad_base  = octx->dst_spad.data  + (ith * octx->dst_spad.size_per_thread);
682    size_t src0_spad_half    = octx->src0_spad.size_per_thread / 2;
683    size_t dst_spad_half     = octx->dst_spad.size_per_thread  / 2;
684
685    dma_queue * q = octx->ctx->dma[ith];
686    uint32_t ir_prefetch = start_row;
687    int spad_idx = 0;
688
689    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
690        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
691        uint32_t i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
692        uint32_t rem = ir_prefetch - i03 * (ne02 * ne01);
693        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
694        uint32_t i01 = rem - i02 * ne01;
695
696        uint8_t * src0_curr = (uint8_t *)src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
697        uint8_t * dst_curr  = (uint8_t *)dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
698
699        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
700        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
701
702        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), 0);
703        dma_queue_push(q, dma_make_ptr(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, ne00 * sizeof(float), current_block_size);
704        ir_prefetch += current_block_size;
705        spad_idx ^= 1;
706    }
707
708    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
709
710    for (uint32_t ir = start_row; ir < end_row; ) {
711        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
712        uint8_t * d_spad = (uint8_t *) dma_queue_pop(q).src;
713        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(q).dst;
714
715        uint32_t i03 = fastdiv(ir, &bctx->src0_dim12_div);
716        uint32_t rem = ir - i03 * (ne02 * ne01);
717        uint32_t i02 = fastdiv(rem, &bctx->src0_dim1_div);
718        uint32_t i01 = rem - i02 * ne01;
719
720        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
721        for (uint32_t r = 0; r < current_block_size; r++) {
722            uint32_t r_i01 = i01 + r; // linear within block since we split at ne01
723
724            const int32_t idx = *(int32_t *)((char *)src2->data + r_i01 * src2->nb[0] + i02 * src2->nb[1]);
725
726            uint8_t * r_src1 = (uint8_t *)src1->data + idx * nb11;
727            uint8_t * r_src0 = s0_spad + r * bctx->src0_row_size_aligned;
728            uint8_t * r_dst  = d_spad + r * bctx->dst_row_size_aligned;
729
730            hvx_add_f32_aau(r_dst, r_src0, r_src1, ne00);
731        }
732        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
733
734        uint8_t * dst_curr = (uint8_t *)dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
735        dma_queue_push(q, dma_make_ptr(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, ne00 * sizeof(float), current_block_size);
736
737        if (ir_prefetch < end_row) {
738             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
739             uint32_t p03  = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
740             uint32_t prem = ir_prefetch - p03 * (ne02 * ne01);
741             uint32_t p02  = fastdiv(prem, &bctx->src0_dim1_div);
742             uint32_t p01  = prem - p02 * ne01;
743             uint8_t * s0_next = (uint8_t *)src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
744             dma_queue_push(q, dma_make_ptr(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, ne00 * sizeof(float), next_block_size);
745             ir_prefetch += next_block_size;
746        }
747        ir += current_block_size;
748    }
749
750    dma_queue_flush(q);
751}
752
753static int execute_op_binary(struct htp_ops_context * octx) {
754    const struct htp_tensor * src0 = octx->src[0];
755    const struct htp_tensor * src1 = octx->src[1];
756    const struct htp_tensor * dst  = octx->dst;
757
758    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
759
760    // Use packed row sizes for VTCM allocation and alignment
761    const uint32_t src0_type = octx->src[0]->type;
762    const size_t elem_size = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
763    const size_t src0_row_size = src0->ne[0] * elem_size;
764    const size_t src1_row_size = src1->ne[0] * elem_size;
765    const size_t dst_row_size  = dst->ne[0]  * elem_size;
766
767    uint32_t row_start = 0;
768    uint32_t nrows     = src0_nrows;
769
770    if (octx->ctx->mdev.count > 1) {
771        uint32_t rows_per_chunk = 0;
772        htp_tensor_mdev_rows_per_chunk(dst, (uint32_t) elem_size, (uint32_t) dst_row_size, &rows_per_chunk);
773        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);
774        row_start = range.start;
775        nrows     = range.count;
776    }
777
778    if (nrows == 0) {
779        return HTP_STATUS_OK;
780    }
781
782    const uint32_t n_threads = octx->n_threads;
783
784    size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
785    size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
786    size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  VLEN);
787
788    bool is_add_id = (octx->op == HTP_OP_ADD_ID);
789    bool is_scalar = !is_add_id && (src1->ne[0] == 1);
790
791    bool is_transposed = (src0->nb[1] < src0_row_size || src1->nb[1] < src1_row_size || dst->nb[1] < dst_row_size);
792
793    bool is_same_shape = !is_add_id && !is_scalar && !is_transposed &&
794               (src1->ne[0] == src0->ne[0] && src0->ne[0] % VLEN == 0) &&
795               (src1->ne[1] == src0->ne[1] || src1->ne[1] == 1) &&
796               (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) &&
797               (src1->ne[3] == src0->ne[3] || src1->ne[3] == 1);
798
799    bool is_row_bcast = is_same_shape && (src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
800    bool is_complex   = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] == src0->ne[0]);
801    bool is_repeat    = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] != src0->ne[0]);
802
803    size_t spad_row_total;
804    if (is_same_shape) {
805        spad_row_total = 2 * (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned);
806    } else {
807        spad_row_total = 2 * (src0_row_size_aligned + dst_row_size_aligned);
808    }
809
810    size_t rows_per_buffer = octx->ctx->vtcm_size / (n_threads * spad_row_total);
811
812    // Adjust for static src1 in row_bcast case
813    if (is_row_bcast) {
814        size_t needed_static = src1_row_size_aligned;
815        if (octx->ctx->vtcm_size < needed_static) return HTP_STATUS_VTCM_TOO_SMALL;
816        size_t avail = octx->ctx->vtcm_size - needed_static;
817        rows_per_buffer = avail / (n_threads * spad_row_total);
818    }
819
820    if (rows_per_buffer < 1) {
821        FARF(ERROR, "binary: VTCM too small\n");
822        return HTP_STATUS_VTCM_TOO_SMALL;
823    }
824
825    octx->src0_spad.size_per_thread = rows_per_buffer * 2 * src0_row_size_aligned;
826    octx->dst_spad.size_per_thread  = rows_per_buffer * 2 * dst_row_size_aligned;
827
828    if (is_add_id || is_scalar || is_complex || is_repeat || is_row_bcast) {
829        octx->src1_spad.size_per_thread = 0;
830    } else {
831        octx->src1_spad.size_per_thread = rows_per_buffer * 2 * src1_row_size_aligned;
832    }
833
834    octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
835    octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
836    if (is_row_bcast) {
837        octx->src1_spad.size = src1_row_size_aligned;
838    } else {
839        octx->src1_spad.size = n_threads * octx->src1_spad.size_per_thread;
840    }
841
842    if (octx->ctx->vtcm_size < (octx->src0_spad.size + octx->src1_spad.size + octx->dst_spad.size)) {
843        return HTP_STATUS_VTCM_TOO_SMALL;
844    }
845
846    octx->src0_spad.data = octx->ctx->vtcm_base;                        octx->src0_spad.src = NULL;
847    octx->src1_spad.data = octx->src0_spad.data + octx->src0_spad.size; octx->src1_spad.src = NULL;
848    octx->dst_spad.data  = octx->src1_spad.data + octx->src1_spad.size; octx->dst_spad.src  = NULL;
849
850    if ((octx->flags & HTP_OPFLAGS_SKIP_COMPUTE)) {
851        return HTP_STATUS_OK;
852    }
853
854    dma_queue * q = octx->ctx->dma[0];
855    if (is_row_bcast) {
856        dma_queue_push(q, dma_make_ptr(octx->src1_spad.data, (const void *) src1->data), src1_row_size_aligned, 0, src1->ne[0] * elem_size, 1);
857    }
858
859    struct htp_binary_context bctx;
860    bctx.octx                  = octx;
861    bctx.nrows_per_thread      = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
862    bctx.total_rows            = nrows;
863    bctx.row_start             = row_start;
864    bctx.block_max             = rows_per_buffer;
865    bctx.src0_row_size_aligned = src0_row_size_aligned;
866    bctx.src1_row_size_aligned = src1_row_size_aligned;
867    bctx.dst_row_size_aligned  = dst_row_size_aligned;
868
869    bctx.src0_dim1_div  = init_fastdiv_values(src0->ne[1]);
870    bctx.src0_dim2_div  = init_fastdiv_values(src0->ne[2]);
871    bctx.src0_dim12_div = init_fastdiv_values(src0->ne[1] * src0->ne[2]);
872
873    bctx.src1_dim1_div  = init_fastdiv_values(src1->ne[1]);
874    bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
875    bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
876
877    bool src0_contig_dim1 = (src0->nb[2] == src0->ne[1] * src0->nb[1]);
878    bool dst_contig_dim1  = (dst->nb[2]  == src0->ne[1] * dst->nb[1]);
879
880    bool src0_contig_dim2 = (src0->nb[3] == src0->ne[2] * src0->nb[2]);
881    bool dst_contig_dim2  = (dst->nb[3]  == src0->ne[2] * dst->nb[2]);
882
883    bctx.split_at_ne01 = (src0->ne[2] > 1) && ((src1->ne[1] > 1) || (src1->ne[2] > 1) || !src0_contig_dim1 || !dst_contig_dim1);
884    bctx.split_at_ne02 = (src0->ne[3] > 1) && ((src1->ne[2] > 1) || (src1->ne[3] > 1) || !src0_contig_dim2 || !dst_contig_dim2);
885
886    worker_callback_t worker_func;
887    if (is_add_id)          worker_func = binary_job_add_id;
888    else if (is_scalar)     worker_func = binary_job_scalar;
889    else if (is_row_bcast)  worker_func = binary_job_vector_row_broadcast;
890    else if (is_same_shape) worker_func = binary_job_vector_same_shape;
891    else if (is_complex)    worker_func = binary_job_vector_complex;
892    else                    worker_func = binary_job_element_repeat;
893
894    if (is_row_bcast) {
895        dma_queue_pop(q);
896    }
897
898    work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
899
900    return HTP_STATUS_OK;
901}
902
903int op_binary(struct htp_ops_context * octx) {
904
905    // Does not support permutations of src1
906    const struct htp_tensor * src1 = octx->src[1];
907    if (src1->nb[1] < src1->nb[0]) {
908        return HTP_STATUS_NO_SUPPORT;
909    }
910
911    const uint32_t src0_type = octx->src[0]->type;
912    if ((src0_type == HTP_TYPE_F32) || (src0_type == HTP_TYPE_F16)) {
913        return execute_op_binary(octx);
914    }
915
916    return HTP_STATUS_NO_SUPPORT;
917}