hexagon: add ARGSORT op

Co-authored-by: Yarden Tal <yardent@qti.qualcomm.com>
This commit is contained in:
Max Krasnyansky 2026-02-03 22:21:19 -08:00
parent b83111815e
commit 5bb5541e2c
8 changed files with 372 additions and 18 deletions

View File

@ -2111,6 +2111,21 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
return true;
}
static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const struct ggml_tensor * src0 = op->src[0]; // values
const struct ggml_tensor * dst = op; // indices
if (src0->type != GGML_TYPE_F32) {
return false;
}
if (dst->type != GGML_TYPE_I32) {
return false;
}
return true;
}
static bool ggml_hexagon_supported_rope(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
const int32_t * op_params = &op->op_params[0];
@ -2316,6 +2331,17 @@ static inline size_t init_get_rows_req(htp_general_req * req, dspqueue_buffer *
return n_bufs;
}
static inline size_t init_argsort_req(htp_general_req * req, dspqueue_buffer * bufs, const ggml_tensor * t) {
req->op = HTP_OP_ARGSORT;
memcpy(&req->op_params, &t->op_params, sizeof(t->op_params));
size_t n_bufs = 0;
n_bufs += htp_req_buff_init(&req->src0, &bufs[n_bufs], t->src[0], DSPQBUF_TYPE_CPU_WRITE_DSP_READ);
n_bufs += htp_req_buff_init(&req->dst, &bufs[n_bufs], t, DSPQBUF_TYPE_DSP_WRITE_CPU_READ);
return n_bufs;
}
template <bool _is_src0_constant>
static inline size_t init_binary_id_req(htp_general_req * req, dspqueue_buffer * bufs, const ggml_tensor * t) {
switch (t->op) {
@ -2564,6 +2590,10 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
ggml_hexagon_dispatch_op<init_cpy_req>(sess, node, flags);
break;
case GGML_OP_ARGSORT:
ggml_hexagon_dispatch_op<init_argsort_req>(sess, node, flags);
break;
default:
GGML_ABORT("\nggml-hex: graph-compute %s is not supported\n", ggml_op_desc(node));
}
@ -2968,6 +2998,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
supp = ggml_hexagon_supported_cpy(sess, op);
break;
case GGML_OP_ARGSORT:
supp = ggml_hexagon_supported_argsort(sess, op);
break;
default:
break;
}

View File

@ -6,6 +6,7 @@ include(${HEXAGON_SDK_ROOT}/build/cmake/hexagon_fun.cmake)
include_directories(
${HEXAGON_SDK_ROOT}/incs
${HEXAGON_SDK_ROOT}/incs/stddef
${CMAKE_CURRENT_SOURCE_DIR}/../../../include
${CMAKE_CURRENT_SOURCE_DIR}/../..
${CMAKE_CURRENT_SOURCE_DIR}/..
${CMAKE_CURRENT_SOURCE_DIR}
@ -28,6 +29,7 @@ add_library(${HTP_LIB} SHARED
set-rows-ops.c
get-rows-ops.c
cpy-ops.c
argsort-ops.c
)
target_compile_definitions(${HTP_LIB} PRIVATE

View File

@ -0,0 +1,281 @@
#include <string.h>
#include <stdlib.h>
#include <math.h>
#include <HAP_farf.h>
#include <HAP_perf.h>
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "ggml.h"
#include "hvx-utils.h"
#include "hex-dma.h"
#include "htp-ctx.h"
#include "htp-msg.h"
#include "htp-ops.h"
#ifndef MIN
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#endif
struct htp_argsort_context {
struct htp_ops_context * octx;
uint32_t nrows_per_thread;
};
static inline bool all_greater_f32(HVX_Vector x, HVX_Vector y)
{
const HVX_Vector one = Q6_V_vsplat_R(1);
const HVX_Vector zero = Q6_V_vzero();
HVX_VectorPred pred = Q6_Q_vcmp_gt_VsfVsf(x, y);
HVX_Vector matches = Q6_V_vmux_QVV(pred, one, zero);
HVX_Vector sum = hvx_vec_reduce_sum_i32(matches);
return hvx_vec_get_i32(sum) == 32;
}
// Sorts values and mirrors swaps to indices.
static void quicksort_values_indices_asc(float * values, int32_t * indices, int left, int right) {
if (left >= right) return;
int pivot_idx = (left + right) / 2;
float pivot = values[pivot_idx];
int i = left;
int j = right;
HVX_Vector pivot_vec = hvx_vec_splat_f32(pivot);
while (i <= j) {
// Vectorized scan for i
while (i <= j) {
// Check if we have at least one full vector
if (i + 32 <= j) {
HVX_Vector vals_vec = *(HVX_UVector *)(values + i);
if (all_greater_f32(pivot_vec, vals_vec)) {
// If all elements are < pivot, we can skip this whole block
i += 32;
continue;
}
}
// Scalar fallback / cleanup
if (values[i] < pivot) {
i++;
} else {
break;
}
}
// Vectorized scan for j
while (i <= j) {
if (j - 32 >= i) {
// Load 32 elements ending at j.
// Since we want `values[j] > pivot`, let's load from j-31 to j.
HVX_Vector vals_vec = *(HVX_UVector *)(values + j - 31);
if (all_greater_f32(vals_vec, pivot_vec)) {
j -= 32;
continue;
}
}
if (values[j] > pivot) {
j--;
} else {
break;
}
}
if (i <= j) {
float tmp_val = values[i];
values[i] = values[j];
values[j] = tmp_val;
int32_t tmp_idx = indices[i];
indices[i] = indices[j];
indices[j] = tmp_idx;
i++;
j--;
}
}
if (left < j) quicksort_values_indices_asc(values, indices, left, j);
if (i < right) quicksort_values_indices_asc(values, indices, i, right);
}
static void quicksort_values_indices_desc(float * values, int32_t * indices, int left, int right) {
if (left >= right) return;
int pivot_idx = (left + right) / 2;
float pivot = values[pivot_idx];
int i = left;
int j = right;
HVX_Vector pivot_vec = hvx_vec_splat_f32(pivot);
while (i <= j) {
// Vectorized scan for i (values[i] > pivot)
while (i <= j) {
if (i + 32 <= j) {
HVX_Vector vals_vec = *(HVX_UVector *)(values + i);
if (all_greater_f32(vals_vec, pivot_vec)) {
i += 32;
continue;
}
}
if (values[i] > pivot) {
i++;
} else {
break;
}
}
// Vectorized scan for j (values[j] < pivot)
while (i <= j) {
if (j - 32 >= i) {
HVX_Vector vals_vec = *(HVX_UVector *)(values + j - 31);
if (all_greater_f32(pivot_vec, vals_vec)) {
j -= 32;
continue;
}
}
if (values[j] < pivot) {
j--;
} else {
break;
}
}
if (i <= j) {
float tmp_val = values[i];
values[i] = values[j];
values[j] = tmp_val;
int32_t tmp_idx = indices[i];
indices[i] = indices[j];
indices[j] = tmp_idx;
i++;
j--;
}
}
if (left < j) quicksort_values_indices_desc(values, indices, left, j);
if (i < right) quicksort_values_indices_desc(values, indices, i, right);
}
static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
struct htp_argsort_context * actx = (struct htp_argsort_context *)data;
struct htp_ops_context * octx = actx->octx;
// Unpack context
const struct htp_tensor * src0 = &octx->src0;
const struct htp_tensor * dst = &octx->dst;
// Scratchpad memory
uint8_t * spad = octx->src0_spad.data + octx->src0_spad.size_per_thread * i;
// Dimensions
uint32_t ne00 = src0->ne[0];
uint32_t ne01 = src0->ne[1];
uint32_t ne02 = src0->ne[2];
uint32_t ne03 = src0->ne[3];
uint32_t nb01 = src0->nb[1];
//uint32_t nb02 = src0->nb[2];
//uint32_t nb03 = src0->nb[3];
uint32_t nb1 = dst->nb[1];
//uint32_t nb2 = dst->nb[2];
//uint32_t nb3 = dst->nb[3];
// Sort order
enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
// Rows to process
uint32_t total_rows = ne01 * ne02 * ne03;
uint32_t rows_per_thread = actx->nrows_per_thread;
uint32_t start_row = rows_per_thread * i;
uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
// Scratchpad layout:
// We need space for one row of float data (values) and one row of int32 indices.
// values: ne00 * sizeof(float)
// indices: ne00 * sizeof(int32_t)
// Padded to 128 bytes.
size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
float * values_buf = (float *) spad;
int32_t * indices_buf = (int32_t *) (spad + values_size);
for (uint32_t r = start_row; r < end_row; r++) {
uint32_t src_offset = r * nb01;
uint32_t dst_offset = r * nb1;
uint8_t * src_ptr = (uint8_t *) src0->data + src_offset;
uint8_t * dst_ptr = (uint8_t *) dst->data + dst_offset;
hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1);
hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00);
// Initialize indices
for (uint32_t j = 0; j < ne00; j++) {
indices_buf[j] = j;
}
// Sort values and mirror swaps to indices
if (order == GGML_SORT_ORDER_ASC) {
quicksort_values_indices_asc(values_buf, indices_buf, 0, ne00 - 1);
} else {
quicksort_values_indices_desc(values_buf, indices_buf, 0, ne00 - 1);
}
// Copy indices back to DDR
hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00);
}
}
int op_argsort(struct htp_ops_context * octx) {
// Check supported types
if (octx->src0.type != HTP_TYPE_F32) {
return HTP_STATUS_NO_SUPPORT;
}
// Allocate scratchpad
// We need 1 row of float + 1 row of int32 per thread.
uint32_t ne00 = octx->src0.ne[0];
size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
size_t indices_size = hex_round_up(ne00 * sizeof(int32_t), 128);
size_t spad_per_thread = values_size + indices_size;
// Make sure we round up to 256 for alignment requirements
spad_per_thread = hex_round_up(spad_per_thread, 256);
size_t total_spad_size = spad_per_thread * octx->n_threads;
if (octx->ctx->vtcm_size < total_spad_size) {
FARF(ERROR, "argsort: VTCM size too small. Needed %zu, have %zu", total_spad_size, octx->ctx->vtcm_size);
return HTP_STATUS_VTCM_TOO_SMALL;
}
octx->src0_spad.data = octx->ctx->vtcm_base;
octx->src0_spad.size = total_spad_size;
octx->src0_spad.size_per_thread = spad_per_thread;
FARF(HIGH, "argsort: %ux%ux%ux%u -> %ux%ux%ux%u (0x%x, 0x%x)",
octx->src0.ne[0], octx->src0.ne[1], octx->src0.ne[2], octx->src0.ne[3],
octx->dst.ne[0], octx->dst.ne[1], octx->dst.ne[2], octx->dst.ne[3],
octx->src0.data, octx->dst.data);
uint32_t total_rows = octx->src0.ne[1] * octx->src0.ne[2] * octx->src0.ne[3];
uint32_t n_jobs = MIN(total_rows, octx->n_threads);
struct htp_argsort_context actx;
actx.octx = octx;
actx.nrows_per_thread = (total_rows + n_jobs - 1) / n_jobs;
// Run jobs
worker_pool_run_func(octx->ctx->worker_pool, htp_argsort_f32, &actx, n_jobs);
return HTP_STATUS_OK;
}

