From e2358dffca60f9f6848a40f4dddee7a89654e4fd Mon Sep 17 00:00:00 2001 From: Todor Boinovski Date: Fri, 18 Sep 2026 14:20:48 -0700 Subject: [PATCH] hexagon: im2col update (llama/29103) * ggml-hexagon: accept 1D and padded IM2COL ops * ggml-hexagon: make pure-DDR IM2COL kernel is_2D-aware * ggml-hexagon: extend IM2COL DMA patch-embed fast path to 1D * ggml-hexagon: add blocked-staging general IM2COL DMA kernel --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 12 - ggml/src/ggml-hexagon/htp/im2col-ops.c | 491 +++++++++++++++++++------ 2 files changed, 369 insertions(+), 134 deletions(-) diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index af8013b08..f6f2fdd28 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -5530,11 +5530,6 @@ static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * se const struct ggml_tensor * src1 = op->src[1]; const struct ggml_tensor * dst = op; - const bool is_2D = ((const int32_t *) op->op_params)[6] == 1; - if (!is_2D) { - return false; - } - // For now support F32->F32 and F32->F16 only. if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) { return false; @@ -5544,13 +5539,6 @@ static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * se return false; } - // For now keep padded OPs on CPU. Will revisit once we expand coverage past patch-embed OPs. - const int32_t p0 = ((const int32_t *) op->op_params)[2]; - const int32_t p1 = ((const int32_t *) op->op_params)[3]; - if (p0 != 0 || p1 != 0) { - return false; - } - GGML_UNUSED(sess); return true; } diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.c b/ggml/src/ggml-hexagon/htp/im2col-ops.c index 52bbc37d1..26af14ed5 100644 --- a/ggml/src/ggml-hexagon/htp/im2col-ops.c +++ b/ggml/src/ggml-hexagon/htp/im2col-ops.c @@ -25,17 +25,20 @@ struct htp_im2col_context { uint32_t npatches; // number of patches assigned to this dev uint32_t npatches_per_thread; // patches = N*OH*OW (pure-DDR kernel) - uint32_t pe_row_base; // first N*OH row index assigned to this dev (DMA path) - uint32_t pe_nrows; // number of N*OH rows assigned to this dev (DMA path) - uint32_t pe_rows_per_thread; // N*OH rows per worker - uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256 - uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256 + uint32_t pe_row_base; // first N*OH row index assigned to this dev (DMA path) + uint32_t pe_nrows; // number of N*OH rows assigned to this dev (DMA path) + uint32_t pe_rows_per_thread; // N*OH rows per worker + uint32_t pe_src_row_bytes; // one output row's source: IC*KH*IW*4, rounded 256 + uint32_t pe_dst_row_bytes; // one output row's dst: OW*patch_stride*2, rounded 256 // Patch-embed DMA path VTCM ping-pong. uint8_t * pe_vtcm_src; // base of the 2x src buffers region uint8_t * pe_vtcm_dst; // base of the 2x dst buffers region uint32_t pe_src_size_per_thread; // 2 * pe_src_row_bytes uint32_t pe_dst_size_per_thread; // 2 * pe_dst_row_bytes + + uint32_t pe_owb; // output-col block size + uint32_t pe_wb; // staged source window width }; // Per-op VTCM layout for the patch-embed DMA path @@ -59,83 +62,253 @@ static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L->total_bytes = L->off_dst + L->dst_bytes_per_thread * n_threads; } -#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ - static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ - struct htp_ops_context * octx = ictx->octx; \ - struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ - const struct htp_tensor * restrict src0 = octx->src[0]; \ - const struct htp_tensor * restrict src1 = octx->src[1]; \ - const struct htp_tensor * restrict dst = octx->dst; \ - const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; \ - const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; \ - const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; \ - const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \ - const uint32_t KH = src0->ne[1], KW = src0->ne[0]; \ - const uint32_t OH = dst->ne[2]; \ - const uint32_t OW = dst->ne[1]; \ - const uint32_t patch_stride = IC * KH * KW; \ - const float * restrict src_data = (const float *) src1->data; \ - DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ - const uint32_t patch_end = ictx->patch_base + ictx->npatches; \ - const uint32_t patch_start = ictx->patch_base + ictx->npatches_per_thread * ith; \ - const uint32_t patch_stop = MIN(patch_start + ictx->npatches_per_thread, patch_end);\ - if (patch_start >= patch_stop) { \ - return; \ - } \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ - for (uint32_t p = patch_start; p < patch_stop; p++) { \ - const uint32_t iow = p % OW; \ - const uint32_t ioh = (p / OW) % OH; \ - const uint32_t in = p / (OW * OH); \ - DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \ - for (uint32_t iic = 0; iic < IC; iic++) { \ - const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \ - for (uint32_t ikh = 0; ikh < KH; ikh++) { \ - const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ - DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ - if (iih < 0 || iih >= (int32_t) IH) { \ - SPLAT_FN(out_run, 0.0f, KW); \ - continue; \ - } \ - const int32_t iiw0 = (int32_t) iow * s0 - p0; \ - const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \ - if (d0 == 1) { \ - /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \ - const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ - int32_t hi = (int32_t) IW - iiw0; \ - if (hi > (int32_t) KW) { \ - hi = (int32_t) KW; \ - } \ - if (hi <= lo) { \ - SPLAT_FN(out_run, 0.0f, KW); \ - } else { \ - if (lo > 0) { \ - SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ - } \ - COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \ - (uint32_t) (hi - lo)); \ - if (hi < (int32_t) KW) { \ - SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ - } \ - } \ - continue; \ - } \ - for (uint32_t ikw = 0; ikw < KW; ikw++) { \ - const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \ - out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \ - (DST_CTYPE) 0.0f : \ - (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \ - } \ - } \ - } \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ +#define IM2COL_PATCHEMBED_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src0 = octx->src[0]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t s0 = octx->op_params[0]; \ + const int32_t s1 = octx->op_params[1]; \ + const int32_t p0 = octx->op_params[2]; \ + const int32_t p1 = octx->op_params[3]; \ + const int32_t d0 = octx->op_params[4]; \ + const int32_t d1 = octx->op_params[5]; \ + const int32_t is_2D = octx->op_params[6] == 1; \ + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ + const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ + const uint32_t IH = is_2D ? src1->ne[1] : 1; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = is_2D ? src0->ne[1] : 1; \ + const uint32_t KW = src0->ne[0]; \ + const uint32_t OH = is_2D ? dst->ne[2] : 1; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + const uint32_t patch_end = ictx->patch_base + ictx->npatches; \ + const uint32_t patch_start = ictx->patch_base + ictx->npatches_per_thread * ith; \ + const uint32_t patch_stop = MIN(patch_start + ictx->npatches_per_thread, patch_end); \ + if (patch_start >= patch_stop) { \ + return; \ + } \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ + for (uint32_t p = patch_start; p < patch_stop; p++) { \ + const uint32_t iow = p % OW; \ + const uint32_t ioh = (p / OW) % OH; \ + const uint32_t in = p / (OW * OH); \ + DST_CTYPE * restrict dst_patch = dst_data + (uint64_t) p * patch_stride; \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + const float * restrict src_plane = src_data + ((uint64_t) in * IC + iic) * IH * IW; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + DST_CTYPE * restrict out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (iih < 0 || iih >= (int32_t) IH) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const int32_t iiw0 = (int32_t) iow * s0 - p0; \ + const float * restrict src_run = src_plane + (uint64_t) iih * IW + iiw0; \ + if (d0 == 1) { \ + /* contiguous source run: [lo,hi) is in-bounds, tails are zero pad */ \ + const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ + int32_t hi = (int32_t) IW - iiw0; \ + if (hi > (int32_t) KW) { \ + hi = (int32_t) KW; \ + } \ + if (hi <= lo) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + } else { \ + if (lo > 0) { \ + SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ + } \ + COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (src_run + lo), \ + (uint32_t) (hi - lo)); \ + if (hi < (int32_t) KW) { \ + SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ + } \ + } \ + continue; \ + } \ + for (uint32_t ikw = 0; ikw < KW; ikw++) { \ + const int32_t iiw = (int32_t) iow * s0 + (int32_t) ikw * d0 - p0; \ + out_run[ikw] = (iiw < 0 || iiw >= (int32_t) IW) ? \ + (DST_CTYPE) 0.0f : \ + (DST_CTYPE) src_plane[(uint64_t) iih * IW + iiw]; \ + } \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, patch_start); \ } IM2COL_PATCHEMBED_BODY(im2col_patchembed_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "f32-f16") IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "f32-f32") +// Software-pipelined 2-deep: while HVX computes block bi from buffer slot +// (bi&1), the DMA engine stages block bi+1 into the other slot concurrently. +// A single dma_queue_flush per iteration (after issuing the next stage-in and +// this block's store-out) waits for both - safe because the ring is strict +// FIFO and each buffer slot is only reused after its prior consumer (compute +// or store-out) already finished in program order. +#define IM2COL_BLOCKED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ + static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ + struct htp_ops_context * octx = ictx->octx; \ + struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ + const struct htp_tensor * restrict src1 = octx->src[1]; \ + const struct htp_tensor * restrict dst = octx->dst; \ + const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1]; \ + const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; \ + const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; \ + const int32_t is_2D = octx->op_params[6] == 1; \ + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ + const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ + const uint32_t IH = is_2D ? src1->ne[1] : 1; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = is_2D ? dst->ne[2] : 1; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t owb = ictx->pe_owb, Wb = ictx->pe_wb; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + dma_queue * dmaq = octx->ctx->dma[ith]; \ + uint8_t * srcb_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dstb_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb2[2] = { (float *) srcb_base, (float *) (srcb_base + ictx->pe_src_row_bytes) }; \ + DST_CTYPE * dstb2[2] = { (DST_CTYPE *) dstb_base, (DST_CTYPE *) (dstb_base + ictx->pe_dst_row_bytes) }; \ + const uint32_t nrows = N * OH; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, nrows); \ + if (row_start >= row_end) \ + return; \ + const uint32_t nbpr = (OW + owb - 1) / owb; \ + const uint32_t nrows_local = row_end - row_start; \ + const uint32_t total_blocks = nrows_local * nbpr; \ + for (uint32_t bi = 0; bi < total_blocks; bi++) { \ + const uint32_t buf = bi & 1u; \ + float * srcb = srcb2[buf]; \ + DST_CTYPE * dstb = dstb2[buf]; \ + const uint32_t r = row_start + bi / nbpr; \ + const uint32_t in = r / OH; \ + const uint32_t ioh = r % OH; \ + const uint32_t c0 = (bi % nbpr) * owb; \ + const uint32_t nb = MIN(owb, OW - c0); \ + const int32_t win0 = (int32_t) c0 * s0 - p0; \ + if (bi == 0) { \ + /* prologue: stage block 0 and wait - nothing to overlap with yet */ \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + if (iih < 0 || iih >= (int32_t) IH) \ + continue; \ + const int32_t lo = win0 < 0 ? -win0 : 0; \ + int32_t hi = (int32_t) IW - win0; \ + if (hi > (int32_t) Wb) \ + hi = (int32_t) Wb; \ + if (hi <= lo) \ + continue; \ + const uint32_t cpw = (uint32_t) (hi - lo); \ + float * vdst = srcb + (uint64_t) ikh * Wb + (uint32_t) lo; \ + const float * vsrc = src_data + ((uint64_t) (in * IC) * IH + iih) * IW + (win0 + lo); \ + while (!dma_queue_push(dmaq, dma_make_ptr((uint8_t *) vdst, (const uint8_t *) vsrc), \ + (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ + cpw * sizeof(float), IC)) { \ + dma_queue_pop(dmaq); \ + } \ + } \ + dma_queue_flush(dmaq); \ + } \ + if (bi + 1 < total_blocks) { \ + /* prefetch: stage block bi+1 into the other slot; overlaps with this block's compute below */ \ + const uint32_t nbuf = 1u - buf; \ + float * nsrcb = srcb2[nbuf]; \ + const uint32_t nr = row_start + (bi + 1) / nbpr; \ + const uint32_t nin = nr / OH; \ + const uint32_t nioh = nr % OH; \ + const uint32_t nc0 = ((bi + 1) % nbpr) * owb; \ + const int32_t nwin0 = (int32_t) nc0 * s0 - p0; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) nioh * s1 + (int32_t) ikh * d1 - p1; \ + if (iih < 0 || iih >= (int32_t) IH) \ + continue; \ + const int32_t lo = nwin0 < 0 ? -nwin0 : 0; \ + int32_t hi = (int32_t) IW - nwin0; \ + if (hi > (int32_t) Wb) \ + hi = (int32_t) Wb; \ + if (hi <= lo) \ + continue; \ + const uint32_t cpw = (uint32_t) (hi - lo); \ + float * vdst = nsrcb + (uint64_t) ikh * Wb + (uint32_t) lo; \ + const float * vsrc = src_data + ((uint64_t) (nin * IC) * IH + iih) * IW + (nwin0 + lo); \ + while (!dma_queue_push(dmaq, dma_make_ptr((uint8_t *) vdst, (const uint8_t *) vsrc), \ + (size_t) KH * Wb * sizeof(float), (size_t) IH * IW * sizeof(float), \ + cpw * sizeof(float), IC)) { \ + dma_queue_pop(dmaq); \ + } \ + } \ + } \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + for (uint32_t j = 0; j < nb; j++) { \ + const uint32_t iow = c0 + j; \ + DST_CTYPE * dst_patch = dstb + (uint64_t) j * patch_stride; \ + const int32_t iiw0 = (int32_t) iow * s0 - p0; \ + for (uint32_t ikh = 0; ikh < KH; ikh++) { \ + const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \ + const int okh = (iih >= 0 && iih < (int32_t) IH); \ + for (uint32_t iic = 0; iic < IC; iic++) { \ + DST_CTYPE * out_run = dst_patch + iic * (KH * KW) + ikh * KW; \ + if (!okh) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + continue; \ + } \ + const float * vrow = srcb + ((uint64_t) (iic * KH + ikh)) * Wb; /* col win0 at idx 0*/ \ + if (d0 == 1) { \ + /* contiguous run within the staged window: [lo,hi) in-bounds, tails zero pad */ \ + const int32_t lo = iiw0 < 0 ? -iiw0 : 0; \ + int32_t hi = (int32_t) IW - iiw0; \ + if (hi > (int32_t) KW) { \ + hi = (int32_t) KW; \ + } \ + if (hi <= lo) { \ + SPLAT_FN(out_run, 0.0f, KW); \ + } else { \ + if (lo > 0) { \ + SPLAT_FN(out_run, 0.0f, (uint32_t) lo); \ + } \ + COPY_FN((uint8_t *) (out_run + lo), (const uint8_t *) (vrow + (iiw0 + lo - win0)), \ + (uint32_t) (hi - lo)); \ + if (hi < (int32_t) KW) { \ + SPLAT_FN(out_run + hi, 0.0f, (KW - (uint32_t) hi)); \ + } \ + } \ + continue; \ + } \ + for (uint32_t ikw = 0; ikw < KW; ikw++) { \ + const int32_t iiw = iiw0 + (int32_t) ikw * d0; \ + out_run[ikw] = \ + (iiw < 0 || iiw >= (int32_t) IW) ? (DST_CTYPE) 0.0f : (DST_CTYPE) vrow[iiw - win0]; \ + } \ + } \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ + DST_CTYPE * ddr = dst_data + ((uint64_t) (in * OH + ioh) * OW + c0) * patch_stride; \ + dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr, (uint8_t *) dstb), \ + nb * patch_stride * (DST_ELEM), nb * patch_stride * (DST_ELEM), 1); \ + dma_queue_flush(dmaq); \ + } \ + } +IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "blk-dma-f16") +IM2COL_BLOCKED_DMA_BODY(im2col_blocked_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "blk-dma-f32") + +// Exact-tiling patch-embed DMA fast path (s0==KW, p0=0, d0=1; and 2D s1==KH, +// p1=0, d1=1). Intentionally reads no stride/pad/dilation params so the inner +// copy stays tight and fully hoisted - do NOT graft the general gather in here. #define IM2COL_PATCHEMBED_DMA_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN, DST_ELEM, TAG) \ static void FNAME(unsigned int nth, unsigned int ith, void * data) { \ struct htp_im2col_context * ictx = (struct htp_im2col_context *) data; \ @@ -143,21 +316,27 @@ IM2COL_PATCHEMBED_BODY(im2col_patchembed_f32_thread, float, hvx_copy_f32_uu, hvx struct htp_thread_trace * restrict tr = &octx->ctx->trace[ith]; \ const struct htp_tensor * restrict src1 = octx->src[1]; \ const struct htp_tensor * restrict dst = octx->dst; \ - const uint32_t N = src1->ne[3], IC = src1->ne[2], IH = src1->ne[1], IW = src1->ne[0]; \ - const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; \ - const uint32_t OH = dst->ne[2], OW = dst->ne[1]; \ - const uint32_t patch_stride = IC * KH * KW; \ - const float * restrict src_data = (const float *) src1->data; \ - DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ - dma_queue * dmaq = octx->ctx->dma[ith]; \ - uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ - uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ - float * srcb = (float *) src_base; \ - DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ - const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows; \ - const uint32_t per_thread = ictx->pe_rows_per_thread; \ - const uint32_t row_start = ictx->pe_row_base + per_thread * ith; \ - const uint32_t row_end = MIN(row_start + per_thread, row_end_max); \ + const int32_t is_2D = octx->op_params[6] == 1; \ + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; \ + const uint32_t IC = is_2D ? src1->ne[2] : src1->ne[1]; \ + const uint32_t IH = is_2D ? src1->ne[1] : 1; \ + const uint32_t IW = src1->ne[0]; \ + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; \ + const uint32_t KW = octx->src[0]->ne[0]; \ + const uint32_t OH = is_2D ? dst->ne[2] : 1; \ + const uint32_t OW = dst->ne[1]; \ + const uint32_t patch_stride = IC * KH * KW; \ + const float * restrict src_data = (const float *) src1->data; \ + DST_CTYPE * restrict dst_data = (DST_CTYPE *) dst->data; \ + dma_queue * dmaq = octx->ctx->dma[ith]; \ + uint8_t * src_base = ictx->pe_vtcm_src + ith * ictx->pe_src_size_per_thread; \ + uint8_t * dst_base = ictx->pe_vtcm_dst + ith * ictx->pe_dst_size_per_thread; \ + float * srcb = (float *) src_base; \ + DST_CTYPE * dstb = (DST_CTYPE *) dst_base; \ + const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows; \ + const uint32_t per_thread = ictx->pe_rows_per_thread; \ + const uint32_t row_start = ictx->pe_row_base + per_thread * ith; \ + const uint32_t row_end = MIN(row_start + per_thread, row_end_max); \ if (row_start >= row_end) \ return; \ for (uint32_t r = row_start; r < row_end; r++) { \ @@ -209,21 +388,30 @@ static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) { const int32_t p0 = octx->op_params[2], p1 = octx->op_params[3]; const int32_t d0 = octx->op_params[4], d1 = octx->op_params[5]; const int is_2D = octx->op_params[6] == 1; - if (!is_2D) { - return false; - } if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) { return false; } - const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; - if (s0 != (int32_t) KW || s1 != (int32_t) KH) { - return false; // non-overlapping + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; + const uint32_t KW = octx->src[0]->ne[0]; + if (s0 != (int32_t) KW) { + return false; // non-overlapping (width) } - if (p0 != 0 || p1 != 0) { - return false; // no padding + if (p0 != 0) { + return false; // no padding (width) } - if (d0 != 1 || d1 != 1) { - return false; // no dilation + if (d0 != 1) { + return false; // no dilation (width) + } + if (is_2D) { + if (s1 != (int32_t) KH) { + return false; // non-overlapping (height) + } + if (p1 != 0) { + return false; // no padding (height) + } + if (d1 != 1) { + return false; // no dilation (height) + } } return true; } @@ -233,8 +421,11 @@ static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) { static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx, struct htp_im2col_context * ictx, uint32_t n_threads) { - const uint32_t IC = octx->src[1]->ne[2], IW = octx->src[1]->ne[0]; - const uint32_t KH = octx->src[0]->ne[1], KW = octx->src[0]->ne[0]; + const int32_t is_2D = octx->op_params[6] == 1; + const uint32_t IC = is_2D ? octx->src[1]->ne[2] : octx->src[1]->ne[1]; + const uint32_t IW = octx->src[1]->ne[0]; + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; + const uint32_t KW = octx->src[0]->ne[0]; const uint32_t OW = octx->dst->ne[1]; const uint32_t patch_stride = IC * KH * KW; @@ -257,6 +448,45 @@ static bool im2col_patchembed_dma_fits(struct htp_ops_context * octx, return true; } +// Sizes a per-thread 2x(src,dst) VTCM ping-pong for the blocked general kernel. +// Stages Wb=(owb-1)*s0+(KW-1)*d0+1 source cols per (iic,ikh) row and owb patches +// of dst. Picks the largest owb that fits; returns false if even owb=1 does not. +static bool im2col_blocked_dma_fits(struct htp_ops_context * octx, + struct htp_im2col_context * ictx, + uint32_t n_threads) { + const int32_t is_2D = octx->op_params[6] == 1; + const int32_t s0 = octx->op_params[0]; + const int32_t d0 = octx->op_params[4]; + const uint32_t IC = is_2D ? octx->src[1]->ne[2] : octx->src[1]->ne[1]; + const uint32_t KH = is_2D ? octx->src[0]->ne[1] : 1; + const uint32_t KW = octx->src[0]->ne[0]; + const uint32_t OW = octx->dst->ne[1]; + const uint32_t patch_stride = IC * KH * KW; + const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); + + for (uint32_t owb = (OW < 256 ? OW : 256); owb >= 1; owb--) { + const uint32_t Wb = (owb - 1) * (uint32_t) s0 + (KW - 1) * (uint32_t) d0 + 1; + const uint32_t src_row_bytes = hex_round_up(IC * KH * Wb * sizeof(float), 256); + const uint32_t dst_row_bytes = hex_round_up(owb * patch_stride * dst_elem, 256); + struct htp_im2col_vtcm_layout L; + htp_im2col_vtcm_layout_build(&L, src_row_bytes, dst_row_bytes, n_threads); + if (L.total_bytes <= octx->ctx->vtcm_size) { + uint8_t * const base = octx->ctx->vtcm_base; + ictx->pe_owb = owb; + ictx->pe_wb = Wb; + ictx->pe_src_row_bytes = src_row_bytes; + ictx->pe_dst_row_bytes = dst_row_bytes; + ictx->pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src); + ictx->pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); + ictx->pe_src_size_per_thread = (uint32_t) L.src_bytes_per_thread; + ictx->pe_dst_size_per_thread = (uint32_t) L.dst_bytes_per_thread; + return true; + } + if (owb == 1) break; // avoid unsigned underflow + } + return false; +} + int op_im2col(struct htp_ops_context * octx) { const struct htp_tensor * src1 = octx->src[1]; const struct htp_tensor * dst = octx->dst; @@ -270,8 +500,9 @@ int op_im2col(struct htp_ops_context * octx) { return HTP_STATUS_OK; } - const uint32_t N = src1->ne[3]; - const uint32_t OH = dst->ne[2]; + const int32_t is_2D = octx->op_params[6] == 1; + const uint32_t N = is_2D ? src1->ne[3] : src1->ne[2]; + const uint32_t OH = is_2D ? dst->ne[2] : 1; const uint32_t OW = dst->ne[1]; const uint32_t total_patches = N * OH * OW; const uint32_t total_rows = N * OH; @@ -280,8 +511,11 @@ int op_im2col(struct htp_ops_context * octx) { uint32_t npatches = total_patches; if (octx->ctx->mdev.count > 1) { const uint32_t patch_size = dst->nb[1]; - const uint32_t patches_per_chunk = (patch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(patch_size, HEX_L2_LINE_SIZE)) : 1; - 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); + const uint32_t patches_per_chunk = + (patch_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(patch_size, HEX_L2_LINE_SIZE)) : 1; + 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); patch_base = range.start; npatches = range.count; } @@ -290,8 +524,11 @@ int op_im2col(struct htp_ops_context * octx) { uint32_t nrows = total_rows; if (octx->ctx->mdev.count > 1) { const uint32_t row_size = dst->nb[2]; - const uint32_t rows_per_chunk = (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1; - 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); + const uint32_t rows_per_chunk = + (row_size > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(row_size, HEX_L2_LINE_SIZE)) : 1; + 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); row_base = range.start; nrows = range.count; } @@ -309,22 +546,32 @@ int op_im2col(struct htp_ops_context * octx) { ictx.npatches_per_thread = (npatches + n_threads - 1) / n_threads; // Clean non-overlapping patch-embed -> DMA kernel (if it fits VTCM); - // everything else (padding/dilation/stride edges) -> pure-DDR kernel. - if (im2col_use_patchembed_dma(octx) && nrows > 0) { + // everything else (padding/dilation/stride edges) -> blocked-staging DMA + // kernel; if neither fits VTCM -> pure-DDR kernel. + if (nrows > 0) { const uint32_t pth = MIN(octx->n_threads, nrows); - if (pth > 0 && im2col_patchembed_dma_fits(octx, &ictx, pth)) { - ictx.pe_row_base = row_base; - ictx.pe_nrows = nrows; - ictx.pe_rows_per_thread = (nrows + pth - 1) / pth; - if (dst->type == HTP_TYPE_F16) { - work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth); - } else { - work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_f32_thread, &ictx, pth); + if (pth > 0) { + ictx.pe_row_base = row_base; + ictx.pe_nrows = nrows; + const bool exact = im2col_use_patchembed_dma(octx); + if (exact && im2col_patchembed_dma_fits(octx, &ictx, pth)) { + ictx.pe_rows_per_thread = (nrows + pth - 1) / pth; + work_queue_run(octx->ctx->work_queue, + dst->type == HTP_TYPE_F16 ? im2col_patchembed_dma_thread + : im2col_patchembed_dma_f32_thread, &ictx, pth); + return HTP_STATUS_OK; + } + if (!exact && im2col_blocked_dma_fits(octx, &ictx, pth)) { + ictx.pe_rows_per_thread = (nrows + pth - 1) / pth; + work_queue_run(octx->ctx->work_queue, + dst->type == HTP_TYPE_F16 ? im2col_blocked_dma_thread + : im2col_blocked_dma_f32_thread, &ictx, pth); + return HTP_STATUS_OK; } - return HTP_STATUS_OK; } - // else: doesn't fit -> fall through to the pure-DDR kernel below. } + // Fall through to pure-DDR. + if (npatches == 0) { return HTP_STATUS_OK;