From 9afe7917bcba161f3733dc59411c1c75516cad39 Mon Sep 17 00:00:00 2001 From: Aparna M P Date: Tue, 6 Oct 2026 04:04:32 +0530 Subject: [PATCH] hexagon: add pool op support (llama/29995) * hexagon: add pool_2d support * hexagon: add pool_1d support * hex-pool: dma changes * hex-pool: Optimize HTP pooling boundaries and DMA pipelining * hex-pool: code cleanup and correctness fixes * hex-pool: re-write the DMA pipeline * hex-pool: pool chunking support * hex-pool: remove/vectorize all scalar paths * hex-pool: simplify chunk solver (no need for a loop) * hex-pool: remove redundant checks --------- Co-authored-by: Max Krasnyansky --- ggml/src/ggml-hexagon/ggml-hexagon.cpp | 185 +++++++++++ ggml/src/ggml-hexagon/htp/CMakeLists.txt | 1 + ggml/src/ggml-hexagon/htp/htp-ctx.h | 2 + ggml/src/ggml-hexagon/htp/htp-ops.h | 2 + ggml/src/ggml-hexagon/htp/main.c | 6 + ggml/src/ggml-hexagon/htp/pool-ops.c | 399 +++++++++++++++++++++++ ggml/src/ggml-hexagon/htp/pool-ops.h | 85 +++++ 7 files changed, 680 insertions(+) create mode 100644 ggml/src/ggml-hexagon/htp/pool-ops.c create mode 100644 ggml/src/ggml-hexagon/htp/pool-ops.h diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 2282a04a8..1e967086b 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -66,6 +66,7 @@ #include "htp/get-rows-ops.h" #include "htp/set-rows-ops.h" #include "htp/softmax-ops.h" +#include "htp/pool-ops.h" #include "htp/rope-ops.h" #include "htp/ssm-conv.h" #include "htp/gated-delta-net-ops.h" @@ -391,6 +392,14 @@ static void ggml_hexagon_precompute_sort_params( struct htp_sort_kernel_params * kparams ); +static void ggml_hexagon_precompute_pool_2d_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * dst, + struct htp_pool_2d_kernel_params * kparams, + bool is_pool_1d +); + static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, @@ -6256,6 +6265,96 @@ static void ggml_hexagon_precompute_sort_params( kparams->n_slots = (int32_t) layout.n_slots; } +static void ggml_hexagon_pool_interior_range( + uint32_t src_x, uint32_t dst_x, uint32_t kernel_x, uint32_t stride_x, int32_t pad_x, + uint32_t * ox_lo, uint32_t * ox_hi) { + const uint32_t lo_raw = ((uint32_t) pad_x + stride_x - 1) / stride_x; + *ox_lo = (lo_raw < dst_x) ? lo_raw : dst_x; + + const int32_t numer_hi = (int32_t) src_x - (int32_t) kernel_x + pad_x; + if (numer_hi < 0) { + *ox_hi = 0; + } else { + const uint32_t hi_raw = (uint32_t) numer_hi / stride_x + 1; + *ox_hi = (hi_raw < dst_x) ? hi_raw : dst_x; + } + if (*ox_hi < *ox_lo) { + *ox_hi = *ox_lo; + } +} + +static bool ggml_hexagon_pool_shape_fits( + const struct ggml_tensor * src0, + const struct ggml_tensor * dst +) { + return ggml_nbytes(src0) <= UINT32_MAX - 256 && ggml_nbytes(dst) <= UINT32_MAX - 256; +} + +static void ggml_hexagon_precompute_pool_2d_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * dst, + struct htp_pool_2d_kernel_params * kparams, + bool is_pool_1d +) { + memset(kparams, 0, sizeof(*kparams)); + + const uint32_t src_x = (uint32_t) src0->ne[0]; + const uint32_t src_y = is_pool_1d ? 1 : (uint32_t) src0->ne[1]; + const uint32_t dst_x = (uint32_t) dst->ne[0]; + const uint32_t dst_y = is_pool_1d ? 1 : (uint32_t) dst->ne[1]; + const uint32_t src_plane_bytes = src_x * src_y * sizeof(float); + const uint32_t dst_plane_bytes = dst_x * dst_y * sizeof(float); + const uint32_t planes = (uint32_t) (is_pool_1d + ? (uint64_t) src0->ne[1] * (uint64_t) src0->ne[2] * (uint64_t) src0->ne[3] + : (uint64_t) src0->ne[2] * (uint64_t) src0->ne[3]); + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1); + + struct htp_pool_vtcm_layout layout; + const bool ok = htp_pool_solve_layout(&layout, src_x, src_y, dst_x, dst_y, n_threads, sess->vtcm_size); + GGML_ASSERT(ok); + + kparams->src_x = src_x; + kparams->src_y = src_y; + kparams->dst_x = dst_x; + kparams->dst_y = dst_y; + kparams->kernel_x = (uint32_t) ggml_get_op_params_i32(dst, 1); + kparams->kernel_y = is_pool_1d ? 1 : (uint32_t) ggml_get_op_params_i32(dst, 2); + kparams->stride_x = (uint32_t) ggml_get_op_params_i32(dst, is_pool_1d ? 2 : 3); + kparams->stride_y = is_pool_1d ? 1 : (uint32_t) ggml_get_op_params_i32(dst, 4); + kparams->pad_x = ggml_get_op_params_i32(dst, is_pool_1d ? 3 : 5); + kparams->pad_y = is_pool_1d ? 0 : ggml_get_op_params_i32(dst, 6); + kparams->src_plane_bytes = src_plane_bytes; + kparams->dst_plane_bytes = dst_plane_bytes; + kparams->src_plane_bytes_aligned = (uint32_t) layout.src_spad_half_size; + kparams->dst_plane_bytes_aligned = (uint32_t) layout.dst_spad_half_size; + kparams->n_threads = n_threads; + kparams->planes = planes; + kparams->pool_op = (uint32_t) ggml_get_op_params_i32(dst, 0); + // Fast HVX path requires exact tiling and a supported kernel width. + kparams->fast_path = (kparams->pad_x == 0 && kparams->pad_y == 0 && + kparams->stride_x == kparams->kernel_x && kparams->stride_y == kparams->kernel_y && + (kparams->kernel_x == 1 || kparams->kernel_x == 2)) ? 1 : 0; + kparams->global_path = (kparams->pad_x == 0 && kparams->pad_y == 0 && + kparams->stride_x == kparams->kernel_x && + kparams->stride_y == kparams->kernel_y && + kparams->kernel_x == kparams->src_x && + kparams->kernel_y == kparams->src_y && + kparams->dst_x == 1 && kparams->dst_y == 1) ? 1 : 0; + kparams->block_path = (kparams->pad_x == 0 && kparams->pad_y == 0 && + kparams->kernel_y == 1 && kparams->stride_y == 1 && + kparams->stride_x == kparams->kernel_x && + kparams->kernel_x >= 32) ? 1 : 0; + kparams->avg_divide_count = (is_pool_1d && kparams->pool_op == GGML_OP_POOL_AVG) ? 1 : 0; + ggml_hexagon_pool_interior_range(kparams->src_x, kparams->dst_x, kparams->kernel_x, + kparams->stride_x, kparams->pad_x, + &kparams->ox_lo, &kparams->ox_hi); + + const bool narrow_ok = (uint64_t) kparams->dst_x * kparams->kernel_x <= 32; + kparams->narrow_path = (kparams->fast_path && narrow_ok) ? 1 : 0; + kparams->inv_kernel_area = 1.0f / (float) (kparams->kernel_x * kparams->kernel_y); +} + static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, // W0 @@ -6701,6 +6800,76 @@ static bool ggml_hexagon_supported_argmax(const struct ggml_hexagon_session * se GGML_UNUSED(sess); } +static bool ggml_hexagon_supported_pool_2d(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { + const struct ggml_tensor * src0 = op->src[0]; + const int32_t * params = op->op_params; + + if (params[0] != GGML_OP_POOL_AVG && params[0] != GGML_OP_POOL_MAX) { + return false; + } + if (src0->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(op)) { + return false; + } + if (!ggml_hexagon_pool_shape_fits(src0, op)) { + return false; + } + + const int32_t kernel_x = params[1]; + const int32_t kernel_y = params[2]; + const int32_t stride_x = params[3]; + const int32_t stride_y = params[4]; + const int32_t pad_x = params[5]; + const int32_t pad_y = params[6]; + + // Keep invalid parameters on CPU, avoiding unsafe unsigned values in HTP. + if (kernel_x <= 0 || kernel_y <= 0 || stride_x <= 0 || stride_y <= 0 || + pad_x < 0 || pad_y < 0) { + return false; + } + + const uint32_t planes = (uint32_t) (src0->ne[2] * src0->ne[3]); + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1); + + struct htp_pool_vtcm_layout layout; + return htp_pool_solve_layout(&layout, (uint32_t) src0->ne[0], (uint32_t) src0->ne[1], + (uint32_t) op->ne[0], (uint32_t) op->ne[1], + n_threads, sess->vtcm_size); +} + +static bool ggml_hexagon_supported_pool_1d(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { + const struct ggml_tensor * src0 = op->src[0]; + const int32_t * params = op->op_params; + + if (params[0] != GGML_OP_POOL_AVG && params[0] != GGML_OP_POOL_MAX) { + return false; + } + if (src0->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) { + return false; + } + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(op)) { + return false; + } + if (!ggml_hexagon_pool_shape_fits(src0, op)) { + return false; + } + + const int32_t kernel = params[1]; + const int32_t stride = params[2]; + const int32_t pad = params[3]; + if (kernel <= 0 || stride <= 0 || pad < 0) { + return false; + } + + const uint32_t planes = (uint32_t) (src0->ne[1] * src0->ne[2] * src0->ne[3]); + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, planes > 0 ? planes : 1); + + struct htp_pool_vtcm_layout layout; + return htp_pool_solve_layout(&layout, (uint32_t) src0->ne[0], 1, (uint32_t) op->ne[0], 1, n_threads, sess->vtcm_size); +} + static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; const struct ggml_tensor * src1 = op->src[1]; @@ -7258,6 +7427,8 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_SET_ROWS: return HTP_OP_SET_ROWS; case GGML_OP_SUM: return HTP_OP_SUM; case GGML_OP_SUM_ROWS: return HTP_OP_SUM_ROWS; + case GGML_OP_POOL_2D: return HTP_OP_POOL_2D; + case GGML_OP_POOL_1D: return HTP_OP_POOL_1D; case GGML_OP_ARGSORT: return HTP_OP_ARGSORT; case GGML_OP_TOP_K: return HTP_OP_TOP_K; case GGML_OP_ARGMAX: return HTP_OP_ARGMAX; @@ -7546,6 +7717,12 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg node.opcode == HTP_OP_TOP_K, (struct htp_sort_kernel_params *) node.kernel_params ); + } else if (node.opcode == HTP_OP_POOL_2D || node.opcode == HTP_OP_POOL_1D) { + ggml_hexagon_precompute_pool_2d_params( + sess, node.node->src[0], node.dst(), + (struct htp_pool_2d_kernel_params *)node.kernel_params, + node.opcode == HTP_OP_POOL_1D + ); } computed_nodes.push_back(std::move(node)); } @@ -8363,6 +8540,14 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons supp = ggml_hexagon_supported_argmax(sess, op); break; + case GGML_OP_POOL_2D: + supp = ggml_hexagon_supported_pool_2d(sess, op); + break; + + case GGML_OP_POOL_1D: + supp = ggml_hexagon_supported_pool_1d(sess, op); + break; + case GGML_OP_SOFT_MAX: supp = ggml_hexagon_supported_softmax(sess, op); break; diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index 821f08c0b..787badf07 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -44,6 +44,7 @@ add_library(${HTP_LIB} SHARED argsort-ops.c im2col-ops.c roll-ops.c + pool-ops.c allreduce-ops.c ) diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index f9f682be4..87a1f4c7b 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -173,6 +173,8 @@ 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_pool_2d(struct htp_ops_context * octx); +int op_pool_1d(struct htp_ops_context * octx); int op_allreduce(struct htp_ops_context * octx); int op_roll(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 a6a3bb85a..e47b42e64 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -115,6 +115,8 @@ enum htp_op_code { HTP_OP_ARGMAX, HTP_OP_UNARY_GELU_ERF, HTP_OP_GLU_GEGLU_ERF, + HTP_OP_POOL_2D, + HTP_OP_POOL_1D, HTP_OP_INVALID }; diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index fcb0a7202..41f78b387 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -931,6 +931,12 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_ROLL: return op_roll(octx); + case HTP_OP_POOL_2D: + return op_pool_2d(octx); + + case HTP_OP_POOL_1D: + return op_pool_1d(octx); + case HTP_OP_CONCAT: return op_concat(octx); diff --git a/ggml/src/ggml-hexagon/htp/pool-ops.c b/ggml/src/ggml-hexagon/htp/pool-ops.c new file mode 100644 index 000000000..ea1f1a597 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/pool-ops.c @@ -0,0 +1,399 @@ +#pragma clang diagnostic ignored "-Wunused-variable" + +#include +#include + +#include "hex-common.h" +#include "dma-queue.h" +#include "hex-profile.h" +#include "htp-ctx.h" +#include "htp-ops.h" +#include "htp-tensor.h" +#include "hvx-inverse.h" +#include "hvx-types.h" +#include "hvx-utils.h" +#include "pool-ops.h" + +#define HTP_POOL_MAX 0 +#define HTP_POOL_AVG 1 + +// Fast path: exact non-overlapping tiling (stride == kernel, no padding), kernel_x in {1,2}. +// Every window is guaranteed fully in-bounds, so this never needs boundary clamping. +static void pool_plane_hvx( + const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) { + const bool is_max = (p->pool_op == HTP_POOL_MAX); + const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area); + const HVX_Vector seed = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0); + + for (uint32_t oy = 0; oy < p->dst_y; ++oy) { + const uint32_t sy = oy * p->kernel_y; + for (uint32_t ox = 0; ox < p->dst_x; ox += VLEN_FP32) { + const uint32_t rem = p->dst_x - ox; + const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN; + HVX_Vector acc = seed; + for (uint32_t ky = 0; ky < p->kernel_y; ++ky) { + const float * row = src + (sy + ky) * p->src_x; + if (p->kernel_x == 2) { + const HVX_Vector v0 = *(const HVX_UVector *) (row + ox * 2); + const HVX_Vector v1 = *(const HVX_UVector *) (row + ox * 2 + VLEN_FP32); + const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(v1, v0, -4); + const HVX_Vector lo = Q6_V_lo_W(deinterleaved); + const HVX_Vector hi = Q6_V_hi_W(deinterleaved); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, Q6_Vsf_vmax_VsfVsf(lo, hi)) + : hvx_vec_add_f32_f32(acc, hvx_vec_add_f32_f32(lo, hi)); + } else if (p->kernel_x == 1) { + const HVX_Vector v = *(const HVX_UVector *) (row + ox); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v); + } + } + hvx_vec_store_u(dst + oy * p->dst_x + ox, nbytes, is_max ? acc : hvx_vec_mul_f32_f32(acc, scale)); + } + } +} + +// Narrow exact-tiling path. The input is staged in VTCM with one vector of +// guard space, so full-width loads are safe even when src_x is below 32. +static void pool_plane_hvx_narrow( + const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) { + const bool is_max = (p->pool_op == HTP_POOL_MAX); + const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area); + const HVX_Vector seed = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0); + + for (uint32_t oy = 0; oy < p->dst_y; ++oy) { + const uint32_t sy = oy * p->kernel_y; + HVX_Vector acc = seed; + for (uint32_t ky = 0; ky < p->kernel_y; ++ky) { + const float * row = src + (sy + ky) * p->src_x; + const HVX_Vector v0 = *(const HVX_UVector *) row; + HVX_Vector v = v0; + if (p->kernel_x == 2) { + // The second vector is zero because a narrow row has fewer + // than 32 input elements. The low lanes still contain the + // complete even/odd pairs needed by the output. + const HVX_Vector zero = Q6_V_vsplat_R(0); + const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(zero, v0, -4); + const HVX_Vector even = Q6_V_lo_W(deinterleaved); + const HVX_Vector odd = Q6_V_hi_W(deinterleaved); + v = is_max ? Q6_Vsf_vmax_VsfVsf(even, odd) + : hvx_vec_add_f32_f32(even, odd); + } + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) + : hvx_vec_add_f32_f32(acc, v); + } + hvx_vec_store_u(dst + oy * p->dst_x, p->dst_x * sizeof(float), + is_max ? acc : hvx_vec_mul_f32_f32(acc, scale)); + } +} + +static void pool_plane_global( + const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) { + const uint32_t n = p->src_x * p->src_y; + float val; + if (p->pool_op == HTP_POOL_MAX) { + val = hvx_reduce_max_f32((const uint8_t *) src, n); + } else { + val = hvx_reduce_sum_f32((const uint8_t *) src, n) * p->inv_kernel_area; + } + hvx_vec_store_u(dst, sizeof(float), hvx_vec_splat_f32(val)); +} + +static void pool_plane_block( + const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) { + for (uint32_t oy = 0; oy < p->dst_y; ++oy) { + const float * row = src + oy * p->src_x; + float * dst_row = dst + oy * p->dst_x; + for (uint32_t ox = 0; ox < p->dst_x; ++ox) { + const uint8_t * block = (const uint8_t *) (row + ox * p->kernel_x); + float val; + if (p->pool_op == HTP_POOL_MAX) { + val = hvx_reduce_max_f32(block, p->kernel_x); + } else { + val = hvx_reduce_sum_f32(block, p->kernel_x) * p->inv_kernel_area; + } + hvx_vec_store_u(dst_row + ox, sizeof(float), hvx_vec_splat_f32(val)); + } + } +} + +// General path: arbitrary kernel/stride/padding + +// Vertical padding is uniform across a whole output row (iy0 depends only on oy, not ox), +// so it collapses to one valid-ky range per row instead of a per-element check. +static inline void pool_row_bounds_y( + const struct htp_pool_2d_kernel_params * p, uint32_t oy, int32_t * iy0, uint32_t * ky_lo, uint32_t * ky_hi) { + *iy0 = (int32_t) (oy * p->stride_y) - p->pad_y; + const int32_t lo = -(*iy0); + const int32_t hi = (int32_t) p->src_y - *iy0; + *ky_lo = (uint32_t) MAX(0, lo); + *ky_hi = (uint32_t) MAX(0, MIN((int32_t) p->kernel_y, hi)); +} + +static inline void pool_pixel_boundary_vec( + const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p, + uint32_t ox, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) { + const int32_t ix0 = (int32_t) (ox * p->stride_x) - p->pad_x; + const int32_t kx_lo = MAX(0, -ix0); + const int32_t kx_hi = MIN((int32_t) p->kernel_x, (int32_t) p->src_x - ix0); + + if (kx_lo >= kx_hi || ky_lo >= ky_hi) { + HVX_Vector empty_val = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0); + hvx_vec_store_u(dst_row + ox, sizeof(float), empty_val); + return; + } + + const uint32_t valid_kx = (uint32_t) (kx_hi - kx_lo); + const HVX_Vector mask_identity = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0); + + HVX_Vector acc = mask_identity; + for (uint32_t ky = ky_lo; ky < ky_hi; ++ky) { + const float * row = src + (uint32_t) (iy0 + (int32_t) ky) * p->src_x; + for (uint32_t k = 0; k < valid_kx; k += VLEN_FP32) { + const uint32_t k_rem = valid_kx - k; + const uint32_t n = (k_rem < VLEN_FP32) ? k_rem : VLEN_FP32; + const HVX_VectorPred q = Q6_Q_vsetq_R(n * sizeof(float)); + const HVX_Vector raw = *(const HVX_UVector *) (row + ix0 + kx_lo + k); + const HVX_Vector v = Q6_V_vmux_QVV(q, raw, mask_identity); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v); + } + } + + HVX_Vector reduced = is_max ? hvx_vec_reduce_max_f32(acc) : hvx_vec_reduce_sum_f32(acc); + if (!is_max) { + HVX_Vector scale_vec; + if (p->avg_divide_count) { + const uint32_t count = (ky_hi - ky_lo) * valid_kx; + scale_vec = hvx_vec_inverse_f32(hvx_vec_splat_f32((float) count)); + } else { + scale_vec = hvx_vec_splat_f32(p->inv_kernel_area); + } + reduced = hvx_vec_mul_f32_f32(reduced, scale_vec); + } + hvx_vec_store_u(dst_row + ox, sizeof(float), reduced); +} + +static inline void pool_row_general_boundary_vec( + const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p, + uint32_t ox_start, uint32_t ox_end, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) { + for (uint32_t ox = ox_start; ox < ox_end; ++ox) { + pool_pixel_boundary_vec(src, dst_row, p, ox, iy0, ky_lo, ky_hi, is_max); + } +} + +// Vectorized interior loop. +static inline void pool_row_general_vec( + const float * src, float * dst_row, const struct htp_pool_2d_kernel_params * p, + uint32_t ox_start, uint32_t ox_end, int32_t iy0, uint32_t ky_lo, uint32_t ky_hi, bool is_max) { + if (ox_start >= ox_end) { + return; + } + + if (p->stride_x != 1 && p->stride_x != 2) { + pool_row_general_boundary_vec(src, dst_row, p, ox_start, ox_end, iy0, ky_lo, ky_hi, is_max); + return; + } + + const HVX_Vector scale = hvx_vec_splat_f32(p->inv_kernel_area); + const HVX_Vector seed = is_max ? hvx_vec_splat_f32(-FLT_MAX) : Q6_V_vsplat_R(0); + + for (uint32_t ox = ox_start; ox < ox_end; ox += VLEN_FP32) { + const uint32_t rem = ox_end - ox; + const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN; + const int32_t ix0 = (int32_t) (ox * p->stride_x) - p->pad_x; + HVX_Vector acc = seed; + + for (uint32_t ky = ky_lo; ky < ky_hi; ++ky) { + const float * row = src + (uint32_t) (iy0 + (int32_t) ky) * p->src_x; + if (p->stride_x == 1) { + for (uint32_t kx = 0; kx < p->kernel_x; ++kx) { + const HVX_Vector v = *(const HVX_UVector *) (row + ix0 + (int32_t) kx); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, v) : hvx_vec_add_f32_f32(acc, v); + } + } else { + for (uint32_t kx = 0; kx < p->kernel_x; kx += 2) { + const HVX_Vector v0 = *(const HVX_UVector *) (row + ix0 + (int32_t) kx); + const HVX_Vector v1 = *(const HVX_UVector *) (row + ix0 + (int32_t) kx + VLEN_FP32); + const HVX_VectorPair deinterleaved = Q6_W_vdeal_VVR(v1, v0, -4); + const HVX_Vector lo = Q6_V_lo_W(deinterleaved); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, lo) : hvx_vec_add_f32_f32(acc, lo); + if (kx + 1 < p->kernel_x) { + const HVX_Vector hi = Q6_V_hi_W(deinterleaved); + acc = is_max ? Q6_Vsf_vmax_VsfVsf(acc, hi) : hvx_vec_add_f32_f32(acc, hi); + } + } + } + } + hvx_vec_store_u(dst_row + ox, nbytes, is_max ? acc : hvx_vec_mul_f32_f32(acc, scale)); + } +} + +static void pool_plane_general( + const float * src, float * dst, const struct htp_pool_2d_kernel_params * p) { + const bool is_max = (p->pool_op == HTP_POOL_MAX); + + for (uint32_t oy = 0; oy < p->dst_y; ++oy) { + int32_t iy0; + uint32_t ky_lo, ky_hi; + pool_row_bounds_y(p, oy, &iy0, &ky_lo, &ky_hi); + float * dst_row = dst + oy * p->dst_x; + + pool_row_general_boundary_vec(src, dst_row, p, 0, p->ox_lo, iy0, ky_lo, ky_hi, is_max); + pool_row_general_vec(src, dst_row, p, p->ox_lo, p->ox_hi, iy0, ky_lo, ky_hi, is_max); + pool_row_general_boundary_vec(src, dst_row, p, p->ox_hi, p->dst_x, iy0, ky_lo, ky_hi, is_max); + } +} + +typedef void (*pool_plane_fn_t)(const float * src, float * dst, const struct htp_pool_2d_kernel_params * p); + +struct pool_2d_context { + struct htp_ops_context * octx; + const struct htp_pool_2d_kernel_params * kparams; + pool_plane_fn_t pool_plane; + uint32_t n_threads; + uint32_t plane_start; + uint32_t plane_count; + uint32_t planes_per_thread; +}; + +static void pool_2d_thread(unsigned int nth, unsigned int ith, void * data) { + struct pool_2d_context * ctx = (struct pool_2d_context *) data; + const struct htp_pool_2d_kernel_params * p = ctx->kparams; + const struct htp_tensor * src0 = ctx->octx->src[0]; + const struct htp_tensor * dst = ctx->octx->dst; + pool_plane_fn_t pool_plane = ctx->pool_plane; + const uint32_t planes_per_thread = ctx->planes_per_thread; + const uint32_t first = ctx->plane_start + ith * planes_per_thread; + const uint32_t last = MIN(first + planes_per_thread, ctx->plane_start + ctx->plane_count); + + if (first >= last) { + return; + } + + struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith]; + dma_queue * dma_queue = ctx->octx->ctx->dma[ith]; + + const uint32_t src_spad_half = p->src_plane_bytes_aligned; + const uint32_t dst_spad_half = p->dst_plane_bytes_aligned; + const uint32_t src_bytes_per_thread = 2 * src_spad_half; + const uint32_t dst_bytes_per_thread = 2 * dst_spad_half; + const size_t off_dst = (size_t) ctx->n_threads * src_bytes_per_thread; + + uint8_t * vtcm_base = (uint8_t *) ctx->octx->ctx->vtcm_base; + uint8_t * src_spad = vtcm_base + ith * src_bytes_per_thread; + uint8_t * dst_spad = vtcm_base + off_dst + ith * dst_bytes_per_thread; + + float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) }; + float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) }; + + const uint32_t total = last - first; + + // Warm up the pipeline: push up to 2 initial (dummy dst, src) transfer pairs. + for (uint32_t i = 0; i < total && i < 2; ++i) { + dma_queue_push(dma_queue, + dma_make_data(dst->data, dstb2[i]), + p->dst_plane_bytes, dst_spad_half, + p->dst_plane_bytes, 0); + + const dma_addr_t src_addr = src0->data + (first + i) * p->src_plane_bytes; + dma_queue_push(dma_queue, + dma_make_data(srcb2[i], src_addr), + src_spad_half, p->src_plane_bytes, + p->src_plane_bytes, 1); + } + + for (uint32_t i = 0; i < total; ++i) { + const uint32_t plane = first + i; + const uint32_t buf = i & 1u; + float * srcb = srcb2[buf]; + float * dstb = dstb2[buf]; + + dma_queue_pop(dma_queue); // dst writeback from plane i - 2 (or dummy on iter 0, 1) + dma_queue_pop(dma_queue); // input for plane i + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane); + pool_plane(srcb, dstb, p); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) plane); + + const dma_addr_t dst_addr = dst->data + plane * p->dst_plane_bytes; + dma_queue_push(dma_queue, + dma_make_data(dst_addr, dstb), + p->dst_plane_bytes, dst_spad_half, + p->dst_plane_bytes, 1); + + if (i + 2 < total) { + const dma_addr_t next_src_addr = src0->data + (plane + 2) * p->src_plane_bytes; + dma_queue_push(dma_queue, + dma_make_data(srcb, next_src_addr), + src_spad_half, p->src_plane_bytes, + p->src_plane_bytes, 1); + } + } + + dma_queue_flush(dma_queue); + + FARF(HIGH, "pool2d-f32-dma %d/%d: %ux%ux%ux%u -> %ux%ux%ux%u (%u:%u)\n", + ith, nth, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], first, last); + (void) nth; +} + +int op_pool_2d(struct htp_ops_context * octx) { + const struct htp_tensor * src0 = octx->src[0]; + const struct htp_tensor * dst = octx->dst; + const struct htp_pool_2d_kernel_params * p = + (const struct htp_pool_2d_kernel_params *) octx->kernel_params; + + if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_F32 || + (p->pool_op != HTP_POOL_AVG && p->pool_op != HTP_POOL_MAX)) { + return HTP_STATUS_NO_SUPPORT; + } + + uint32_t plane_start = 0; + uint32_t plane_count = p->planes; + if (octx->ctx->mdev.count > 1) { + const uint32_t planes_per_chunk = (p->dst_plane_bytes > 0) ? (HEX_L2_LINE_SIZE / hex_gcd_u32(p->dst_plane_bytes, HEX_L2_LINE_SIZE)) : 1; + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition( + plane_count, htp_tensor_mdev_data_aligned(dst) ? planes_per_chunk : 0, + octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + plane_start = range.start; + plane_count = range.count; + } + if (plane_count == 0) { + return HTP_STATUS_OK; + } + + const uint32_t n_threads = MIN(p->n_threads, plane_count); + if (!htp_ops_context_set_n_threads(octx, n_threads)) { + return HTP_STATUS_INVAL_PARAMS; + } + + const uint32_t planes_per_thread = fastdiv(plane_count + n_threads - 1, &octx->n_threads_div); + + pool_plane_fn_t pool_plane; + if (p->global_path) { + pool_plane = pool_plane_global; + } else if (p->block_path) { + pool_plane = pool_plane_block; + } else if (p->narrow_path) { + pool_plane = pool_plane_hvx_narrow; + } else if (p->fast_path) { + pool_plane = pool_plane_hvx; + } else { + pool_plane = pool_plane_general; + } + + struct pool_2d_context ctx = { + .octx = octx, + .kparams = p, + .pool_plane = pool_plane, + .n_threads = n_threads, + .plane_start = plane_start, + .plane_count = plane_count, + .planes_per_thread = planes_per_thread, + }; + work_queue_run(octx->ctx->work_queue, pool_2d_thread, &ctx, n_threads); + return HTP_STATUS_OK; +} + +int op_pool_1d(struct htp_ops_context * octx) { + return op_pool_2d(octx); +} diff --git a/ggml/src/ggml-hexagon/htp/pool-ops.h b/ggml/src/ggml-hexagon/htp/pool-ops.h new file mode 100644 index 000000000..66b405a8b --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/pool-ops.h @@ -0,0 +1,85 @@ +#ifndef HTP_POOL_OPS_H +#define HTP_POOL_OPS_H + +#include +#include +#include +#include + +#include "hex-common.h" + +struct htp_pool_2d_kernel_params { + uint32_t src_x; + uint32_t src_y; + uint32_t dst_x; + uint32_t dst_y; + uint32_t kernel_x; + uint32_t kernel_y; + uint32_t stride_x; + uint32_t stride_y; + int32_t pad_x; + int32_t pad_y; + uint32_t src_plane_bytes; + uint32_t dst_plane_bytes; + uint32_t src_plane_bytes_aligned; + uint32_t dst_plane_bytes_aligned; + uint32_t n_threads; + uint32_t planes; + uint32_t pool_op; + uint32_t fast_path; + uint32_t narrow_path; + uint32_t global_path; + uint32_t block_path; + uint32_t avg_divide_count; + uint32_t ox_lo; + uint32_t ox_hi; + float inv_kernel_area; +}; + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_pool_2d_kernel_params) <= 128, "htp_pool_2d_kernel_params is too large"); +#else +_Static_assert(sizeof(struct htp_pool_2d_kernel_params) <= 128, "htp_pool_2d_kernel_params is too large"); +#endif + +struct htp_pool_vtcm_layout { + size_t total_bytes; + size_t off_src; + size_t off_dst; + size_t src_bytes_per_thread; + size_t dst_bytes_per_thread; + size_t src_spad_half_size; + size_t dst_spad_half_size; +}; + +static inline bool htp_pool_solve_layout( + struct htp_pool_vtcm_layout * layout, + uint32_t src_x, + uint32_t src_y, + uint32_t dst_x, + uint32_t dst_y, + uint32_t n_threads, + size_t vtcm_budget +) { + // Full-plane double buffering (256 bytes guard space for vector loads) + const size_t src_plane_bytes = (size_t) src_x * src_y * sizeof(float); + const size_t dst_plane_bytes = (size_t) dst_x * dst_y * sizeof(float); + const size_t src_plane_aligned = hex_round_up((uint32_t) src_plane_bytes + 256, 128); + const size_t dst_plane_aligned = hex_round_up((uint32_t) dst_plane_bytes, 128); + + const size_t spad_per_thread = 2 * (src_plane_aligned + dst_plane_aligned); + if (spad_per_thread * n_threads > vtcm_budget) { + return false; + } + + layout->src_spad_half_size = src_plane_aligned; + layout->dst_spad_half_size = dst_plane_aligned; + layout->src_bytes_per_thread = 2 * src_plane_aligned; + layout->dst_bytes_per_thread = 2 * dst_plane_aligned; + layout->off_src = 0; + layout->off_dst = layout->src_bytes_per_thread * n_threads; + layout->total_bytes = layout->off_dst + layout->dst_bytes_per_thread * n_threads; + return true; +} + +#endif // HTP_POOL_OPS_H