View File

@ -64,6 +64,7 @@ enum htp_op {
HTP_OP_SCALE = 16,
HTP_OP_GET_ROWS = 17,
HTP_OP_CPY = 18,
HTP_OP_ARGSORT = 19,
INVALID
};
@ -103,22 +104,6 @@ static inline size_t htp_type_nbytes(uint32_t t) {
return 0;
}
static const char * htp_type_name(uint32_t t) {
switch (t) {
case HTP_TYPE_F32:
return "fp32";
case HTP_TYPE_F16:
return "fp16";
case HTP_TYPE_Q4_0:
return "q4_0";
case HTP_TYPE_Q8_0:
return "q8_0";
case HTP_TYPE_MXFP4:
return "mxfp4";
}
return 0;
}
// Internal types
#define QK_Q4_0x4x2 256 // 4x Q4_0 blocks packed with next 4x Q4_0 blocks (size in bytes 128)
#define QK_Q8_0x4x2 256 // 4x Q8_0 blocks concat with next 4x Q8_0 blocks

View File

@ -98,5 +98,6 @@ int op_flash_attn_ext(struct htp_ops_context * octx);
int op_set_rows(struct htp_ops_context * octx);
int op_get_rows(struct htp_ops_context * octx);
int op_cpy(struct htp_ops_context * octx);
int op_argsort(struct htp_ops_context * octx);
#endif /* HTP_OPS_H */

