hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

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)eafe15a5e2026-09-11

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 <hexagon_protos.h>
7#include <hexagon_types.h>
8#include <string.h>
9
10#include "hex-common.h"
11
12#define GGML_COMMON_DECL_C
13#include "ggml-common.h"
14#include "htp-ctx.h"
15#include "htp-ops.h"
16#include "hvx-utils.h"
17#include "hex-dma.h"
18#include "hex-profile.h"
19#include "htp-vtcm.h"
20#include "htp-tensor.h"
21
22struct htp_im2col_context {
23    struct htp_ops_context * octx;
24    uint32_t                 patch_base;           // first patch index assigned to this dev
25    uint32_t                 npatches;             // number of patches assigned to this dev
26    uint32_t                 npatches_per_thread;  // patches = N*OH*OW (pure-DDR kernel)
27
28    uint32_t pe_row_base;              // first N*OH row index assigned to this dev (DMA path)
29    uint32_t pe_nrows;                 // number of N*OH rows assigned to this dev (DMA path)
30    uint32_t pe_rows_per_thread;       // N*OH rows per worker
31    uint32_t pe_src_row_bytes;         // one output row's source: IC*KH*IW*4, rounded 256
32    uint32_t pe_dst_row_bytes;         // one output row's dst: OW*patch_stride*2, rounded 256
33
34    // Patch-embed DMA path VTCM ping-pong.
35    uint8_t * pe_vtcm_src;                         // base of the 2x src buffers region
36    uint8_t * pe_vtcm_dst;                         // base of the 2x dst buffers region
37    uint32_t  pe_src_size_per_thread;              // 2 * pe_src_row_bytes
38    uint32_t  pe_dst_size_per_thread;              // 2 * pe_dst_row_bytes
39};
40
41// Per-op VTCM layout for the patch-embed DMA path
42struct htp_im2col_vtcm_layout {
43    size_t off_src;
44    size_t off_dst;
45    size_t src_bytes_per_thread;
46    size_t dst_bytes_per_thread;
47    size_t total_bytes;
48};
49
50static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L,
51                                                size_t                          src_row_bytes,
52                                                size_t                          dst_row_bytes,
53                                                uint32_t                        n_threads) {
54    L->src_bytes_per_thread = 2 * src_row_bytes;
55    L->dst_bytes_per_thread = 2 * dst_row_bytes;
56
57    L->off_src     = 0;
58    L->off_dst     = L->off_src + L->src_bytes_per_thread * n_threads;
59    L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads;
60}
61
62#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                        \
63    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                  \
64        struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                     \
65        struct htp_ops_context *    octx        = ictx->octx;                                             \
66        struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                 \
67        const struct htp_tensor * restrict src0 = octx->src[0];                                           \
68        const struct htp_tensor * restrict src1 = octx->src[1];                                           \
69        const struct htp_tensor * restrict dst  = octx->dst;                                              \
70        const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];                                   \
71        const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];                                   \
72        const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];                                   \
73        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];             \
74        const uint32_t KH                       = src0->ne[1], KW = src0->ne[0];                          \
75        const uint32_t OH                       = dst->ne[2];                                             \
76        const uint32_t OW                       = dst->ne[1];                                             \
77        const uint32_t patch_stride             = IC * KH * KW;                                           \
78        const float * restrict src_data         = (const float *) src1->data;                             \
79        DST_CTYPE * restrict dst_data           = (DST_CTYPE *) dst->data;                                \
80        const uint32_t patch_end                = ictx->patch_base + ictx->npatches;                      \
81        const uint32_t patch_start              = ictx->patch_base + ictx->npatches_per_thread * ith;     \
82        const uint32_t patch_stop               = MIN(patch_start + ictx->npatches_per_thread, patch_end);\
83        if (patch_start >= patch_stop) {                                                                  \
84            return;                                                                                       \
85        }                                                                                                 \
86        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                   \
87        for (uint32_t p = patch_start; p < patch_stop; p++) {                                             \
88            const uint32_t iow             = p % OW;                                                      \
89            const uint32_t ioh             = (p / OW) % OH;                                               \
90            const uint32_t in              = p / (OW * OH);                                               \
91            DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride;                      \
92            for (uint32_t iic = 0; iic < IC; iic++) {                                                     \
93                const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW;       \
94                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                 \
95                    const int32_t iih            = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;          \
96                    DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW;                \
97                    if (iih < 0 || iih >= (int32_t) IH) {                                                 \
98                        SPLAT_FN(out_run, 0.0f, KW);                                                      \
99                        continue;                                                                         \
100                    }                                                                                     \
101                    const int32_t iiw0             = (int32_t) iow * s0 - p0;                             \
102                    const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0;              \
103                    if (d0 == 1) {                                                                        \
104                        /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */             \
105                        const int32_t lo = iiw0 < 0 ? -iiw0 : 0;                                          \
106                        int32_t       hi = (int32_t) IW - iiw0;                                           \
107                        if (hi > (int32_t) KW) {                                                          \
108                            hi = (int32_t) KW;                                                            \
109                        }                                                                                 \
110                        if (hi <= lo) {                                                                   \
111                            SPLAT_FN(out_run, 0.0f, KW);                                                  \
112                        } else {                                                                          \
113                            if (lo > 0) {                                                                 \
114                                SPLAT_FN(out_run, 0.0f, (uint32_t) lo);                                   \
115                            }                                                                             \
116                            COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo),         \
117                                    (uint32_t) (hi - lo));                                                \
118                            if (hi < (int32_t) KW) {                                                      \
119                                SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi));                       \
120                            }                                                                             \
121                        }                                                                                 \
122                        continue;                                                                         \
123                    }                                                                                     \
124                    for (uint32_t ikw = 0; ikw < KW; ikw++) {                                             \
125                        const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0;                 \
126                        out_run[ikw]      = (iiw < 0 || iiw >= (int32_t) IW) ?                            \
127                                                (DST_CTYPE) 0.0f :                                        \
128                                                (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw];         \
129                    }                                                                                     \
130                }                                                                                         \
131            }                                                                                             \
132        }                                                                                                 \
133        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start);                                    \
134    }
135
136IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16")
137IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32")
138
139#define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG)                               \
140    static void FNAME(unsigned int nth, unsigned int ith, void * data) {                                             \
141        struct htp_im2col_context * ictx        = (struct htp_im2col_context *) data;                                \
142        struct htp_ops_context *    octx        = ictx->octx;                                                        \
143        struct htp_thread_trace * restrict tr   = &octx->ctx->trace[ith];                                            \
144        const struct htp_tensor * restrict src1 = octx->src[1];                                                      \
145        const struct htp_tensor * restrict dst  = octx->dst;                                                         \
146        const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0];                        \
147        const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];                                           \
148        const uint32_t OH = dst->ne[2], OW = dst->ne[1];                                                             \
149        const uint32_t patch_stride     = IC * KH * KW;                                                              \
150        const float * restrict src_data = (const float *) src1->data;                                                \
151        DST_CTYPE * restrict dst_data   = (DST_CTYPE *) dst->data;                                                   \
152        dma_queue *    dmaq             = octx->ctx->dma[ith];                                                       \
153        uint8_t *      src_base         = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread;                    \
154        uint8_t *      dst_base         = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread;                    \
155        float *        srcb             = (float *) src_base;                                                        \
156        DST_CTYPE *    dstb             = (DST_CTYPE *) dst_base;                                                    \
157        const uint32_t row_end_max      = ictx->pe_row_base + ictx->pe_nrows;                                        \
158        const uint32_t per_thread       = ictx->pe_rows_per_thread;                                                  \
159        const uint32_t row_start        = ictx->pe_row_base + per_thread * ith;                                      \
160        const uint32_t row_end          = MIN(row_start + per_thread, row_end_max);                                  \
161        if (row_start >= row_end)                                                                                    \
162            return;                                                                                                  \
163        for (uint32_t r = row_start; r < row_end; r++) {                                                             \
164            const uint32_t in  = r / OH;                                                                             \
165            const uint32_t ioh = r % OH;                                                                             \
166            for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                                \
167                int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                          \
168                int     ok  = (iih >= 0 && iih < (int32_t) IH);                                                      \
169                for (uint32_t iic = 0; iic < IC; iic++) {                                                            \
170                    float *       vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW;                                  \
171                    const float * _vsrc =                                                                            \
172                        ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst;       \
173                    dma_queue_push_ddr_to_vtcm(                                                                      \
174                        dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \
175                        IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0);                                         \
176                }                                                                                                    \
177            }                                                                                                        \
178            for (uint32_t i = 0; i < IC * KH; i++)                                                                   \
179                dma_queue_pop(dmaq);                                                                                 \
180            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                    \
181            for (uint32_t iow = 0; iow < OW; iow++) {                                                                \
182                DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride;                                        \
183                for (uint32_t ikh = 0; ikh < KH; ikh++) {                                                            \
184                    int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;                                      \
185                    for (uint32_t iic = 0; iic < IC; iic++) {                                                        \
186                        DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW;                                \
187                        if (iih < 0 || iih >= (int32_t) IH) {                                                        \
188                            SPLAT_FN(out_run, 0.0f, KW);                                                             \
189                            continue;                                                                                \
190                        }                                                                                            \
191                        const float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW;     \
192                        COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW);                                 \
193                    }                                                                                                \
194                }                                                                                                    \
195            }                                                                                                        \
196            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r);                                                     \
197            DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride;                       \
198            dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb),                    \
199                                       OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1);           \
200            dma_queue_flush(dmaq);                                                                                   \
201        }                                                                                                            \
202    }
203
204IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread,     __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16")
205IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float,  hvx_copy_f32_uu,     hvx_splat_f32_u, sizeof(float),  "pe-dma-f32")
206
207static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) {
208    const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1];
209    const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3];
210    const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5];
211    const int     is_2D = octx->op_params[6] == 1;
212    if (!is_2D) {
213        return false;
214    }
215    if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
216        return false;
217    }
218    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
219    if (s0 != (int32_t) KW || s1 != (int32_t) KH) {
220        return false;  // non-overlapping
221    }
222    if (p0 != 0 || p1 != 0) {
223        return false;  // no padding
224    }
225    if (d0 != 1 || d1 != 1) {
226        return false;  // no dilation
227    }
228    return true;
229}
230
231// Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path.
232// Returns false if it doesn't fit the VTCM budget (caller falls back).
233static bool im2col_patchembed_dma_fits(struct htp_ops_context *    octx,
234                                       struct htp_im2col_context * ictx,
235                                       uint32_t                    n_threads) {
236    const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0];
237    const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0];
238    const uint32_t OW           = octx->dst->ne[1];
239    const uint32_t patch_stride = IC * KH * KW;
240
241    ictx->pe_src_row_bytes  = hex_round_up(IC * KH * IW * sizeof(float), 256);
242    const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float);
243    ictx->pe_dst_row_bytes  = hex_round_up(OW * patch_stride * dst_elem, 256);
244
245    // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst.
246    struct htp_im2col_vtcm_layout L;
247    htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads);
248    if (L.total_bytes > octx->ctx->vtcm_size) {
249        return false;
250    }
251
252    uint8_t * const base        = octx->ctx->vtcm_base;
253    ictx->pe_vtcm_src           = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src);
254    ictx->pe_vtcm_dst           = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
255    ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread;
256    ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread;
257    return true;
258}
259
260int op_im2col(struct htp_ops_context * octx) {
261    const struct htp_tensor * src1 = octx->src[1];
262    const struct htp_tensor * dst  = octx->dst;
263
264    if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) {
265        FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported");
266        return HTP_STATUS_NO_SUPPORT;
267    }
268
269    if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) {
270        return HTP_STATUS_OK;
271    }
272
273    const uint32_t N             = src1->ne[3];
274    const uint32_t OH            = dst->ne[2];
275    const uint32_t OW            = dst->ne[1];
276    const uint32_t total_patches = N * OH * OW;
277    const uint32_t total_rows    = N * OH;
278
279    uint32_t patch_base = 0;
280    uint32_t npatches   = total_patches;
281    if (octx->ctx->mdev.count > 1) {
282        const uint32_t patch_size = dst->nb[1];
283        const uint32_t patches_per_chunk = (patch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(patch_size, HEX_L2_LINE_SIZE)) : 1;
284        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_patches, htp_tensor_mdev_data_aligned(dst) ? patches_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
285        patch_base = range.start;
286        npatches   = range.count;
287    }
288
289    uint32_t row_base = 0;
290    uint32_t nrows    = total_rows;
291    if (octx->ctx->mdev.count > 1) {
292        const uint32_t row_size = dst->nb[2];
293        const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1;
294        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_rows, htp_tensor_mdev_data_aligned(dst) ? rows_per_chunk : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
295        row_base = range.start;
296        nrows    = range.count;
297    }
298
299    if (npatches == 0 && nrows == 0) {
300        return HTP_STATUS_OK;
301    }
302
303    const uint32_t n_threads = MIN(octx->n_threads, MAX(npatches, 1));
304
305    struct htp_im2col_context ictx = { 0 };
306    ictx.octx                = octx;
307    ictx.patch_base          = patch_base;
308    ictx.npatches            = npatches;
309    ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads;
310
311    // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM);
312    // everything else (padding/dilation/stride edges) -> pure-DDR kernel.
313    if (im2col_use_patchembed_dma(octx) && nrows > 0) {
314        const uint32_t pth = MIN(octx->n_threads, nrows);
315        if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) {
316            ictx.pe_row_base        = row_base;
317            ictx.pe_nrows           = nrows;
318            ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
319            if (dst->type == HTP_TYPE_F16) {
320                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
321            } else {
322                work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth);
323            }
324            return HTP_STATUS_OK;
325        }
326        // else: doesn't fit -> fall through to the pure-DDR kernel below.
327    }
328
329    if (npatches == 0) {
330        return HTP_STATUS_OK;
331    }
332
333    if (dst->type == HTP_TYPE_F16) {
334        work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads);
335    } else {
336        work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads);
337    }
338    return HTP_STATUS_OK;
339}