Commit b13a93a3d for llama.cpp
commit b13a93a3d5fa44b57899d86467a4de463e7415a7
Author: kurquhar <kurquhar@qti.qualcomm.com>
Date: Sun Oct 11 20:24:47 2026 -0700
hexagon: add IM2COL_3D (#30305)
* hexagon: add IM2COL_3D support
Assisted-by: OpenCode
* hexagon: add DMA paths for IM2COL_3D
Assisted-by: OpenCode
* hexagon: avoid scalar fp16 conversions in IM2COL
Assisted-by: OpenCode
* remove in-loop spills in pre-v81 devices
Assisted-by: OpenCode
* hex-im2col: re-write to use canonical dma dbl-buff pipeline, remove divs and very inefficient fp conversions
* hex-im2col: move kernel_params and vtcm layout precompute into the host
* hex-im2col: remove explicit tracking of number of dma-pushes
* hex-im2col: remove incorrect uint64_t casts
* hex-im2col: no need for macros for dma loops and pass compute func ptr via local context
* hex-im2col: use cursor pattern, tracing and proper n_threads handling
* hex-im2col: support non-contigue views
* hex-im2col: minor reformatting
---------
Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>
diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index f3ca73124..b214b1543 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -74,6 +74,7 @@
#include "htp/argsort-ops.h"
#include "htp/concat-ops.h"
#include "htp/cpy-ops.h"
+#include "htp/im2col-ops.h"
#include "htp_iface.h"
#include "htp-drv.h"
@@ -418,6 +419,18 @@ static bool ggml_hexagon_precompute_cpy_params(
struct htp_copy_kernel_params * kparams
);
+static bool ggml_hexagon_precompute_im2col_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_im2col_kernel_params * kparams
+);
+
+static bool ggml_hexagon_precompute_im2col_3d_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_im2col_kernel_params * kparams
+);
+
static void ggml_hexagon_precompute_fused_mmnx_params(
const struct ggml_hexagon_session * sess,
const struct ggml_tensor * src0,
@@ -7696,21 +7709,397 @@ static bool ggml_hexagon_supported_ssm_conv(const struct ggml_hexagon_session *
return true;
}
-static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+static bool im2col_3d_calc_output_size(uint32_t input, uint32_t kernel, int32_t stride, int32_t padding,
+ int32_t dilation, uint32_t * output) {
+ if (stride <= 0 || padding < 0 || dilation <= 0 || kernel == 0) {
+ return false;
+ }
+
+ const int64_t numerator = (int64_t) input + 2LL * padding - (int64_t) dilation * (kernel - 1) - 1;
+ if (numerator < 0) {
+ return false;
+ }
+
+ const uint64_t size = (uint64_t) (numerator / stride) + 1;
+ if (size == 0 || size > UINT32_MAX) {
+ return false;
+ }
+
+ *output = (uint32_t) size;
+ return true;
+}
+
+static bool ggml_hexagon_precompute_im2col_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_im2col_kernel_params * kparams
+) {
+ memset(kparams, 0, sizeof(*kparams));
+ kparams->kernel_type = HTP_IM2COL_KERNEL_UNSUPPORTED;
+
+ const struct ggml_tensor * src0 = op->src[0];
const struct ggml_tensor * src1 = op->src[1];
const struct ggml_tensor * dst = op;
- // For now support F32->F32 and F32->F16 only.
+ if (!src0 || !src1 || !dst) {
+ return false;
+ }
+
if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
return false;
}
- if (!ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
+ if (src1->nb[0] != sizeof(float) || !ggml_is_contiguous(dst)) {
return false;
}
- GGML_UNUSED(sess);
- return true;
+ const int32_t * op_params = (const int32_t *) op->op_params;
+ const int32_t s0 = op_params[0];
+ const int32_t s1 = op_params[1];
+ const int32_t p0 = op_params[2];
+ const int32_t p1 = op_params[3];
+ const int32_t d0 = op_params[4];
+ const int32_t d1 = op_params[5];
+ const bool is_2D = (op_params[6] == 1);
+
+ if (s0 <= 0 || d0 <= 0 || p0 < 0) {
+ return false;
+ }
+ if (is_2D && (s1 <= 0 || d1 <= 0 || p1 < 0)) {
+ return false;
+ }
+
+ const uint32_t N = is_2D ? (uint32_t) src1->ne[3] : (uint32_t) src1->ne[2];
+ const uint32_t IC = is_2D ? (uint32_t) src1->ne[2] : (uint32_t) src1->ne[1];
+ const uint32_t IH = is_2D ? (uint32_t) src1->ne[1] : 1;
+ const uint32_t IW = (uint32_t) src1->ne[0];
+ const uint32_t KH = is_2D ? (uint32_t) src0->ne[1] : 1;
+ const uint32_t KW = (uint32_t) src0->ne[0];
+ const uint32_t OH = is_2D ? (uint32_t) dst->ne[2] : 1;
+ const uint32_t OW = (uint32_t) dst->ne[1];
+
+ if (KW == 0 || KH == 0 || IW == 0 || IH == 0 || IC == 0 || OW == 0 || OH == 0 || N == 0) {
+ return false;
+ }
+
+ const uint64_t patch_stride_64 = (uint64_t) IC * KH * KW;
+ const uint64_t total_rows_64 = (uint64_t) N * OH;
+ if (patch_stride_64 > UINT32_MAX || total_rows_64 > UINT32_MAX) {
+ return false;
+ }
+ const uint32_t patch_stride = (uint32_t) patch_stride_64;
+ const uint32_t total_rows = (uint32_t) total_rows_64;
+
+ if (dst->ne[0] != patch_stride || dst->ne[1] != OW) {
+ return false;
+ }
+ if (is_2D && dst->ne[2] != OH) {
+ return false;
+ }
+
+ const uint32_t dst_elem = (dst->type == GGML_TYPE_F16) ? sizeof(ggml_fp16_t) : sizeof(float);
+
+ bool is_patchembed = ggml_is_contiguous(src1) && (s0 == (int32_t) KW && p0 == 0 && d0 == 1);
+ if (is_2D) {
+ is_patchembed = is_patchembed && (s1 == (int32_t) KH && p1 == 0 && d1 == 1);
+ }
+
+ const uint32_t max_threads = (sess && sess->n_threads > 0) ? (uint32_t) sess->n_threads : 8;
+ const uint32_t n_threads = (total_rows < max_threads) ? total_rows : max_threads;
+ const uint32_t vtcm_size = (sess && sess->vtcm_size > 0) ? sess->vtcm_size : 0;
+
+ if (is_patchembed) {
+ const uint64_t src_row_bytes_raw = (uint64_t) IC * KH * IW * sizeof(float);
+ const uint64_t dst_row_bytes_raw = (uint64_t) OW * patch_stride * dst_elem;
+ if (src_row_bytes_raw <= UINT32_MAX && dst_row_bytes_raw <= UINT32_MAX) {
+ const uint32_t src_row_bytes = hex_round_up((uint32_t) src_row_bytes_raw, 256);
+ const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);
+
+ struct htp_im2col_vtcm_layout layout;
+ htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+
+ if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
+ kparams->kernel_type = HTP_IM2COL_KERNEL_PATCHEMBED_DMA;
+ kparams->is_2d = is_2D ? 1 : 0;
+ kparams->is_3d = 0;
+ kparams->n_threads = (uint8_t) n_threads;
+ kparams->vtcm_size = layout.total_bytes;
+ kparams->off_src = layout.off_src;
+ kparams->off_dst = layout.off_dst;
+ kparams->src_bytes_per_thread = layout.src_bytes_per_thread;
+ kparams->dst_bytes_per_thread = layout.dst_bytes_per_thread;
+ kparams->src_row_bytes = src_row_bytes;
+ kparams->dst_row_bytes = dst_row_bytes;
+ kparams->owb = OW;
+ kparams->wb = IW;
+ kparams->nbpr = 1;
+ kparams->N = N;
+ kparams->IC = IC;
+ kparams->ID = 1;
+ kparams->IH = IH;
+ kparams->IW = IW;
+ kparams->KD = 1;
+ kparams->KH = KH;
+ kparams->KW = KW;
+ kparams->OD = 1;
+ kparams->OH = OH;
+ kparams->OW = OW;
+ kparams->patch_stride = patch_stride;
+ kparams->total_rows = total_rows;
+ return true;
+ }
+ }
+ }
+
+ for (uint32_t owb = (OW < 256 ? OW : 256); owb >= 1; owb--) {
+ const uint64_t Wb = (uint64_t) (owb - 1) * s0 + (uint64_t) (KW - 1) * d0 + 1;
+ if (Wb > UINT32_MAX) {
+ if (owb == 1) break;
+ continue;
+ }
+
+ const uint64_t src_row_bytes_raw = (uint64_t) IC * KH * Wb * sizeof(float);
+ const uint64_t dst_row_bytes_raw = (uint64_t) owb * patch_stride * dst_elem;
+ if (src_row_bytes_raw > UINT32_MAX || dst_row_bytes_raw > UINT32_MAX) {
+ if (owb == 1) break;
+ continue;
+ }
+
+ const uint32_t src_row_bytes = hex_round_up((uint32_t) src_row_bytes_raw, 256);
+ const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);
+
+ struct htp_im2col_vtcm_layout layout;
+ htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+
+ if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
+ kparams->kernel_type = HTP_IM2COL_KERNEL_BLOCKED_DMA;
+ kparams->is_2d = is_2D ? 1 : 0;
+ kparams->is_3d = 0;
+ kparams->n_threads = (uint8_t) n_threads;
+ kparams->vtcm_size = layout.total_bytes;
+ kparams->off_src = layout.off_src;
+ kparams->off_dst = layout.off_dst;
+ kparams->src_bytes_per_thread = layout.src_bytes_per_thread;
+ kparams->dst_bytes_per_thread = layout.dst_bytes_per_thread;
+ kparams->src_row_bytes = src_row_bytes;
+ kparams->dst_row_bytes = dst_row_bytes;
+ kparams->owb = owb;
+ kparams->wb = (uint32_t) Wb;
+ kparams->nbpr = (OW + owb - 1) / owb;
+ kparams->N = N;
+ kparams->IC = IC;
+ kparams->ID = 1;
+ kparams->IH = IH;
+ kparams->IW = IW;
+ kparams->KD = 1;
+ kparams->KH = KH;
+ kparams->KW = KW;
+ kparams->OD = 1;
+ kparams->OH = OH;
+ kparams->OW = OW;
+ kparams->patch_stride = patch_stride;
+ kparams->total_rows = total_rows;
+ return true;
+ }
+ if (owb == 1) break;
+ }
+
+ return false;
+}
+
+static bool ggml_hexagon_precompute_im2col_3d_params(
+ const struct ggml_hexagon_session * sess,
+ const struct ggml_tensor * op,
+ struct htp_im2col_kernel_params * kparams
+) {
+ memset(kparams, 0, sizeof(*kparams));
+ kparams->kernel_type = HTP_IM2COL_KERNEL_UNSUPPORTED;
+
+ const struct ggml_tensor * src0 = op->src[0];
+ const struct ggml_tensor * src1 = op->src[1];
+ const struct ggml_tensor * dst = op;
+
+ if (!src0 || !src1 || !dst) {
+ return false;
+ }
+
+ if (src1->type != GGML_TYPE_F32 || (dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_F32)) {
+ return false;
+ }
+
+ if (src1->nb[0] != sizeof(float) || !ggml_is_contiguous(dst)) {
+ return false;
+ }
+
+ const int32_t * params = (const int32_t *) op->op_params;
+ const int32_t s0 = params[0], s1 = params[1], s2 = params[2];
+ const int32_t p0 = params[3], p1 = params[4], p2 = params[5];
+ const int32_t d0 = params[6], d1 = params[7], d2 = params[8];
+ const int32_t ic = params[9];
+
+ if (s0 <= 0 || s1 <= 0 || s2 <= 0 ||
+ p0 < 0 || p1 < 0 || p2 < 0 ||
+ d0 <= 0 || d1 <= 0 || d2 <= 0 || ic <= 0) {
+ return false;
+ }
+
+ if (src1->ne[3] % ic != 0) {
+ return false;
+ }
+
+ const uint32_t IC = (uint32_t) ic;
+ const uint32_t N = (uint32_t) (src1->ne[3] / IC);
+ if (N == 0) {
+ return false;
+ }
+
+ uint32_t OW, OH, OD;
+ if (!im2col_3d_calc_output_size(src1->ne[0], src0->ne[0], s0, p0, d0, &OW) ||
+ !im2col_3d_calc_output_size(src1->ne[1], src0->ne[1], s1, p1, d1, &OH) ||
+ !im2col_3d_calc_output_size(src1->ne[2], src0->ne[2], s2, p2, d2, &OD)) {
+ return false;
+ }
+
+ const uint32_t KD = (uint32_t) src0->ne[2];
+ const uint32_t KH = (uint32_t) src0->ne[1];
+ const uint32_t KW = (uint32_t) src0->ne[0];
+ const uint32_t ID = (uint32_t) src1->ne[2];
+ const uint32_t IH = (uint32_t) src1->ne[1];
+ const uint32_t IW = (uint32_t) src1->ne[0];
+
+ const uint64_t patch_stride_64 = (uint64_t) IC * KD * KH * KW;
+ const uint64_t total_rows_64 = (uint64_t) N * OD * OH;
+ const uint64_t dst_ne3_64 = (uint64_t) OD * N;
+
+ if (patch_stride_64 > UINT32_MAX || total_rows_64 > UINT32_MAX || dst_ne3_64 > UINT32_MAX) {
+ return false;
+ }
+ const uint32_t patch_stride = (uint32_t) patch_stride_64;
+ const uint32_t total_rows = (uint32_t) total_rows_64;
+ const uint32_t dst_ne3 = (uint32_t) dst_ne3_64;
+
+ if (dst->ne[0] != patch_stride || dst->ne[1] != OW || dst->ne[2] != OH || dst->ne[3] != dst_ne3) {
+ return false;
+ }
+
+ const uint32_t dst_elem = (dst->type == GGML_TYPE_F16) ? sizeof(ggml_fp16_t) : sizeof(float);
+
+ const bool is_patchembed = ggml_is_contiguous(src1) &&
+ (s0 == (int32_t) KW && s1 == (int32_t) KH && s2 == (int32_t) KD &&
+ p0 == 0 && p1 == 0 && p2 == 0 &&
+ d0 == 1 && d1 == 1 && d2 == 1);
+
+ const uint32_t max_threads = (sess && sess->n_threads > 0) ? (uint32_t) sess->n_threads : 8;
+ const uint32_t n_threads = (total_rows < max_threads) ? total_rows : max_threads;
+ const uint32_t vtcm_size = (sess && sess->vtcm_size > 0) ? sess->vtcm_size : 0;
+
+ if (is_patchembed) {
+ const uint64_t src_row_bytes_raw = (uint64_t) IC * KD * KH * IW * sizeof(float);
+ const uint64_t dst_row_bytes_raw = (uint64_t) OW * patch_stride * dst_elem;
+ if (src_row_bytes_raw <= UINT32_MAX && dst_row_bytes_raw <= UINT32_MAX) {
+ const uint32_t src_row_bytes = hex_round_up((uint32_t) src_row_bytes_raw, 256);
+ const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);
+
+ struct htp_im2col_vtcm_layout layout;
+ htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+
+ if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
+ kparams->kernel_type = HTP_IM2COL_KERNEL_PATCHEMBED_DMA;
+ kparams->is_2d = 0;
+ kparams->is_3d = 1;
+ kparams->n_threads = (uint8_t) n_threads;
+ kparams->vtcm_size = layout.total_bytes;
+ kparams->off_src = layout.off_src;
+ kparams->off_dst = layout.off_dst;
+ kparams->src_bytes_per_thread = layout.src_bytes_per_thread;
+ kparams->dst_bytes_per_thread = layout.dst_bytes_per_thread;
+ kparams->src_row_bytes = src_row_bytes;
+ kparams->dst_row_bytes = dst_row_bytes;
+ kparams->owb = OW;
+ kparams->wb = IW;
+ kparams->nbpr = 1;
+ kparams->N = N;
+ kparams->IC = IC;
+ kparams->ID = ID;
+ kparams->IH = IH;
+ kparams->IW = IW;
+ kparams->KD = KD;
+ kparams->KH = KH;
+ kparams->KW = KW;
+ kparams->OD = OD;
+ kparams->OH = OH;
+ kparams->OW = OW;
+ kparams->patch_stride = patch_stride;
+ kparams->total_rows = total_rows;
+ return true;
+ }
+ }
+ }
+
+ for (uint32_t owb = (OW < 256 ? OW : 256); owb >= 1; owb--) {
+ const uint64_t Wb = (uint64_t) (owb - 1) * s0 + (uint64_t) (KW - 1) * d0 + 1;
+ if (Wb > UINT32_MAX) {
+ if (owb == 1) break;
+ continue;
+ }
+
+ const uint64_t src_row_bytes_raw = (uint64_t) IC * KD * KH * Wb * sizeof(float);
+ const uint64_t dst_row_bytes_raw = (uint64_t) owb * patch_stride * dst_elem;
+ if (src_row_bytes_raw > UINT32_MAX || dst_row_bytes_raw > UINT32_MAX) {
+ if (owb == 1) break;
+ continue;
+ }
+
+ const uint32_t src_row_bytes = hex_round_up((uint32_t) src_row_bytes_raw, 256);
+ const uint32_t dst_row_bytes = hex_round_up((uint32_t) dst_row_bytes_raw, 256);
+
+ struct htp_im2col_vtcm_layout layout;
+ htp_im2col_vtcm_layout_build(&layout, src_row_bytes, dst_row_bytes, n_threads);
+
+ if (vtcm_size == 0 || layout.total_bytes <= vtcm_size) {
+ kparams->kernel_type = HTP_IM2COL_KERNEL_BLOCKED_DMA;
+ kparams->is_2d = 0;
+ kparams->is_3d = 1;
+ kparams->n_threads = (uint8_t) n_threads;
+ kparams->vtcm_size = layout.total_bytes;
+ kparams->off_src = layout.off_src;
+ kparams->off_dst = layout.off_dst;
+ kparams->src_bytes_per_thread = layout.src_bytes_per_thread;
+ kparams->dst_bytes_per_thread = layout.dst_bytes_per_thread;
+ kparams->src_row_bytes = src_row_bytes;
+ kparams->dst_row_bytes = dst_row_bytes;
+ kparams->owb = owb;
+ kparams->wb = (uint32_t) Wb;
+ kparams->nbpr = (OW + owb - 1) / owb;
+ kparams->N = N;
+ kparams->IC = IC;
+ kparams->ID = ID;
+ kparams->IH = IH;
+ kparams->IW = IW;
+ kparams->KD = KD;
+ kparams->KH = KH;
+ kparams->KW = KW;
+ kparams->OD = OD;
+ kparams->OH = OH;
+ kparams->OW = OW;
+ kparams->patch_stride = patch_stride;
+ kparams->total_rows = total_rows;
+ return true;
+ }
+ if (owb == 1) break;
+ }
+
+ return false;
+}
+
+static bool ggml_hexagon_supported_im2col(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+ struct htp_im2col_kernel_params kparams;
+ return ggml_hexagon_precompute_im2col_params(sess, op, &kparams);
+}
+
+static bool ggml_hexagon_supported_im2col_3d(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+ struct htp_im2col_kernel_params kparams;
+ return ggml_hexagon_precompute_im2col_3d_params(sess, op, &kparams);
}
static bool ggml_hexagon_supported_pad(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -7875,6 +8264,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
case GGML_OP_TRI: return HTP_OP_TRI;
case GGML_OP_PAD: return HTP_OP_PAD;
case GGML_OP_IM2COL: return HTP_OP_IM2COL;
+ case GGML_OP_IM2COL_3D: return HTP_OP_IM2COL_3D;
case GGML_OP_ROLL: return HTP_OP_ROLL;
case GGML_OP_UNARY:
@@ -8157,6 +8547,16 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
if (ok && node.opcode == HTP_OP_CPY && kparams->total_elems == 0) {
continue;
}
+ } else if (node.opcode == HTP_OP_IM2COL) {
+ ggml_hexagon_precompute_im2col_params(sess,
+ node.node,
+ (struct htp_im2col_kernel_params *) node.kernel_params
+ );
+ } else if (node.opcode == HTP_OP_IM2COL_3D) {
+ ggml_hexagon_precompute_im2col_3d_params(sess,
+ node.node,
+ (struct htp_im2col_kernel_params *) node.kernel_params
+ );
}
computed_nodes.push_back(std::move(node));
}
@@ -9018,6 +9418,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
supp = ggml_hexagon_supported_im2col(sess, op);
break;
+ case GGML_OP_IM2COL_3D:
+ supp = ggml_hexagon_supported_im2col_3d(sess, op);
+ break;
+
case GGML_OP_GATED_DELTA_NET:
supp = ggml_hexagon_supported_gated_delta_net(sess, op);
break;
diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h
index 87a1f4c7b..b36322ae8 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ctx.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h
@@ -173,6 +173,7 @@ int op_solve_tri(struct htp_ops_context * octx);
int op_gated_delta_net(struct htp_ops_context * octx);
int op_pad(struct htp_ops_context * octx);
int op_im2col(struct htp_ops_context * octx);
+int op_im2col_3d(struct htp_ops_context * octx);
int op_pool_2d(struct htp_ops_context * octx);
int op_pool_1d(struct htp_ops_context * octx);
int op_allreduce(struct htp_ops_context * octx);
diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h
index 5d2571fef..5405c83b6 100644
--- a/ggml/src/ggml-hexagon/htp/htp-ops.h
+++ b/ggml/src/ggml-hexagon/htp/htp-ops.h
@@ -118,6 +118,7 @@ enum htp_op_code {
HTP_OP_POOL_2D,
HTP_OP_POOL_1D,
HTP_OP_UNARY_GELU_QUICK,
+ HTP_OP_IM2COL_3D,
HTP_OP_INVALID
};
diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.c b/ggml/src/ggml-hexagon/htp/im2col-ops.c
index b99d711c9..ee3c7da11 100644
--- a/ggml/src/ggml-hexagon/htp/im2col-ops.c
+++ b/ggml/src/ggml-hexagon/htp/im2col-ops.c
@@ -18,17 +18,15 @@
#include "hex-profile.h"
#include "htp-vtcm.h"
#include "htp-tensor.h"
+#include "im2col-ops.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_row_base; // first output row index assigned to this dev (DMA path)
+ uint32_t pe_nrows; // number of output rows assigned to this dev (DMA path)
+ uint32_t pe_rows_per_thread; // output rows per worker
+ uint32_t pe_src_row_bytes; // one output row's source, 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.
@@ -39,551 +37,892 @@ struct htp_im2col_context {
uint32_t pe_owb; // output-col block size
uint32_t pe_wb; // staged source window width
+ uint32_t pe_nbpr; // blocks per row
+
+ // Precomputed dimensions for worker threads
+ uint32_t N;
+ uint32_t IC;
+ uint32_t ID, IH, IW;
+ uint32_t KD, KH, KW;
+ uint32_t OD, OH, OW;
+ uint32_t patch_stride;
+
+ // Per-thread initial coordinates (removes division from thread prologue)
+ uint32_t thread_in[HTP_MAX_NTHREADS];
+ uint32_t thread_iod[HTP_MAX_NTHREADS];
+ uint32_t thread_ioh[HTP_MAX_NTHREADS];
+
+ void * compute;
+ uint32_t dst_elem;
};
-// 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;
+typedef void (*im2col_blocked_compute_fn_t)(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb, const float * restrict srcb,
+ float * restrict scratch,
+ uint32_t nb, uint32_t c0, uint32_t ioh);
+
+typedef void (*im2col_patchembed_compute_fn_t)(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb, const float * restrict srcb, uint32_t ioh);
+
+typedef void (*im2col_3d_blocked_compute_fn_t)(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb, const float * restrict srcb,
+ float * restrict scratch,
+ uint32_t nb, uint32_t c0, uint32_t iod, uint32_t ioh);
+
+typedef void (*im2col_3d_patchembed_compute_fn_t)(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb, const float * restrict srcb);
+
+struct im2col_cursor {
+ uint32_t r;
+ uint32_t in;
+ uint32_t iod;
+ uint32_t ioh;
+ uint32_t c0;
};
-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;
+static inline void im2col_step_row(struct im2col_cursor * c, const struct htp_im2col_context * ictx) {
+ c->r++;
+ c->ioh++;
+ if (c->ioh >= ictx->OH) {
+ c->ioh = 0;
+ c->in++;
+ }
+}
- 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;
+static inline void im2col_step_block(struct im2col_cursor * c, const struct htp_im2col_context * ictx) {
+ c->c0 += ictx->pe_owb;
+ if (c->c0 >= ictx->OW) {
+ c->c0 = 0;
+ c->r++;
+ c->ioh++;
+ if (c->ioh >= ictx->OH) {
+ c->ioh = 0;
+ c->in++;
+ }
+ }
}
-#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 *) (uintptr_t) src1->data; \
- DST_CTYPE * restrict dst_data = (DST_CTYPE *) (uintptr_t) 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 dma_addr_t src_data = src1->data; \
- const dma_addr_t dst_data = 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 + (size_t) ikh * Wb + lo; \
- const dma_addr_t vsrc = src_data + (size_t) (((in * IC) * IH + iih) * IW + (win0 + lo)) * sizeof(float); \
- while (!dma_queue_push(dmaq, dma_make_data(vdst, 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 + (size_t) ikh * Wb + lo; \
- const dma_addr_t vsrc = src_data + (size_t) (((nin * IC) * IH + iih) * IW + (nwin0 + lo)) * sizeof(float); \
- while (!dma_queue_push(dmaq, dma_make_data(vdst, 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); \
- const dma_addr_t ddr = dst_data + (size_t) ((in * OH + ioh) * OW + c0) * patch_stride * (DST_ELEM); \
- dma_queue_push(dmaq, dma_make_data(ddr, dstb), \
- nb * patch_stride * (DST_ELEM), 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; \
- 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 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 dma_addr_t src_data = src1->data; \
- const dma_addr_t dst_data = dst->data; \
- dma_queue * dma_q = 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 + (size_t) (iic * KH + ikh) * IW; \
- const dma_addr_t vsrc = ok \
- ? (src_data + (size_t) ((in * IC + iic) * IH + iih) * IW * sizeof(float)) \
- : src_data; \
- /* IC*KH descriptors per row can exceed the ring capacity: retire the oldest when full */ \
- while (!dma_queue_push(dma_q, dma_make_data(vdst, vsrc), \
- IW * sizeof(float), IW * sizeof(float), IW * sizeof(float), \
- ok ? 1 : 0)) { \
- dma_queue_pop(dma_q); \
- } \
- } \
- } \
- dma_queue_flush(dma_q); \
- 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); \
- const dma_addr_t ddr_row = dst_data + (size_t) (in * OH + ioh) * OW * patch_stride * (DST_ELEM); \
- dma_queue_push(dma_q, dma_make_data(ddr_row, dstb), \
- OW * patch_stride * (DST_ELEM), OW * patch_stride * (DST_ELEM), \
- OW * patch_stride * (DST_ELEM), 1); \
- dma_queue_flush(dma_q); \
- } \
- }
-
-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) {
+static void im2col_blocked_dma_stage(struct htp_im2col_context * ictx, dma_queue * dma_q,
+ float * srcb, const struct im2col_cursor * c) {
+ const uint32_t in = c->in, ioh = c->ioh, c0 = c->c0;
+ const struct htp_ops_context * octx = ictx->octx;
+ const struct htp_tensor * src1 = octx->src[1];
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 (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
- return false;
+ const int32_t d1 = octx->op_params[5];
+ const bool is_2D = (octx->op_params[6] == 1);
+ const uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KH = ictx->KH;
+ const uint32_t Wb = ictx->pe_wb;
+ const dma_addr_t src_data = src1->data;
+ const int32_t win0 = (int32_t) c0 * s0 - p0;
+ const size_t src_channel_stride = is_2D ? src1->nb[2] : src1->nb[1];
+ const size_t batch_offset = is_2D ? (size_t) in * src1->nb[3] : (size_t) in * src1->nb[2];
+
+ 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) {
+ dma_queue_push(dma_q, dma_make_data(srcb, 0), 0, 0, 0, 0);
+ 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) {
+ dma_queue_push(dma_q, dma_make_data(srcb, 0), 0, 0, 0, 0);
+ continue;
+ }
+ const uint32_t cpw = (uint32_t) (hi - lo);
+ float * vdst = srcb + (size_t) ikh * Wb + lo;
+ const dma_addr_t vsrc = src_data + batch_offset + (size_t) iih * src1->nb[1] + (size_t) (win0 + lo) * src1->nb[0];
+ dma_queue_push(dma_q, dma_make_data(vdst, vsrc),
+ (size_t) KH * Wb * sizeof(float), src_channel_stride,
+ cpw * sizeof(float), IC);
}
- 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)
+}
+
+#define IM2COL_BLOCKED_COMPUTE_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN) \
+static void FNAME( \
+ const struct htp_im2col_context * ictx, \
+ void * restrict dstb_raw, const float * restrict srcb, \
+ float * restrict scratch, \
+ uint32_t nb, uint32_t c0, uint32_t ioh) { \
+ DST_CTYPE * restrict dstb = (DST_CTYPE *) dstb_raw; \
+ const struct htp_ops_context * octx = ictx->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 uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW; \
+ const uint32_t KH = ictx->KH, KW = ictx->KW; \
+ const uint32_t Wb = ictx->pe_wb; \
+ const uint32_t patch_stride = ictx->patch_stride; \
+ const int32_t win0 = (int32_t) c0 * s0 - p0; \
+ \
+ for (uint32_t j = 0; j < nb; j++) { \
+ const uint32_t iow = c0 + j; \
+ DST_CTYPE * dst_patch = dstb + (size_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 + (size_t) (iic * KH + ikh) * KW; \
+ if (!okh) { \
+ SPLAT_FN(out_run, 0, KW); \
+ continue; \
+ } \
+ const float * vrow = srcb + (size_t) (iic * KH + ikh) * Wb; \
+ if (d0 == 1) { \
+ 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, KW); \
+ } else { \
+ if (lo > 0) { \
+ SPLAT_FN(out_run, 0, (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, (uint32_t) (KW - hi)); \
+ } \
+ } \
+ continue; \
+ } \
+ for (uint32_t ikw = 0; ikw < KW; ikw += 64) { \
+ const uint32_t chunk = MIN(KW - ikw, 64); \
+ for (uint32_t c = 0; c < chunk; c++) { \
+ const int32_t iiw = iiw0 + (int32_t) (ikw + c) * d0; \
+ scratch[c] = (iiw < 0 || iiw >= (int32_t) IW) ? 0.0f : vrow[iiw - win0]; \
+ } \
+ COPY_FN((uint8_t *) (out_run + ikw), (const uint8_t *) scratch, chunk); \
+ } \
+ } \
+ } \
+ } \
+}
+
+IM2COL_BLOCKED_COMPUTE_BODY(im2col_blocked_compute_f16, uint16_t, hvx_copy_f16_f32_uu, hvx_splat_u16_u)
+IM2COL_BLOCKED_COMPUTE_BODY(im2col_blocked_compute_f32, float, hvx_copy_f32_uu, hvx_splat_f32_u)
+
+static void im2col_blocked_dma_thread(unsigned int nth, unsigned int ith, void * data) {
+ (void) nth;
+ 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 uint32_t OW = ictx->OW;
+ const uint32_t KH = ictx->KH;
+ const uint32_t owb = ictx->pe_owb;
+ const uint32_t nbpr = ictx->pe_nbpr;
+ const uint32_t patch_stride = ictx->patch_stride;
+ const uint32_t dst_elem = ictx->dst_elem;
+ im2col_blocked_compute_fn_t compute = (im2col_blocked_compute_fn_t) ictx->compute;
+ const dma_addr_t dst_data = octx->dst->data;
+ dma_queue * dma_q = 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) };
+ void * dstb2[2] = { dstb_base, dstb_base + ictx->pe_dst_row_bytes };
+ float * scratch = (float *) (srcb_base + 2 * ictx->pe_src_row_bytes);
+ const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows;
+ const uint32_t row_start = ictx->pe_row_base + ictx->pe_rows_per_thread * ith;
+ const uint32_t row_end = MIN(row_start + ictx->pe_rows_per_thread, row_end_max);
+ if (row_start >= row_end) {
+ return;
}
- if (p0 != 0) {
- return false; // no padding (width)
+ const uint32_t nrows_local = row_end - row_start;
+ const uint32_t total_blocks = nrows_local * nbpr;
+ if (total_blocks == 0) {
+ return;
}
- if (d0 != 1) {
- return false; // no dilation (width)
+ struct im2col_cursor cur = {
+ .r = row_start,
+ .in = ictx->thread_in[ith],
+ .iod = 0,
+ .ioh = ictx->thread_ioh[ith],
+ .c0 = 0,
+ };
+ struct im2col_cursor stage = cur;
+ for (uint32_t i = 0; i < total_blocks && i < 2; i++) {
+ dma_queue_push(dma_q, dma_make_data(dst_data, dstb2[i]), 0, 0, 0, 0);
+ im2col_blocked_dma_stage(ictx, dma_q, srcb2[i], &stage);
+ im2col_step_block(&stage, ictx);
}
- if (is_2D) {
- if (s1 != (int32_t) KH) {
- return false; // non-overlapping (height)
+ for (uint32_t bi = 0; bi < total_blocks; bi++) {
+ const uint32_t buf = bi & 1u;
+ float * srcb = srcb2[buf];
+ void * dstb = dstb2[buf];
+ const uint32_t nb = MIN(owb, OW - cur.c0);
+ dma_queue_pop(dma_q);
+ for (uint32_t k = 0; k < KH; k++) {
+ dma_queue_pop(dma_q);
}
- if (p1 != 0) {
- return false; // no padding (height)
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ compute(ictx, dstb, srcb, scratch, nb, cur.c0, cur.ioh);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ const dma_addr_t ddr = dst_data + (size_t) (cur.r * OW + cur.c0) * patch_stride * dst_elem;
+ dma_queue_push(dma_q, dma_make_data(ddr, dstb),
+ nb * patch_stride * dst_elem, nb * patch_stride * dst_elem,
+ nb * patch_stride * dst_elem, 1);
+ if (bi + 2 < total_blocks) {
+ im2col_blocked_dma_stage(ictx, dma_q, srcb, &stage);
+ im2col_step_block(&stage, ictx);
}
- if (d1 != 1) {
- return false; // no dilation (height)
+ im2col_step_block(&cur, ictx);
+ }
+ dma_queue_flush(dma_q);
+}
+
+static void im2col_patchembed_dma_stage(struct htp_im2col_context * ictx, dma_queue * dma_q,
+ float * srcb, const struct im2col_cursor * c) {
+ const struct htp_ops_context * octx = ictx->octx;
+ const uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KH = ictx->KH;
+ const uint32_t in = c->in, ioh = c->ioh;
+ const size_t src_channel_stride = (size_t) IH * IW * sizeof(float);
+ const size_t vtcm_channel_stride = (size_t) KH * IW * sizeof(float);
+ const dma_addr_t src_data = octx->src[1]->data;
+ const dma_addr_t vsrc = src_data + (size_t) (((in * IC) * IH + ioh * KH) * IW) * sizeof(float);
+
+ dma_queue_push(dma_q, dma_make_data(srcb, vsrc),
+ vtcm_channel_stride, src_channel_stride,
+ vtcm_channel_stride, IC);
+}
+
+static void im2col_patchembed_compute_f16(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb_raw, const float * restrict srcb, uint32_t ioh) {
+ uint16_t * restrict dstb = (uint16_t *) dstb_raw;
+ const uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KH = ictx->KH, KW = ictx->KW;
+ const uint32_t OW = ictx->OW, patch_stride = ictx->patch_stride;
+
+ for (uint32_t iow = 0; iow < OW; iow++) {
+ uint16_t * dst_patch = dstb + (size_t) iow * patch_stride;
+ for (uint32_t ikh = 0; ikh < KH; ikh++) {
+ const int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;
+ for (uint32_t iic = 0; iic < IC; iic++) {
+ uint16_t * out_run = dst_patch + (size_t) (iic * KH + ikh) * KW;
+ if (iih < 0 || iih >= (int32_t) IH) {
+ hvx_splat_u16_u(out_run, 0, KW);
+ continue;
+ }
+ const float * src_run = srcb + (size_t) (iic * KH + ikh) * IW + (size_t) iow * KW;
+ hvx_copy_f16_f32_uu((uint8_t *) out_run, (const uint8_t *) src_run, KW);
+ }
}
}
- 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 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;
-
- 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;
+static void im2col_patchembed_compute_f32(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb_raw, const float * restrict srcb, uint32_t ioh) {
+ float * restrict dstb = (float *) dstb_raw;
+ const uint32_t IC = ictx->IC, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KH = ictx->KH, KW = ictx->KW;
+ const uint32_t OW = ictx->OW, patch_stride = ictx->patch_stride;
+
+ for (uint32_t iow = 0; iow < OW; iow++) {
+ float * dst_patch = dstb + (size_t) iow * patch_stride;
+ for (uint32_t ikh = 0; ikh < KH; ikh++) {
+ const int32_t iih = (int32_t) ioh * (int32_t) KH + (int32_t) ikh;
+ for (uint32_t iic = 0; iic < IC; iic++) {
+ float * out_run = dst_patch + (size_t) (iic * KH + ikh) * KW;
+ if (iih < 0 || iih >= (int32_t) IH) {
+ hvx_splat_f32_u(out_run, 0, KW);
+ continue;
+ }
+ const float * src_run = srcb + (size_t) (iic * KH + ikh) * IW + (size_t) iow * KW;
+ hvx_copy_f32_uu((uint8_t *) out_run, (const uint8_t *) src_run, KW);
+ }
+ }
+ }
}
-// 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;
+static void im2col_patchembed_dma_thread(unsigned int nth, unsigned int ith, void * data) {
+ (void) nth;
+ 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 uint32_t OW = ictx->OW;
+ const uint32_t patch_stride = ictx->patch_stride;
+ const uint32_t dst_elem = ictx->dst_elem;
+ im2col_patchembed_compute_fn_t compute = (im2col_patchembed_compute_fn_t) ictx->compute;
+ const dma_addr_t dst_data = octx->dst->data;
+ dma_queue * dma_q = 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) };
+ void * dstb2[2] = { dstb_base, dstb_base + ictx->pe_dst_row_bytes };
+ const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows;
+ const uint32_t row_start = ictx->pe_row_base + ictx->pe_rows_per_thread * ith;
+ const uint32_t row_end = MIN(row_start + ictx->pe_rows_per_thread, row_end_max);
+ if (row_start >= row_end) {
+ return;
+ }
+ const uint32_t total_rows = row_end - row_start;
+ if (total_rows == 0) {
+ return;
+ }
+ struct im2col_cursor cur = {
+ .r = row_start,
+ .in = ictx->thread_in[ith],
+ .iod = 0,
+ .ioh = ictx->thread_ioh[ith],
+ .c0 = 0,
+ };
+ struct im2col_cursor stage = cur;
+ for (uint32_t i = 0; i < total_rows && i < 2; i++) {
+ dma_queue_push(dma_q, dma_make_data(dst_data, dstb2[i]), 0, 0, 0, 0);
+ im2col_patchembed_dma_stage(ictx, dma_q, srcb2[i], &stage);
+ im2col_step_row(&stage, ictx);
+ }
+ for (uint32_t ri = 0; ri < total_rows; ri++) {
+ const uint32_t buf = ri & 1u;
+ float * srcb = srcb2[buf];
+ void * dstb = dstb2[buf];
+ dma_queue_pop(dma_q);
+ dma_queue_pop(dma_q);
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ compute(ictx, dstb, srcb, cur.ioh);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ const dma_addr_t ddr_row = dst_data + (size_t) cur.r * OW * patch_stride * dst_elem;
+ dma_queue_push(dma_q, dma_make_data(ddr_row, dstb),
+ OW * patch_stride * dst_elem, OW * patch_stride * dst_elem,
+ OW * patch_stride * dst_elem, 1);
+ if (ri + 2 < total_rows) {
+ im2col_patchembed_dma_stage(ictx, dma_q, srcb, &stage);
+ im2col_step_row(&stage, ictx);
}
- if (owb == 1) break; // avoid unsigned underflow
+ im2col_step_row(&cur, ictx);
}
- return false;
+ dma_queue_flush(dma_q);
}
-int op_im2col(struct htp_ops_context * octx) {
+static inline void im2col_3d_step_row(struct im2col_cursor * c, const struct htp_im2col_context * ictx) {
+ c->r++;
+ c->ioh++;
+ if (c->ioh >= ictx->OH) {
+ c->ioh = 0;
+ c->iod++;
+ if (c->iod >= ictx->OD) {
+ c->iod = 0;
+ c->in++;
+ }
+ }
+}
+
+static inline void im2col_3d_step_block(struct im2col_cursor * c, const struct htp_im2col_context * ictx) {
+ c->c0 += ictx->pe_owb;
+ if (c->c0 >= ictx->OW) {
+ c->c0 = 0;
+ c->r++;
+ c->ioh++;
+ if (c->ioh >= ictx->OH) {
+ c->ioh = 0;
+ c->iod++;
+ if (c->iod >= ictx->OD) {
+ c->iod = 0;
+ c->in++;
+ }
+ }
+ }
+}
+
+static void im2col_3d_blocked_dma_stage(struct htp_im2col_context * ictx, dma_queue * dma_q,
+ float * srcb, const struct im2col_cursor * c) {
+ const uint32_t in = c->in, iod = c->iod, ioh = c->ioh, c0 = c->c0;
+ const struct htp_ops_context * octx = ictx->octx;
const struct htp_tensor * src1 = octx->src[1];
- const struct htp_tensor * dst = octx->dst;
+ const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1], s2 = octx->op_params[2];
+ const int32_t p0 = octx->op_params[3], p1 = octx->op_params[4], p2 = octx->op_params[5];
+ const int32_t d1 = octx->op_params[7], d2 = octx->op_params[8];
+ const uint32_t IC = ictx->IC;
+ const uint32_t ID = ictx->ID, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KD = ictx->KD, KH = ictx->KH;
+ const uint32_t Wb = ictx->pe_wb;
+ const int32_t win0 = (int32_t) c0 * s0 - p0;
+ const dma_addr_t src_data = src1->data;
+ const size_t src_channel_stride = src1->nb[3];
+ const size_t vtcm_channel_stride = (size_t) KD * KH * Wb * sizeof(float);
+ const size_t batch_offset = (size_t) (in * IC) * src1->nb[3];
+
+ for (uint32_t ikd = 0; ikd < KD; ikd++) {
+ const int32_t iid = (int32_t) iod * s2 + (int32_t) ikd * d2 - p2;
+ const bool valid_d = (iid >= 0 && iid < (int32_t) ID);
+ for (uint32_t ikh = 0; ikh < KH; ikh++) {
+ const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1;
+ if (!valid_d || iih < 0 || iih >= (int32_t) IH) {
+ dma_queue_push(dma_q, dma_make_data(srcb, 0), 0, 0, 0, 0);
+ 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) {
+ dma_queue_push(dma_q, dma_make_data(srcb, 0), 0, 0, 0, 0);
+ continue;
+ }
+ float * vdst = srcb + (size_t) (ikd * KH + ikh) * Wb + lo;
+ const dma_addr_t vsrc = src_data + batch_offset +
+ (size_t) iid * src1->nb[2] +
+ (size_t) iih * src1->nb[1] +
+ (size_t) (win0 + lo) * src1->nb[0];
+ dma_queue_push(dma_q, dma_make_data(vdst, vsrc),
+ vtcm_channel_stride, src_channel_stride,
+ (size_t) (hi - lo) * sizeof(float), IC);
+ }
+ }
+}
- 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;
+#define IM2COL_3D_BLOCKED_COMPUTE_BODY(FNAME, DST_CTYPE, COPY_FN, SPLAT_FN) \
+static void FNAME(const struct htp_im2col_context * ictx, \
+ void * restrict dstb_raw, const float * restrict srcb, \
+ float * restrict scratch, \
+ uint32_t nb, uint32_t c0, uint32_t iod, uint32_t ioh) { \
+ DST_CTYPE * restrict dstb = (DST_CTYPE *) dstb_raw; \
+ const struct htp_ops_context * octx = ictx->octx; \
+ const int32_t s0 = octx->op_params[0], s1 = octx->op_params[1], s2 = octx->op_params[2]; \
+ const int32_t p0 = octx->op_params[3], p1 = octx->op_params[4], p2 = octx->op_params[5]; \
+ const int32_t d0 = octx->op_params[6], d1 = octx->op_params[7], d2 = octx->op_params[8]; \
+ const uint32_t IC = ictx->IC; \
+ const uint32_t ID = ictx->ID, IH = ictx->IH, IW = ictx->IW; \
+ const uint32_t KD = ictx->KD, KH = ictx->KH, KW = ictx->KW; \
+ const uint32_t patch_stride = ictx->patch_stride; \
+ const uint32_t Wb = ictx->pe_wb; \
+ const int32_t win0 = (int32_t) c0 * s0 - p0; \
+ for (uint32_t j = 0; j < nb; j++) { \
+ const uint32_t iow = c0 + j; \
+ const int32_t iiw0 = (int32_t) iow * s0 - p0; \
+ DST_CTYPE * dst_patch = dstb + (size_t) j * patch_stride; \
+ for (uint32_t iic = 0; iic < IC; iic++) { \
+ for (uint32_t ikd = 0; ikd < KD; ikd++) { \
+ const int32_t iid = (int32_t) iod * s2 + (int32_t) ikd * d2 - p2; \
+ for (uint32_t ikh = 0; ikh < KH; ikh++) { \
+ const int32_t iih = (int32_t) ioh * s1 + (int32_t) ikh * d1 - p1; \
+ DST_CTYPE * out_run = dst_patch + \
+ (size_t) (iic * KD * KH + ikd * KH + ikh) * KW; \
+ if (iid < 0 || iid >= (int32_t) ID || iih < 0 || iih >= (int32_t) IH) { \
+ SPLAT_FN(out_run, 0, KW); \
+ continue; \
+ } \
+ const float * vrow = srcb + (size_t) (iic * KD * KH + ikd * KH + ikh) * Wb; \
+ if (d0 == 1) { \
+ 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, KW); \
+ } else { \
+ if (lo > 0) { \
+ SPLAT_FN(out_run, 0, (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, (uint32_t) (KW - hi)); \
+ } \
+ } \
+ continue; \
+ } \
+ for (uint32_t ikw = 0; ikw < KW; ikw += 64) { \
+ const uint32_t chunk = MIN(KW - ikw, 64); \
+ for (uint32_t c = 0; c < chunk; c++) { \
+ const int32_t iiw = iiw0 + (int32_t) (ikw + c) * d0; \
+ scratch[c] = (iiw < 0 || iiw >= (int32_t) IW) ? 0.0f : vrow[iiw - win0]; \
+ } \
+ COPY_FN((uint8_t *) (out_run + ikw), (const uint8_t *) scratch, chunk); \
+ } \
+ } \
+ } \
+ } \
+ } \
+}
+
+IM2COL_3D_BLOCKED_COMPUTE_BODY(im2col_3d_blocked_compute_f16, uint16_t, hvx_copy_f16_f32_uu, hvx_splat_u16_u)
+IM2COL_3D_BLOCKED_COMPUTE_BODY(im2col_3d_blocked_compute_f32, float, hvx_copy_f32_uu, hvx_splat_f32_u)
+
+static void im2col_3d_blocked_dma_thread(unsigned int nth, unsigned int ith, void * data) {
+ (void) nth;
+ 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 dst = octx->dst;
+ const uint32_t OW = ictx->OW;
+ const uint32_t KD = ictx->KD, KH = ictx->KH;
+ const uint32_t kd_kh = KD * KH;
+ const uint32_t patch_stride = ictx->patch_stride;
+ const uint32_t owb = ictx->pe_owb;
+ const uint32_t nbpr = ictx->pe_nbpr;
+ const uint32_t dst_elem = ictx->dst_elem;
+ const dma_addr_t dst_data = dst->data;
+ dma_queue * dma_q = 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) };
+ void * dstb2[2] = { dstb_base, dstb_base + ictx->pe_dst_row_bytes };
+ float * scratch = (float *) (srcb_base + 2 * ictx->pe_src_row_bytes);
+ const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows;
+ const uint32_t row_start = ictx->pe_row_base + ictx->pe_rows_per_thread * ith;
+ const uint32_t row_end = MIN(row_start + ictx->pe_rows_per_thread, row_end_max);
+ if (row_start >= row_end) {
+ return;
+ }
+ const uint32_t nrows_local = row_end - row_start;
+ const uint32_t total_blocks = nrows_local * nbpr;
+ if (total_blocks == 0) {
+ return;
+ }
+ struct im2col_cursor cur = {
+ .r = row_start,
+ .in = ictx->thread_in[ith],
+ .iod = ictx->thread_iod[ith],
+ .ioh = ictx->thread_ioh[ith],
+ .c0 = 0,
+ };
+ struct im2col_cursor stage = cur;
+ for (uint32_t i = 0; i < total_blocks && i < 2; i++) {
+ dma_queue_push(dma_q, dma_make_data(dst_data, dstb2[i]), 0, 0, 0, 0);
+ im2col_3d_blocked_dma_stage(ictx, dma_q, srcb2[i], &stage);
+ im2col_3d_step_block(&stage, ictx);
}
- 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;
+ im2col_3d_blocked_compute_fn_t compute = (im2col_3d_blocked_compute_fn_t) ictx->compute;
+ for (uint32_t bi = 0; bi < total_blocks; bi++) {
+ const uint32_t buf = bi & 1u;
+ float * srcb = srcb2[buf];
+ void * dstb = dstb2[buf];
+ const uint32_t nb = MIN(owb, OW - cur.c0);
+ dma_queue_pop(dma_q);
+ for (uint32_t p = 0; p < kd_kh; p++) {
+ dma_queue_pop(dma_q);
+ }
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ compute(ictx, dstb, srcb, scratch, nb, cur.c0, cur.iod, cur.ioh);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ const dma_addr_t ddr = dst_data + (size_t) (cur.r * OW + cur.c0) * patch_stride * dst_elem;
+ dma_queue_push(dma_q, dma_make_data(ddr, dstb),
+ nb * patch_stride * dst_elem, nb * patch_stride * dst_elem,
+ nb * patch_stride * dst_elem, 1);
+ if (bi + 2 < total_blocks) {
+ im2col_3d_blocked_dma_stage(ictx, dma_q, srcb, &stage);
+ im2col_3d_step_block(&stage, ictx);
+ }
+ im2col_3d_step_block(&cur, ictx);
+ }
+ dma_queue_flush(dma_q);
+}
+static void im2col_3d_patchembed_dma_stage(struct htp_im2col_context * ictx, dma_queue * dma_q,
+ float * srcb, const struct im2col_cursor * c) {
+ const struct htp_ops_context * octx = ictx->octx;
+ const uint32_t IC = ictx->IC;
+ const uint32_t ID = ictx->ID, IH = ictx->IH, IW = ictx->IW;
+ const uint32_t KD = ictx->KD, KH = ictx->KH;
+ const dma_addr_t src_data = octx->src[1]->data;
+ const uint32_t in = c->in, iod = c->iod, ioh = c->ioh;
+ const size_t src_channel_stride = (size_t) ID * IH * IW * sizeof(float);
+ const size_t vtcm_channel_stride = (size_t) KD * KH * IW * sizeof(float);
+ const size_t row_size = (size_t) KH * IW * sizeof(float);
+
+ for (uint32_t ikd = 0; ikd < KD; ikd++) {
+ const uint32_t iid = iod * KD + ikd;
+ float * vdst = srcb + (size_t) ikd * KH * IW;
+ const dma_addr_t vsrc = src_data + (size_t) (((in * IC) * ID + iid) * IH + ioh * KH) * IW * sizeof(float);
+ dma_queue_push(dma_q, dma_make_data(vdst, vsrc),
+ vtcm_channel_stride, src_channel_stride, row_size, IC);
+ }
+}
+static void im2col_3d_patchembed_compute_f16(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb_raw, const float * restrict srcb) {
+ uint16_t * restrict dstb = (uint16_t *) dstb_raw;
+ const uint32_t OW = ictx->OW;
+ const uint32_t IC = ictx->IC;
+ const uint32_t KD = ictx->KD;
+ const uint32_t KH = ictx->KH;
+ const uint32_t KW = ictx->KW;
+ const uint32_t IW = ictx->IW;
+ const uint32_t patch_stride = ictx->patch_stride;
+
+ for (uint32_t iow = 0; iow < OW; iow++) {
+ uint16_t * dst_patch = dstb + (size_t) iow * patch_stride;
+ for (uint32_t iic = 0; iic < IC; iic++) {
+ for (uint32_t ikd = 0; ikd < KD; ikd++) {
+ for (uint32_t ikh = 0; ikh < KH; ikh++) {
+ uint16_t * out_run = dst_patch + (size_t) (iic * KD * KH + ikd * KH + ikh) * KW;
+ const float * src_run = srcb + (size_t) (iic * KD * KH + ikd * KH + ikh) * IW + (size_t) iow * KW;
+ hvx_copy_f16_f32_uu((uint8_t *) out_run, (const uint8_t *) src_run, KW);
+ }
+ }
+ }
+ }
+}
- 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;
+static void im2col_3d_patchembed_compute_f32(
+ const struct htp_im2col_context * ictx,
+ void * restrict dstb_raw, const float * restrict srcb) {
+ float * restrict dstb = (float *) dstb_raw;
+ const uint32_t OW = ictx->OW;
+ const uint32_t IC = ictx->IC;
+ const uint32_t KD = ictx->KD;
+ const uint32_t KH = ictx->KH;
+ const uint32_t KW = ictx->KW;
+ const uint32_t IW = ictx->IW;
+ const uint32_t patch_stride = ictx->patch_stride;
+
+ for (uint32_t iow = 0; iow < OW; iow++) {
+ float * dst_patch = dstb + (size_t) iow * patch_stride;
+ for (uint32_t iic = 0; iic < IC; iic++) {
+ for (uint32_t ikd = 0; ikd < KD; ikd++) {
+ for (uint32_t ikh = 0; ikh < KH; ikh++) {
+ float * out_run = dst_patch + (size_t) (iic * KD * KH + ikd * KH + ikh) * KW;
+ const float * src_run = srcb + (size_t) (iic * KD * KH + ikd * KH + ikh) * IW + (size_t) iow * KW;
+ hvx_copy_f32_uu((uint8_t *) out_run, (const uint8_t *) src_run, KW);
+ }
+ }
+ }
+ }
+}
+
+static void im2col_3d_patchembed_dma_thread(unsigned int nth, unsigned int ith, void * data) {
+ (void) nth;
+ 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 dst = octx->dst;
+ const uint32_t KD = ictx->KD;
+ const uint32_t OW = ictx->OW;
+ const uint32_t patch_stride = ictx->patch_stride;
+ const uint32_t dst_elem = ictx->dst_elem;
+ const dma_addr_t dst_data = dst->data;
+ dma_queue * dma_q = octx->ctx->dma[ith];
+ const uint32_t row_end_max = ictx->pe_row_base + ictx->pe_nrows;
+ const uint32_t row_start = ictx->pe_row_base + ictx->pe_rows_per_thread * ith;
+ const uint32_t row_end = MIN(row_start + ictx->pe_rows_per_thread, row_end_max);
+ if (row_start >= row_end) {
+ return;
+ }
+
+ const uint32_t nrows_local = row_end - row_start;
+ uint8_t * src_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 *) src_base, (float *) (src_base + ictx->pe_src_row_bytes) };
+ void * dstb2[2] = { dstb_base, dstb_base + ictx->pe_dst_row_bytes };
+ struct im2col_cursor cur = {
+ .r = row_start,
+ .in = ictx->thread_in[ith],
+ .iod = ictx->thread_iod[ith],
+ .ioh = ictx->thread_ioh[ith],
+ .c0 = 0,
+ };
+ struct im2col_cursor stage = cur;
+ for (uint32_t i = 0; i < nrows_local && i < 2; i++) {
+ dma_queue_push(dma_q, dma_make_data(dst_data, dstb2[i]), 0, 0, 0, 0);
+ im2col_3d_patchembed_dma_stage(ictx, dma_q, srcb2[i], &stage);
+ im2col_3d_step_row(&stage, ictx);
}
+ im2col_3d_patchembed_compute_fn_t compute = (im2col_3d_patchembed_compute_fn_t) ictx->compute;
+ for (uint32_t idx = 0; idx < nrows_local; idx++) {
+ const uint32_t buf = idx & 1u;
+ float * srcb = srcb2[buf];
+ void * dstb = dstb2[buf];
+ dma_queue_pop(dma_q);
+ for (uint32_t p = 0; p < KD; p++) {
+ dma_queue_pop(dma_q);
+ }
+ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ compute(ictx, dstb, srcb);
+ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, cur.r);
+ const dma_addr_t ddr_row = dst_data + (size_t) cur.r * OW * patch_stride * dst_elem;
+ dma_queue_push(dma_q, dma_make_data(ddr_row, dstb),
+ OW * patch_stride * dst_elem, OW * patch_stride * dst_elem,
+ OW * patch_stride * dst_elem, 1);
+ if (idx + 2 < nrows_local) {
+ im2col_3d_patchembed_dma_stage(ictx, dma_q, srcb, &stage);
+ im2col_3d_step_row(&stage, ictx);
+ }
+ im2col_3d_step_row(&cur, ictx);
+ }
+ dma_queue_flush(dma_q);
+}
+
+int op_im2col(struct htp_ops_context * octx) {
+ const struct htp_im2col_kernel_params * kparams = (const struct htp_im2col_kernel_params *) octx->kernel_params;
+ if (kparams->kernel_type == HTP_IM2COL_KERNEL_UNSUPPORTED) {
+ return HTP_STATUS_NO_SUPPORT;
+ }
+
+ if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+ return HTP_STATUS_INVAL_PARAMS;
+ }
+
+ const struct htp_tensor * dst = octx->dst;
+
uint32_t row_base = 0;
- uint32_t nrows = total_rows;
+ uint32_t nrows = kparams->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 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,
+ htp_tensor_mdev_partition(kparams->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) {
+ if (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) -> 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) {
- 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;
- }
- }
+ uint8_t * const base = octx->ctx->vtcm_base;
+
+ const uint32_t pth = MIN(octx->n_threads, nrows);
+ ictx.octx = octx;
+ ictx.pe_row_base = row_base;
+ ictx.pe_nrows = nrows;
+ ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
+ ictx.pe_src_row_bytes = kparams->src_row_bytes;
+ ictx.pe_dst_row_bytes = kparams->dst_row_bytes;
+ ictx.pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
+ ictx.pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
+ ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
+ ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+ ictx.pe_owb = kparams->owb;
+ ictx.pe_wb = kparams->wb;
+ ictx.pe_nbpr = kparams->nbpr;
+ ictx.N = kparams->N;
+ ictx.IC = kparams->IC;
+ ictx.ID = kparams->ID;
+ ictx.IH = kparams->IH;
+ ictx.IW = kparams->IW;
+ ictx.KD = kparams->KD;
+ ictx.KH = kparams->KH;
+ ictx.KW = kparams->KW;
+ ictx.OD = kparams->OD;
+ ictx.OH = kparams->OH;
+ ictx.OW = kparams->OW;
+ ictx.patch_stride = kparams->patch_stride;
+
+ for (uint32_t ith = 0; ith < pth; ith++) {
+ uint32_t r = row_base + ictx.pe_rows_per_thread * ith;
+ ictx.thread_ioh[ith] = r % kparams->OH;
+ ictx.thread_in[ith] = r / kparams->OH;
}
- // Fall through to pure-DDR.
- if (npatches == 0) {
+
+ ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);
+
+ if (kparams->kernel_type == HTP_IM2COL_KERNEL_PATCHEMBED_DMA) {
+ ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_patchembed_compute_f16
+ : (void *) im2col_patchembed_compute_f32;
+ work_queue_run(octx->ctx->work_queue, im2col_patchembed_dma_thread, &ictx, pth);
+ return HTP_STATUS_OK;
+ } else if (kparams->kernel_type == HTP_IM2COL_KERNEL_BLOCKED_DMA) {
+ ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_blocked_compute_f16
+ : (void *) im2col_blocked_compute_f32;
+ work_queue_run(octx->ctx->work_queue, im2col_blocked_dma_thread, &ictx, pth);
return HTP_STATUS_OK;
}
- if (htp_tensor_is_extended(src1) || htp_tensor_is_extended(dst)) {
+ return HTP_STATUS_NO_SUPPORT;
+}
+
+int op_im2col_3d(struct htp_ops_context * octx) {
+ const struct htp_im2col_kernel_params * kparams =
+ (const struct htp_im2col_kernel_params *) octx->kernel_params;
+ if (kparams->kernel_type == HTP_IM2COL_KERNEL_UNSUPPORTED) {
return HTP_STATUS_NO_SUPPORT;
}
- 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);
+ if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+ return HTP_STATUS_INVAL_PARAMS;
}
- return HTP_STATUS_OK;
+
+ const struct htp_tensor * dst = octx->dst;
+
+ uint32_t row_base = 0;
+ uint32_t nrows = kparams->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(kparams->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 (nrows == 0) {
+ return HTP_STATUS_OK;
+ }
+
+ struct htp_im2col_context ictx = { 0 };
+ uint8_t * const base = octx->ctx->vtcm_base;
+
+ const uint32_t pth = MIN(octx->n_threads, nrows);
+ ictx.octx = octx;
+ ictx.pe_row_base = row_base;
+ ictx.pe_nrows = nrows;
+ ictx.pe_rows_per_thread = (nrows + pth - 1) / pth;
+ ictx.pe_src_row_bytes = kparams->src_row_bytes;
+ ictx.pe_dst_row_bytes = kparams->dst_row_bytes;
+ ictx.pe_vtcm_src = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_src);
+ ictx.pe_vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, kparams->off_dst);
+ ictx.pe_src_size_per_thread = kparams->src_bytes_per_thread;
+ ictx.pe_dst_size_per_thread = kparams->dst_bytes_per_thread;
+ ictx.pe_owb = kparams->owb;
+ ictx.pe_wb = kparams->wb;
+ ictx.pe_nbpr = kparams->nbpr;
+ ictx.N = kparams->N;
+ ictx.IC = kparams->IC;
+ ictx.ID = kparams->ID;
+ ictx.IH = kparams->IH;
+ ictx.IW = kparams->IW;
+ ictx.KD = kparams->KD;
+ ictx.KH = kparams->KH;
+ ictx.KW = kparams->KW;
+ ictx.OD = kparams->OD;
+ ictx.OH = kparams->OH;
+ ictx.OW = kparams->OW;
+ ictx.patch_stride = kparams->patch_stride;
+
+ for (uint32_t ith = 0; ith < pth; ith++) {
+ uint32_t rem = row_base + ictx.pe_rows_per_thread * ith;
+ ictx.thread_ioh[ith] = rem % kparams->OH;
+ rem /= kparams->OH;
+ ictx.thread_iod[ith] = rem % kparams->OD;
+ ictx.thread_in[ith] = rem / kparams->OD;
+ }
+
+ ictx.dst_elem = dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float);
+
+ if (kparams->kernel_type == HTP_IM2COL_KERNEL_PATCHEMBED_DMA) {
+ ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_3d_patchembed_compute_f16
+ : (void *) im2col_3d_patchembed_compute_f32;
+ work_queue_run(octx->ctx->work_queue, im2col_3d_patchembed_dma_thread, &ictx, pth);
+ return HTP_STATUS_OK;
+ } else if (kparams->kernel_type == HTP_IM2COL_KERNEL_BLOCKED_DMA) {
+ ictx.compute = dst->type == HTP_TYPE_F16 ? (void *) im2col_3d_blocked_compute_f16
+ : (void *) im2col_3d_blocked_compute_f32;
+ work_queue_run(octx->ctx->work_queue, im2col_3d_blocked_dma_thread, &ictx, pth);
+ return HTP_STATUS_OK;
+ }
+
+ return HTP_STATUS_NO_SUPPORT;
}
diff --git a/ggml/src/ggml-hexagon/htp/im2col-ops.h b/ggml/src/ggml-hexagon/htp/im2col-ops.h
new file mode 100644
index 000000000..9ccc21fd7
--- /dev/null
+++ b/ggml/src/ggml-hexagon/htp/im2col-ops.h
@@ -0,0 +1,68 @@
+#ifndef HTP_IM2COL_OPS_H
+#define HTP_IM2COL_OPS_H
+
+#include "hex-common.h"
+#include <stdint.h>
+
+enum htp_im2col_kernel_type {
+ HTP_IM2COL_KERNEL_UNSUPPORTED = 0,
+ HTP_IM2COL_KERNEL_PATCHEMBED_DMA = 1,
+ HTP_IM2COL_KERNEL_BLOCKED_DMA = 2,
+};
+
+struct htp_im2col_vtcm_layout {
+ uint32_t off_src;
+ uint32_t off_dst;
+ uint32_t src_bytes_per_thread;
+ uint32_t dst_bytes_per_thread;
+ uint32_t total_bytes;
+};
+
+static inline void htp_im2col_vtcm_layout_build(
+ struct htp_im2col_vtcm_layout * layout,
+ uint32_t src_row_bytes,
+ uint32_t dst_row_bytes,
+ uint32_t n_threads) {
+
+ layout->src_bytes_per_thread = 2 * src_row_bytes + 256;
+ layout->dst_bytes_per_thread = 2 * dst_row_bytes;
+
+ layout->off_src = 0;
+ layout->off_dst = layout->off_src + layout->src_bytes_per_thread * n_threads;
+ layout->total_bytes = layout->off_dst + layout->dst_bytes_per_thread * n_threads;
+}
+
+struct htp_im2col_kernel_params {
+ uint8_t kernel_type;
+ uint8_t is_2d;
+ uint8_t is_3d;
+ uint8_t n_threads;
+
+ uint32_t vtcm_size;
+ uint32_t off_src;
+ uint32_t off_dst;
+ uint32_t src_bytes_per_thread;
+ uint32_t dst_bytes_per_thread;
+ uint32_t src_row_bytes;
+ uint32_t dst_row_bytes;
+
+ uint32_t owb;
+ uint32_t wb;
+ uint32_t nbpr;
+
+ uint32_t N;
+ uint32_t IC;
+ uint32_t ID, IH, IW;
+ uint32_t KD, KH, KW;
+ uint32_t OD, OH, OW;
+ uint32_t patch_stride;
+ uint32_t total_rows;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_im2col_kernel_params) <= 128, "htp_im2col_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_im2col_kernel_params) <= 128, "htp_im2col_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_IM2COL_OPS_H
diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c
index 2dd1b3fc1..6e3482b91 100644
--- a/ggml/src/ggml-hexagon/htp/main.c
+++ b/ggml/src/ggml-hexagon/htp/main.c
@@ -929,6 +929,9 @@ static int execute_op(struct htp_ops_context * octx) {
case HTP_OP_IM2COL:
return op_im2col(octx);
+ case HTP_OP_IM2COL_3D:
+ return op_im2col_3d(octx);
+
case HTP_OP_ROLL:
return op_roll(octx);