View File

@ -66,6 +66,12 @@ static inline float hvx_vec_get_f32(HVX_Vector v) {
return x;
}
static inline int32_t hvx_vec_get_i32(HVX_Vector v) {
int32_t __attribute__((aligned(128))) x;
hvx_vec_store_a(&x, 4, v);
return x;
}
static inline HVX_Vector hvx_vec_abs_f16(HVX_Vector v) {
// abs by clearing the fp16 sign bit
HVX_Vector mask = Q6_Vh_vsplat_R(0x7fff);

View File

@ -136,8 +136,6 @@ static inline void hvx_copy_f32_uu(uint8_t * restrict dst, const uint8_t * restr
dst_type * restrict vdst = (dst_type *) dst; \
src_type * restrict vsrc = (src_type *) src; \
\
const HVX_Vector zero = Q6_V_vsplat_R(0); \
\
const uint32_t elem_size = sizeof(__fp16); \
const uint32_t epv = 128 / elem_size; \
const uint32_t nvec = n / epv; \

View File

@ -440,6 +440,45 @@ static void proc_matmul_req(struct htp_context * ctx,
send_htp_rsp(ctx, req->op, rsp_status, rsp_bufs, 1, &prof);
}
static void proc_argsort_req(struct htp_context * ctx, struct htp_general_req * req, struct dspqueue_buffer * bufs) {
struct dspqueue_buffer rsp_bufs[1];
// We had written to the output buffer, we'd also need to flush it
rsp_bufs[0].fd = bufs[1].fd;
rsp_bufs[0].ptr = bufs[1].ptr;
rsp_bufs[0].offset = bufs[1].offset;
rsp_bufs[0].size = bufs[1].size;
rsp_bufs[0].flags = (DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | // Flush HTP
DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT); // Invalidate CPU
// Setup Op context
struct htp_ops_context octx = { 0 };
octx.ctx = ctx;
octx.src0 = req->src0;
octx.dst = req->dst;
octx.flags = req->flags;
octx.op = req->op;
memcpy(octx.op_params, req->op_params, sizeof(octx.op_params));
// Update data pointers
octx.src0.data = (uint32_t) bufs[0].ptr;
octx.dst.data = (uint32_t) bufs[1].ptr;
octx.n_threads = ctx->n_threads;
struct profile_data prof;
profile_start(&prof);
uint32_t rsp_status = HTP_STATUS_INTERNAL_ERR;
if (vtcm_acquire(ctx) == AEE_SUCCESS) {
rsp_status = op_argsort(&octx);
vtcm_release(ctx);
}
profile_stop(&prof);
send_htp_rsp(ctx, req->op, rsp_status, rsp_bufs, 1, &prof);
}
static void proc_cpy_req(struct htp_context * ctx, struct htp_general_req * req, struct dspqueue_buffer * bufs) {
struct dspqueue_buffer rsp_bufs[1];
@ -1035,6 +1074,14 @@ static void htp_packet_callback(dspqueue_t queue, int error, void * context) {
proc_cpy_req(ctx, &req, bufs);
break;
case HTP_OP_ARGSORT:
if (n_bufs != 2) {
FARF(ERROR, "Bad argsort-req buffer list");
continue;
}
proc_argsort_req(ctx, &req, bufs);
break;
default:
FARF(ERROR, "Unknown Op %u", req.op);
break;