#pragma clang diagnostic ignored "-Wunused-variable" #pragma clang diagnostic ignored "-Wunused-function" #pragma clang diagnostic ignored "-Wunused-but-set-variable" #include #include #include #include #include "hex-common.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" #include "hvx-utils.h" #include "hex-dma.h" #include "hex-profile.h" #include "htp-vtcm.h" #include "htp-tensor.h" struct htp_im2col_context { struct htp_ops_context * octx; uint32_t patch_base; // first patch index assigned to this dev 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 // 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 }; // Per-op VTCM layout for the patch-embed DMA path struct htp_im2col_vtcm_layout { size_t off_src; size_t off_dst; size_t src_bytes_per_thread; size_t dst_bytes_per_thread; size_t total_bytes; }; static inline void htp_im2col_vtcm_layout_build(struct htp_im2col_vtcm_layout * L, size_t src_row_bytes, size_t dst_row_bytes, uint32_t n_threads) { L->src_bytes_per_thread = 2 * src_row_bytes; L->dst_bytes_per_thread = 2 * dst_row_bytes; L->off_src = 0; L->off_dst = L->off_src + L->src_bytes_per_thread * n_threads; 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); \ } 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") #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; \ 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 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); \ if (row_start >= row_end) \ return; \ for (uint32_t r = row_start; r < row_end; r++) { \ const uint32_t in = r / OH; \ const uint32_t ioh = r % OH; \ for (uint32_t ikh = 0; ikh < KH; ikh++) { \ int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ int ok = (iih >= 0 && iih < (int32_t) IH); \ for (uint32_t iic = 0; iic < IC; iic++) { \ float * vdst = srcb + ((uint64_t) (iic * KH + ikh)) * IW; \ const float * _vsrc = \ ok ? (src_data + ((uint64_t) (in * IC + iic) * IH + iih) * IW) : (const float *) vdst; \ dma_queue_push_ddr_to_vtcm( \ dmaq, dma_make_ptr((uint8_t *) vdst, ok ? (const uint8_t *) _vsrc : (const uint8_t *) vdst), \ IW * sizeof(float), IW * sizeof(float), ok ? 1 : 0); \ } \ } \ for (uint32_t i = 0; i < IC * KH; i++) \ dma_queue_pop(dmaq); \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, r); \ for (uint32_t iow = 0; iow < OW; iow++) { \ DST_CTYPE * dst_patch = dstb + (uint64_t) iow * patch_stride; \ for (uint32_t ikh = 0; ikh < KH; ikh++) { \ int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh; \ for (uint32_t iic = 0; iic < IC; iic++) { \ DST_CTYPE * 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 float * src_run = srcb + ((uint64_t) (iic * KH + ikh)) * IW + (uint64_t) iow * KW; \ COPY_FN((uint8_t *) out_run, (const uint8_t *) src_run, KW); \ } \ } \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, r); \ DST_CTYPE * ddr_row = dst_data + ((uint64_t) (in * OH + ioh) * OW) * patch_stride; \ dma_queue_push_vtcm_to_ddr(dmaq, dma_make_ptr((uint8_t *) ddr_row, (uint8_t *) dstb), \ OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), 1); \ dma_queue_flush(dmaq); \ } \ } IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_thread, __fp16, hvx_copy_f16_f32_uu, hvx_splat_f16_u, sizeof(__fp16), "pe-dma-f16") IM2COL_PATCHEMBED_DMA_BODY(im2col_patchembed_dma_f32_thread, float, hvx_copy_f32_uu, hvx_splat_f32_u, sizeof(float), "pe-dma-f32") static bool im2col_use_patchembed_dma(const struct htp_ops_context * octx) { 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 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 } if (p0 != 0 || p1 != 0) { return false; // no padding } if (d0 != 1 || d1 != 1) { return false; // no dilation } return true; } // Sizes the per-thread 2x(src,dst) VTCM ping-pong for the patch-embed DMA path. // Returns false if it doesn't fit the VTCM budget (caller falls back). 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 uint32_t OW = octx->dst->ne[1]; const uint32_t patch_stride = IC * KH * KW; ictx->pe_src_row_bytes = hex_round_up(IC * KH * IW * sizeof(float), 256); const uint32_t dst_elem = (octx->dst->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); ictx->pe_dst_row_bytes = hex_round_up(OW * patch_stride * dst_elem, 256); // 2 src + 2 dst buffers per thread (ping-pong), src region first then dst. struct htp_im2col_vtcm_layout L; htp_im2col_vtcm_layout_build(&L, ictx->pe_src_row_bytes, ictx->pe_dst_row_bytes, n_threads); if (L.total_bytes > octx->ctx->vtcm_size) { return false; } uint8_t * const base = octx->ctx->vtcm_base; 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; } int op_im2col(struct htp_ops_context * octx) { const struct htp_tensor * src1 = octx->src[1]; const struct htp_tensor * dst = octx->dst; if (src1->type != HTP_TYPE_F32 || (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32)) { FARF(ERROR, "im2col: only (F32 image -> F16/F32 columns) supported"); return HTP_STATUS_NO_SUPPORT; } if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { return HTP_STATUS_OK; } const uint32_t N = src1->ne[3]; const uint32_t OH = dst->ne[2]; const uint32_t OW = dst->ne[1]; const uint32_t total_patches = N * OH * OW; const uint32_t total_rows = N * OH; uint32_t patch_base = 0; 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); patch_base = range.start; npatches = range.count; } uint32_t row_base = 0; 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); row_base = range.start; nrows = range.count; } if (npatches == 0 && nrows == 0) { return HTP_STATUS_OK; } const uint32_t n_threads = MIN(octx->n_threads, MAX(npatches, 1)); struct htp_im2col_context ictx = { 0 }; ictx.octx = octx; ictx.patch_base = patch_base; ictx.npatches = npatches; 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) { 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); } return HTP_STATUS_OK; } // else: doesn't fit -> fall through to the pure-DDR kernel below. } if (npatches == 0) { return HTP_STATUS_OK; } if (dst->type == HTP_TYPE_F16) { work_queue_run(octx->ctx->work_queue, im2col_patchembed_thread, &ictx, n_threads); } else { work_queue_run(octx->ctx->work_queue, im2col_patchembed_f32_thread, &ictx, n_threads); } return HTP_STATUS_OK; }