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 →
24.2 KiB558 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 <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 "hex-common.h"
16#include "hex-profile.h"
17#include "htp-ctx.h"
18#include "htp-ops.h"
19#include "htp-tensor.h"
20
21/* Circular wrap: maps any integer x into [0, n) */
22static inline uint32_t wrap_around(int32_t x, uint32_t n) {
23    return (uint32_t)(((x % (int32_t)n) + (int32_t)n) % (int32_t)n);
24}
25
26/* Decompose a flat dst row index into (i1, i2, i3) */
27static inline void pad_decompose_row(uint32_t ir, uint32_t ne1, uint32_t ne2,
28                                     uint32_t *i1, uint32_t *i2, uint32_t *i3) {
29    *i1 = ir % ne1;
30    *i2 = (ir / ne1) % ne2;
31    *i3 = ir / (ne1 * ne2);
32}
33
34/* Return non-zero if row (i1,i2,i3) falls in the non-padded interior */
35static inline int pad_is_interior(uint32_t i1, uint32_t i2, uint32_t i3,
36                                   int32_t lp1, int32_t rp1, uint32_t ne1,
37                                   int32_t lp2, int32_t rp2, uint32_t ne2,
38                                   int32_t lp3, int32_t rp3, uint32_t ne3) {
39    return ((int32_t)i1 >= lp1 && (int32_t)i1 < (int32_t)ne1 - rp1) &&
40           ((int32_t)i2 >= lp2 && (int32_t)i2 < (int32_t)ne2 - rp2) &&
41           ((int32_t)i3 >= lp3 && (int32_t)i3 < (int32_t)ne3 - rp3);
42}
43
44/* Compute the DDR src row pointer for a zero-pad interior row */
45static inline const uint8_t * pad_src_row_ptr(const struct htp_tensor * src,
46                                               uint32_t i1, uint32_t i2, uint32_t i3,
47                                               int32_t lp1, int32_t lp2, int32_t lp3) {
48    return (const uint8_t *) src->data
49        + (i1 - (uint32_t)lp1) * src->nb[1]
50        + (i2 - (uint32_t)lp2) * src->nb[2]
51        + (i3 - (uint32_t)lp3) * src->nb[3];
52}
53
54/* Compute the DDR src row pointer for a circular row (wrap-around indexing) */
55static inline const uint8_t * pad_circ_src_row_ptr(const struct htp_tensor * src,
56                                                    uint32_t i1, uint32_t i2, uint32_t i3,
57                                                    int32_t lp1, int32_t lp2, int32_t lp3) {
58    return (const uint8_t *) src->data
59        + wrap_around((int32_t)i1 - lp1, src->ne[1]) * src->nb[1]
60        + wrap_around((int32_t)i2 - lp2, src->ne[2]) * src->nb[2]
61        + wrap_around((int32_t)i3 - lp3, src->ne[3]) * src->nb[3];
62}
63
64struct htp_pad_context {
65    struct htp_ops_context * octx;
66
67    int32_t  lp0, rp0;
68    int32_t  lp1, rp1;
69    int32_t  lp2, rp2;
70    int32_t  lp3, rp3;
71
72    uint32_t nrows_per_thread;
73    uint32_t total_dst_rows;
74    uint32_t row_start;
75
76    size_t   type_size;
77
78    // Row sizes for DMA kernel (populated when VTCM is available)
79    size_t   src_row_size;
80    size_t   src_row_size_aligned;
81    size_t   dst_row_size;
82    size_t   dst_row_size_aligned;
83};
84
85#define htp_pad_preamble                                                             \
86    const struct htp_tensor * src = octx->src[0];                                    \
87    const struct htp_tensor * dst = octx->dst;                                       \
88                                                                                     \
89    const uint32_t ne00 = src->ne[0];                                                \
90    const uint32_t nb00 = src->nb[0];                                                \
91                                                                                     \
92    const uint32_t ne0 = dst->ne[0];                                                 \
93    const uint32_t ne1 = dst->ne[1];                                                 \
94    const uint32_t ne2 = dst->ne[2];                                                 \
95    const uint32_t ne3 = dst->ne[3];                                                 \
96                                                                                     \
97    const uint32_t nb1 = dst->nb[1];                                                 \
98    const uint32_t nb2 = dst->nb[2];                                                 \
99    const uint32_t nb3 = dst->nb[3];                                                 \
100                                                                                     \
101    const int32_t lp0 = pctx->lp0, rp0 = pctx->rp0;                                  \
102    const int32_t lp1 = pctx->lp1, rp1 = pctx->rp1;                                  \
103    const int32_t lp2 = pctx->lp2, rp2 = pctx->rp2;                                  \
104    const int32_t lp3 = pctx->lp3, rp3 = pctx->rp3;                                  \
105                                                                                     \
106    const size_t type_size = pctx->type_size;                                        \
107                                                                                     \
108    const uint32_t row_start = pctx->row_start + pctx->nrows_per_thread * ith;       \
109    const uint32_t row_end   = MIN(row_start + pctx->nrows_per_thread, pctx->row_start + pctx->total_dst_rows);
110
111
112#define htp_pad_dma_preamble                                                                \
113    const size_t src_row_size         = pctx->src_row_size;                                 \
114    const size_t src_row_size_aligned = pctx->src_row_size_aligned;                         \
115    const size_t dst_row_size         = pctx->dst_row_size;                                 \
116    const size_t dst_row_size_aligned = pctx->dst_row_size_aligned;                         \
117                                                                                            \
118    uint8_t * src_spad_base = octx->src0_spad.data + ith * octx->src0_spad.size_per_thread; \
119    uint8_t * dst_spad_base = octx->dst_spad.data  + ith * octx->dst_spad.size_per_thread;  \
120                                                                                            \
121    dma_queue * dma = octx->ctx->dma[ith];
122
123// ---------------------------------------------------------------------------
124// HVX vectorized PAD kernel
125// ---------------------------------------------------------------------------
126
127static void pad_job_per_thread_hvx(unsigned int nth, unsigned int ith, void * data) {
128    const struct htp_pad_context * pctx = (const struct htp_pad_context *) data;
129    struct htp_ops_context * octx = pctx->octx;
130    htp_pad_preamble;
131
132    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
133    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
134
135    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
136        uint32_t i1, i2, i3;
137        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
138
139        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
140
141        const int interior = pad_is_interior(i1, i2, i3,
142                                             lp1, rp1, ne1,
143                                             lp2, rp2, ne2,
144                                             lp3, rp3, ne3);
145
146        if (!interior) {
147            hvx_splat_f32_u(dst_ptr, 0.0f, ne0);
148        } else {
149            const uint8_t * src_ptr = pad_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3);
150
151            if (lp0 > 0) {
152                hvx_splat_f32_u(dst_ptr, 0.0f, (uint32_t)lp0);
153            }
154
155            uint8_t * dst_row_start = dst_ptr + (size_t)lp0 * type_size;
156            if (nb00 == type_size) {
157                hvx_copy_f32_uu(dst_row_start, src_ptr, ne00);
158            } else {
159                for (uint32_t i = 0; i < ne00; i++) {
160                    memcpy(dst_row_start + i * type_size,
161                           src_ptr + (size_t)i * nb00,
162                           type_size);
163                }
164            }
165
166            if (rp0 > 0) {
167                hvx_splat_f32_u(dst_ptr + ((size_t)lp0 + ne00) * type_size, 0.0f, (uint32_t)rp0);
168            }
169        }
170    }
171
172    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
173
174    FARF(HIGH, "pad-hvx %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
175         ith, nth,
176         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
177         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
178         row_start, row_end);
179}
180
181// ---------------------------------------------------------------------------
182// HVX + DMA PAD kernel - aligned, double-buffered
183// ---------------------------------------------------------------------------
184
185static void pad_job_per_thread_hvx_dma(unsigned int nth, unsigned int ith, void * data) {
186    const struct htp_pad_context * pctx = (const struct htp_pad_context *) data;
187    struct htp_ops_context * octx = pctx->octx;
188    htp_pad_preamble;
189    htp_pad_dma_preamble;
190
191    // -----------------------------------------------------------------------
192    // Priming phase: push 2 pairs of (dummy_dst_DMA, src_DMA) to seed the
193    // double-buffer pipeline before the main loop begins.
194    // -----------------------------------------------------------------------
195    for (uint32_t ir = row_start, spad_idx = 0; ir < row_end && spad_idx < 2; ir++, spad_idx++) {
196        uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned;
197        uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned;
198
199        dma_queue_push_vtcm_to_ddr(dma,
200            dma_make_ptr((uint8_t *)dst->data, dst_spad_cur),
201            dst_row_size, dst_row_size_aligned, 0);
202
203        uint32_t i1, i2, i3;
204        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
205        const int interior = pad_is_interior(i1, i2, i3,
206                                             lp1, rp1, ne1,
207                                             lp2, rp2, ne2,
208                                             lp3, rp3, ne3);
209
210        const uint8_t * src_ptr = interior
211            ? pad_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3) : NULL;
212
213        // Interior row: real DMA (1 row) from DDR to VTCM.
214        // Border row: null DMA (nrows=0)
215        dma_queue_push_ddr_to_vtcm(dma,
216            dma_make_ptr(src_spad_cur,
217                         src_ptr ? src_ptr : (const uint8_t *)src_spad_cur),
218            src_row_size_aligned, src_row_size, src_ptr ? 1 : 0);
219    }
220
221    // -----------------------------------------------------------------------
222    // Main loop: pop completed DMAs, compute in VTCM with aligned HVX ops,
223    // push dst DMA and prefetch src for the next+1 row.
224    // -----------------------------------------------------------------------
225    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
226
227    for (uint32_t ir = row_start; ir < row_end; ir++) {
228        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
229        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
230
231        uint32_t i1, i2, i3;
232        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
233
234        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
235
236        const int interior = pad_is_interior(i1, i2, i3,
237                                             lp1, rp1, ne1,
238                                             lp2, rp2, ne2,
239                                             lp3, rp3, ne3);
240
241        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
242        if (!interior) {
243            hvx_splat_f32_a(dst_spad_cur, 0.0f, ne0);
244        } else {
245            hvx_splat_f32_a(dst_spad_cur, 0.0f, ne0);
246
247            uint8_t * dst_interior = dst_spad_cur + (size_t)lp0 * type_size;
248
249            if ((uintptr_t)dst_interior % VLEN == 0) {
250                hvx_copy_f32_aa(dst_interior, src_spad_cur, ne00);
251            } else {
252                hvx_copy_f32_ua(dst_interior, src_spad_cur, ne00);
253            }
254        }
255        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
256
257        dma_queue_push_vtcm_to_ddr(dma,
258            dma_make_ptr(dst_ptr, dst_spad_cur),
259            dst_row_size, dst_row_size_aligned, 1);
260
261        const uint32_t next_row = ir + 2;
262        if (next_row < row_end) {
263            uint32_t ni1, ni2, ni3;
264            pad_decompose_row(next_row, ne1, ne2, &ni1, &ni2, &ni3);
265            const int next_interior = pad_is_interior(ni1, ni2, ni3,
266                                                      lp1, rp1, ne1,
267                                                      lp2, rp2, ne2,
268                                                      lp3, rp3, ne3);
269            const uint8_t * next_src_ptr = next_interior
270                ? pad_src_row_ptr(src, ni1, ni2, ni3, lp1, lp2, lp3) : NULL;
271
272            dma_queue_push_ddr_to_vtcm(dma,
273                dma_make_ptr(src_spad_cur,
274                             next_src_ptr ? next_src_ptr : (const uint8_t *)src_spad_cur),
275                src_row_size_aligned, src_row_size, next_src_ptr ? 1 : 0);
276        }
277    }
278
279    dma_queue_flush(dma);
280
281    FARF(HIGH, "pad-hvx-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
282         ith, nth,
283         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
284         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
285         row_start, row_end);
286}
287
288// ---------------------------------------------------------------------------
289// HVX circular PAD kernel
290// ---------------------------------------------------------------------------
291
292static void pad_job_per_thread_hvx_circular(unsigned int nth, unsigned int ith, void * data) {
293    const struct htp_pad_context * pctx = (const struct htp_pad_context *) data;
294    struct htp_ops_context * octx = pctx->octx;
295    htp_pad_preamble;
296
297    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
298    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
299
300    for (uint32_t dst_row = row_start; dst_row < row_end; dst_row++) {
301        uint32_t i1, i2, i3;
302        pad_decompose_row(dst_row, ne1, ne2, &i1, &i2, &i3);
303
304        uint8_t       * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
305        const uint8_t * src_row = pad_circ_src_row_ptr(src, i1, i2, i3, lp1, lp2, lp3);
306
307        if (nb00 == type_size) {
308
309            if (lp0 > 0) {
310                if ((uint32_t)lp0 < 32) {
311                    memcpy(dst_ptr,
312                           src_row + (size_t)(ne00 - (uint32_t)lp0) * type_size,
313                           (size_t)lp0 * type_size);
314                } else {
315                    hvx_copy_f32_uu(dst_ptr,
316                                    src_row + (size_t)(ne00 - (uint32_t)lp0) * type_size,
317                                    (uint32_t)lp0);
318                }
319            }
320            hvx_copy_f32_uu(dst_ptr + (size_t)lp0 * type_size, src_row, ne00);
321            if (rp0 > 0) {
322                if ((uint32_t)rp0 < 32) {
323                    memcpy(dst_ptr + ((size_t)lp0 + ne00) * type_size,
324                           src_row,
325                           (size_t)rp0 * type_size);
326                } else {
327                    hvx_copy_f32_uu(dst_ptr + ((size_t)lp0 + ne00) * type_size,
328                                    src_row,
329                                    (uint32_t)rp0);
330                }
331            }
332        } else {
333            for (uint32_t i = 0; i < (uint32_t)lp0; i++) {
334                *(float *)(dst_ptr + i * type_size) =
335                    *(const float *)(src_row + (size_t)(ne00 - (uint32_t)lp0 + i) * nb00);
336            }
337            for (uint32_t i = 0; i < ne00; i++) {
338                *(float *)(dst_ptr + ((size_t)lp0 + i) * type_size) =
339                    *(const float *)(src_row + (size_t)i * nb00);
340            }
341            for (uint32_t i = 0; i < (uint32_t)rp0; i++) {
342                *(float *)(dst_ptr + ((size_t)lp0 + ne00 + i) * type_size) =
343                    *(const float *)(src_row + (size_t)i * nb00);
344            }
345        }
346    }
347
348    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, row_start);
349
350    FARF(HIGH, "pad-hvx-circ %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
351         ith, nth,
352         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
353         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
354         row_start, row_end);
355}
356
357// ---------------------------------------------------------------------------
358// HVX + DMA circular PAD kernel - aligned, double-buffered
359// ---------------------------------------------------------------------------
360
361static void pad_job_per_thread_hvx_circular_dma(unsigned int nth, unsigned int ith, void * data) {
362    const struct htp_pad_context * pctx = (const struct htp_pad_context *) data;
363    struct htp_ops_context * octx = pctx->octx;
364    htp_pad_preamble;
365    htp_pad_dma_preamble;
366
367    // -----------------------------------------------------------------------
368    // Priming phase: push 2 pairs of (dummy_dst_DMA, src_DMA) to seed the
369    // double-buffer pipeline.  Every row is a real src DMA (no null DMAs).
370    // -----------------------------------------------------------------------
371    for (uint32_t ir = row_start, spad_idx = 0; ir < row_end && spad_idx < 2; ir++, spad_idx++) {
372        uint8_t * src_spad_cur = src_spad_base + spad_idx * src_row_size_aligned;
373        uint8_t * dst_spad_cur = dst_spad_base + spad_idx * dst_row_size_aligned;
374
375        dma_queue_push_vtcm_to_ddr(dma,
376            dma_make_ptr((uint8_t *)dst->data, dst_spad_cur),
377            dst_row_size, dst_row_size_aligned, 0);
378
379        uint32_t pi1, pi2, pi3;
380        pad_decompose_row(ir, ne1, ne2, &pi1, &pi2, &pi3);
381        dma_queue_push_ddr_to_vtcm(dma,
382            dma_make_ptr(src_spad_cur, pad_circ_src_row_ptr(src, pi1, pi2, pi3, lp1, lp2, lp3)),
383            src_row_size_aligned, src_row_size, 1);
384    }
385
386    // -----------------------------------------------------------------------
387    // Main loop: pop completed DMAs, assemble circular row in VTCM with
388    // aligned HVX ops, push dst DMA and prefetch src for the next+1 row.
389    // -----------------------------------------------------------------------
390    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
391
392    for (uint32_t ir = row_start; ir < row_end; ir++) {
393        uint8_t * dst_spad_cur = (uint8_t *) dma_queue_pop(dma).src;
394        uint8_t * src_spad_cur = (uint8_t *) dma_queue_pop(dma).dst;
395
396        uint32_t i1, i2, i3;
397        pad_decompose_row(ir, ne1, ne2, &i1, &i2, &i3);
398        uint8_t * dst_ptr = (uint8_t *) dst->data + i1 * nb1 + i2 * nb2 + i3 * nb3;
399
400        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
401        if (lp0 > 0) {
402            uint8_t * dst_left       = dst_spad_cur;
403            const uint8_t * src_left = src_spad_cur + (size_t)(ne00 - (uint32_t)lp0) * type_size;
404            if ((uint32_t)lp0 < 32) {
405                memcpy(dst_left, src_left, (size_t)lp0 * type_size);
406            } else {
407                hvx_copy_f32_uu(dst_left, src_left, (uint32_t)lp0);
408            }
409        }
410
411        {
412            uint8_t * dst_mid = dst_spad_cur + (size_t)lp0 * type_size;
413            if ((uintptr_t)dst_mid % VLEN == 0) {
414                hvx_copy_f32_aa(dst_mid, src_spad_cur, ne00);
415            } else {
416                hvx_copy_f32_ua(dst_mid, src_spad_cur, ne00);
417            }
418        }
419
420        if (rp0 > 0) {
421            uint8_t * dst_right = dst_spad_cur + ((size_t)lp0 + ne00) * type_size;
422            if ((uint32_t)rp0 < 32) {
423                memcpy(dst_right, src_spad_cur, (size_t)rp0 * type_size);
424            } else {
425                if ((uintptr_t)dst_right % VLEN == 0) {
426                    hvx_copy_f32_aa(dst_right, src_spad_cur, (uint32_t)rp0);
427                } else {
428                    hvx_copy_f32_ua(dst_right, src_spad_cur, (uint32_t)rp0);
429                }
430            }
431        }
432        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
433
434        dma_queue_push_vtcm_to_ddr(dma,
435            dma_make_ptr(dst_ptr, dst_spad_cur),
436            dst_row_size, dst_row_size_aligned, 1);
437
438        const uint32_t next_row = ir + 2;
439        if (next_row < row_end) {
440            uint32_t nri1, nri2, nri3;
441            pad_decompose_row(next_row, ne1, ne2, &nri1, &nri2, &nri3);
442            dma_queue_push_ddr_to_vtcm(dma,
443                dma_make_ptr(src_spad_cur,
444                             pad_circ_src_row_ptr(src, nri1, nri2, nri3, lp1, lp2, lp3)),
445                src_row_size_aligned, src_row_size, 1);
446        }
447    }
448
449    dma_queue_flush(dma);
450
451    FARF(HIGH, "pad-hvx-circ-dma %d/%d: (%ux%ux%ux%u) -> (%ux%ux%ux%u) rows %u:%u\n",
452         ith, nth,
453         src->ne[0], src->ne[1], src->ne[2], src->ne[3],
454         dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
455         row_start, row_end);
456}
457
458int op_pad(struct htp_ops_context * octx) {
459    const struct htp_tensor * src0 = octx->src[0];
460    const struct htp_tensor * dst  = octx->dst;
461
462    // Only F32 supported
463    size_t type_size;
464    switch (src0->type) {
465        case HTP_TYPE_F32: type_size = 4; break;
466        default:
467            FARF(ERROR, "pad-hvx: unsupported type %u\n", src0->type);
468            return HTP_STATUS_NO_SUPPORT;
469    }
470
471    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
472        return HTP_STATUS_OK;
473    }
474
475    const int32_t lp0 = octx->op_params[0];
476    const int32_t rp0 = octx->op_params[1];
477    const int32_t lp1 = octx->op_params[2];
478    const int32_t rp1 = octx->op_params[3];
479    const int32_t lp2 = octx->op_params[4];
480    const int32_t rp2 = octx->op_params[5];
481    const int32_t lp3 = octx->op_params[6];
482    const int32_t rp3 = octx->op_params[7];
483    const int32_t circular = octx->op_params[8];
484
485    const uint32_t ne0  = dst->ne[0];
486    const uint32_t ne00 = src0->ne[0];
487
488    const uint32_t total_dst_rows = dst->ne[1] * dst->ne[2] * dst->ne[3];
489    const size_t dst_row_size     = (size_t)ne0 * type_size;
490
491    uint32_t row_start = 0;
492    uint32_t nrows     = total_dst_rows;
493
494    if (octx->ctx->mdev.count > 1) {
495        uint32_t rows_per_chunk = 0;
496        htp_tensor_mdev_rows_per_chunk(dst, type_size, (uint32_t) dst_row_size, &rows_per_chunk);
497        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_dst_rows, rows_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
498        row_start = range.start;
499        nrows     = range.count;
500    }
501
502    if (nrows == 0) {
503        return HTP_STATUS_OK;
504    }
505
506    const uint32_t n_threads = octx->n_threads;
507
508    const size_t src_row_size         = (size_t)ne00 * type_size;
509    const size_t src_row_size_aligned = hex_round_up(src_row_size, VLEN);
510    const size_t dst_row_size_aligned = hex_round_up(dst_row_size, VLEN);
511
512    // Total VTCM needed: 2 buffers (ping+pong) for src and dst, per thread
513    const size_t vtcm_needed = (size_t)n_threads * 2 * (src_row_size_aligned + dst_row_size_aligned);
514
515    const int use_dma = (src0->nb[0] == (uint32_t)type_size) && (ne00 >= 512) &&
516                        (octx->ctx->vtcm_size >= vtcm_needed);
517
518    if (use_dma) {
519        octx->src0_spad.size_per_thread = 2 * src_row_size_aligned;
520        octx->dst_spad.size_per_thread  = 2 * dst_row_size_aligned;
521        octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
522        octx->dst_spad.size  = n_threads * octx->dst_spad.size_per_thread;
523        octx->src0_spad.data = octx->ctx->vtcm_base;
524        octx->dst_spad.data  = octx->src0_spad.data + octx->src0_spad.size;
525        octx->src0_spad.src  = NULL;
526        octx->dst_spad.src   = NULL;
527    }
528
529    struct htp_pad_context pctx = {
530        .octx             = octx,
531        .lp0 = lp0, .rp0 = rp0,
532        .lp1 = lp1, .rp1 = rp1,
533        .lp2 = lp2, .rp2 = rp2,
534        .lp3 = lp3, .rp3 = rp3,
535        .nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div),
536        .total_dst_rows   = nrows,
537        .row_start        = row_start,
538        .type_size        = type_size,
539        .src_row_size         = src_row_size,
540        .src_row_size_aligned = src_row_size_aligned,
541        .dst_row_size         = dst_row_size,
542        .dst_row_size_aligned = dst_row_size_aligned,
543    };
544
545    FARF(HIGH, "pad-hvx%s%s: (%ux%ux%ux%u) -> (%ux%ux%ux%u) pads=(%d,%d,%d,%d,%d,%d,%d,%d)\n",
546         circular ? "-circ" : "",
547         use_dma   ? "-dma"  : "",
548         src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
549         dst->ne[0],  dst->ne[1],  dst->ne[2],  dst->ne[3],
550         lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
551
552    if      (circular && use_dma) { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular_dma, &pctx, n_threads); }
553    else if (circular)            { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_circular,     &pctx, n_threads); }
554    else if (use_dma)             { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx_dma,          &pctx, n_threads); }
555    else                          { work_queue_run(octx->ctx->work_queue, pad_job_per_thread_hvx,              &pctx, n_threads); }
556
557    return HTP_STATUS_OK;
558}