Skip to content

Commit 67776ea

Browse files
max-krasnyanskymax-krasnyansky
authored
hexagon: improve ARGSORT performance for small tensors (ggml-org#25512)
* hex-sort: add efficient bitomic sort in hvx regs up to 1024 elements * hex-sort: fix inverted vrors * hex-sort: specialize sort functions for the common cases * hex-sort: add tracing and local context
1 parent 22b69b6 commit 67776ea

2 files changed

Lines changed: 241 additions & 21 deletions

File tree

‎ggml/src/ggml-hexagon/htp/CMakeLists.txt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,16 +31,16 @@ add_library(${HTP_LIB} SHARED
3131
get-rows-ops.c
3232
cpy-ops.c
3333
repeat-ops.c
34-
argsort-ops.c
3534
ssm-conv.c
3635
cumsum-ops.c
3736
fill-ops.c
3837
concat-ops.c
3938
diag-ops.c
4039
solve-tri-ops.c
4140
pad-ops.c
42-
matmul-ops.c
4341
flash-attn-ops.c
42+
matmul-ops.c
43+
argsort-ops.c
4444
)
4545

4646
target_compile_definitions(${HTP_LIB} PRIVATE

‎ggml/src/ggml-hexagon/htp/argsort-ops.c‎

Lines changed: 239 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
struct htp_argsort_context {
2323
struct htp_ops_context * octx;
2424
uint32_t nrows_per_thread;
25+
uint8_t * vtcm_base;
26+
size_t vtcm_per_thread;
2527
};
2628

2729
static inline bool all_greater_f32(HVX_Vector x, HVX_Vector y)
@@ -170,7 +172,208 @@ int32_t argosrt_ramp_lut[32] __attribute__((aligned(VLEN))) = {
170172
16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
171173
};
172174

173-
static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
175+
__attribute__((always_inline))
176+
static inline void vec_cas(HVX_Vector * X_val, HVX_Vector * X_idx, HVX_Vector * Y_val, HVX_Vector * Y_idx, bool asc) {
177+
HVX_VectorPred pred = asc ? Q6_Q_vcmp_gt_VsfVsf(*X_val, *Y_val)
178+
: Q6_Q_vcmp_gt_VsfVsf(*Y_val, *X_val);
179+
HVX_Vector next_X_val = Q6_V_vmux_QVV(pred, *Y_val, *X_val);
180+
HVX_Vector next_Y_val = Q6_V_vmux_QVV(pred, *X_val, *Y_val);
181+
HVX_Vector next_X_idx = Q6_V_vmux_QVV(pred, *Y_idx, *X_idx);
182+
HVX_Vector Y_tmp_idx = Q6_V_vmux_QVV(pred, *X_idx, *Y_idx);
183+
*X_val = next_X_val;
184+
*Y_val = next_Y_val;
185+
*X_idx = next_X_idx;
186+
*Y_idx = Y_tmp_idx;
187+
}
188+
189+
__attribute__((always_inline))
190+
static inline void bitonic_cas_32(HVX_Vector * V, HVX_Vector * I, int d, HVX_VectorPred dir_mask, HVX_Vector idx_vec, HVX_Vector zero_vec) {
191+
HVX_VectorPred mask_left;
192+
HVX_Vector V_rot_left, V_rot_right;
193+
HVX_Vector I_rot_left, I_rot_right;
194+
195+
if (d == 1) {
196+
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1)), zero_vec);
197+
V_rot_left = Q6_V_vror_VR(*V, 4);
198+
V_rot_right = Q6_V_vror_VR(*V, 124);
199+
I_rot_left = Q6_V_vror_VR(*I, 4);
200+
I_rot_right = Q6_V_vror_VR(*I, 124);
201+
} else if (d == 2) {
202+
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(2)), zero_vec);
203+
V_rot_left = Q6_V_vror_VR(*V, 8);
204+
V_rot_right = Q6_V_vror_VR(*V, 120);
205+
I_rot_left = Q6_V_vror_VR(*I, 8);
206+
I_rot_right = Q6_V_vror_VR(*I, 120);
207+
} else if (d == 4) {
208+
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(4)), zero_vec);
209+
V_rot_left = Q6_V_vror_VR(*V, 16);
210+
V_rot_right = Q6_V_vror_VR(*V, 112);
211+
I_rot_left = Q6_V_vror_VR(*I, 16);
212+
I_rot_right = Q6_V_vror_VR(*I, 112);
213+
} else if (d == 8) {
214+
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(8)), zero_vec);
215+
V_rot_left = Q6_V_vror_VR(*V, 32);
216+
V_rot_right = Q6_V_vror_VR(*V, 96);
217+
I_rot_left = Q6_V_vror_VR(*I, 32);
218+
I_rot_right = Q6_V_vror_VR(*I, 96);
219+
} else { // d == 16
220+
mask_left = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(16)), zero_vec);
221+
V_rot_left = Q6_V_vror_VR(*V, 64);
222+
V_rot_right = Q6_V_vror_VR(*V, 64);
223+
I_rot_left = Q6_V_vror_VR(*I, 64);
224+
I_rot_right = Q6_V_vror_VR(*I, 64);
225+
}
226+
227+
HVX_Vector V_paired = Q6_V_vmux_QVV(mask_left, V_rot_left, V_rot_right);
228+
HVX_Vector I_paired = Q6_V_vmux_QVV(mask_left, I_rot_left, I_rot_right);
229+
230+
HVX_VectorPred V_gt_Vpaired = Q6_Q_vcmp_gt_VsfVsf(*V, V_paired);
231+
HVX_VectorPred Vpaired_gt_V = Q6_Q_vcmp_gt_VsfVsf(V_paired, *V);
232+
HVX_VectorPred mask_right = Q6_Q_not_Q(mask_left);
233+
HVX_VectorPred Q_asc = Q6_Q_or_QQ(
234+
Q6_Q_and_QQ(mask_left, V_gt_Vpaired),
235+
Q6_Q_and_QQ(Vpaired_gt_V, mask_right)
236+
);
237+
HVX_VectorPred Q_swap = Q6_Q_or_QQ(
238+
Q6_Q_and_QQ(dir_mask, Q_asc),
239+
Q6_Q_and_QQ(Q6_Q_not_Q(dir_mask), Q6_Q_not_Q(Q_asc))
240+
);
241+
242+
*V = Q6_V_vmux_QVV(Q_swap, V_paired, *V);
243+
*I = Q6_V_vmux_QVV(Q_swap, I_paired, *I);
244+
}
245+
246+
__attribute__((always_inline))
247+
static inline void bitonic_sort_generic_hvx(uint8_t * values, uint8_t * indices, int K, bool asc_order) {
248+
HVX_Vector V[32];
249+
HVX_Vector I[32];
250+
251+
HVX_Vector zero_vec = Q6_V_vzero();
252+
HVX_Vector idx_vec = *(HVX_Vector *)argosrt_ramp_lut;
253+
254+
// Load values and initialize indices
255+
for (int v = 0; v < K; v++) {
256+
V[v] = *(HVX_Vector *)(values + v * 128);
257+
I[v] = Q6_Vw_vadd_VwVw(idx_vec, Q6_V_vsplat_R(v * 32));
258+
}
259+
260+
HVX_VectorPred pred_all_1s = Q6_Q_vcmp_eq_VwVw(zero_vec, zero_vec);
261+
HVX_VectorPred pred_all_0s = Q6_Q_not_Q(pred_all_1s);
262+
263+
int M = 5;
264+
while ((1 << (M - 5)) < K) M++;
265+
266+
for (int s = 1; s <= M; s++) {
267+
for (int stage_d = s - 1; stage_d >= 0; stage_d--) {
268+
int d = 1 << stage_d;
269+
if (d >= 32) {
270+
int v_dist = d / 32;
271+
for (int v1 = 0; v1 < K; v1++) {
272+
if ((v1 & v_dist) == 0) {
273+
int v2 = v1 + v_dist;
274+
bool asc = (s < M) ? ((((v1 * 32) >> s) % 2) == 0) : asc_order;
275+
vec_cas(&V[v1], &I[v1], &V[v2], &I[v2], asc);
276+
}
277+
}
278+
} else {
279+
if (s < 5) {
280+
HVX_VectorPred dir_mask = Q6_Q_vcmp_eq_VwVw(Q6_V_vand_VV(idx_vec, Q6_V_vsplat_R(1 << s)), zero_vec);
281+
for (int v = 0; v < K; v++) {
282+
bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
283+
}
284+
} else {
285+
for (int v = 0; v < K; v++) {
286+
bool asc = (s < M) ? ((((v * 32) >> s) % 2) == 0) : asc_order;
287+
HVX_VectorPred dir_mask = asc ? pred_all_1s : pred_all_0s;
288+
bitonic_cas_32(&V[v], &I[v], d, dir_mask, idx_vec, zero_vec);
289+
}
290+
}
291+
}
292+
}
293+
}
294+
295+
// Write back sorted values and indices
296+
for (int v = 0; v < K; v++) {
297+
*(HVX_Vector *)(values + v * 128) = V[v];
298+
*(HVX_Vector *)(indices + v * 128) = I[v];
299+
}
300+
}
301+
302+
__attribute__((always_inline))
303+
static inline void sort32_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
304+
bitonic_sort_generic_hvx(values, indices, 1, order == GGML_SORT_ORDER_ASC);
305+
}
306+
307+
__attribute__((always_inline))
308+
static inline void sort64_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
309+
bitonic_sort_generic_hvx(values, indices, 2, order == GGML_SORT_ORDER_ASC);
310+
}
311+
312+
__attribute__((always_inline))
313+
static inline void sort128_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
314+
bitonic_sort_generic_hvx(values, indices, 4, order == GGML_SORT_ORDER_ASC);
315+
}
316+
317+
__attribute__((always_inline))
318+
static inline void sort256_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
319+
bitonic_sort_generic_hvx(values, indices, 8, order == GGML_SORT_ORDER_ASC);
320+
}
321+
322+
__attribute__((always_inline))
323+
static inline void sort512_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
324+
bitonic_sort_generic_hvx(values, indices, 16, order == GGML_SORT_ORDER_ASC);
325+
}
326+
327+
__attribute__((always_inline))
328+
static inline void sort1024_f32_hvx(uint8_t * values, uint8_t * indices, enum ggml_sort_order order) {
329+
bitonic_sort_generic_hvx(values, indices, 32, order == GGML_SORT_ORDER_ASC);
330+
}
331+
332+
#define HTP_ARGSORT_FN(ne00, order_name, order_enum, sort_fn) \
333+
static void htp_argsort_f32_##ne00##_##order_name(unsigned int n, unsigned int i, void * data) { \
334+
struct htp_argsort_context * actx = (struct htp_argsort_context *)data; \
335+
struct htp_ops_context * octx = actx->octx; \
336+
const struct htp_tensor * src0 = octx->src[0]; \
337+
const struct htp_tensor * dst = octx->dst; \
338+
uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i; \
339+
uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3]; \
340+
uint32_t rows_per_thread = actx->nrows_per_thread; \
341+
uint32_t start_row = rows_per_thread * i; \
342+
uint32_t end_row = MIN(start_row + rows_per_thread, total_rows); \
343+
size_t values_size = hex_round_up(ne00 * sizeof(float), 128); \
344+
float * values_buf = (float *) spad; \
345+
int32_t * indices_buf = (int32_t *) (spad + values_size); \
346+
uint32_t nb01 = src0->nb[1]; \
347+
uint32_t nb1 = dst->nb[1]; \
348+
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL; \
349+
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row); \
350+
for (uint32_t r = start_row; r < end_row; r++) { \
351+
uint32_t src_offset = r * nb01; \
352+
uint32_t dst_offset = r * nb1; \
353+
uint8_t * src_ptr = (uint8_t *) src0->data + src_offset; \
354+
uint8_t * dst_ptr = (uint8_t *) dst->data + dst_offset; \
355+
hex_l2fetch(src_ptr, ne00 * sizeof(float), ne00 * sizeof(float), 1); \
356+
hvx_copy_f32_au((uint8_t*)values_buf, src_ptr, ne00); \
357+
sort_fn((uint8_t*)values_buf, (uint8_t*)indices_buf, order_enum); \
358+
hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00); \
359+
} \
360+
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row); \
361+
}
362+
363+
HTP_ARGSORT_FN(32, asc, GGML_SORT_ORDER_ASC, sort32_f32_hvx)
364+
HTP_ARGSORT_FN(32, dsc, GGML_SORT_ORDER_DESC, sort32_f32_hvx)
365+
HTP_ARGSORT_FN(64, asc, GGML_SORT_ORDER_ASC, sort64_f32_hvx)
366+
HTP_ARGSORT_FN(64, dsc, GGML_SORT_ORDER_DESC, sort64_f32_hvx)
367+
HTP_ARGSORT_FN(128, asc, GGML_SORT_ORDER_ASC, sort128_f32_hvx)
368+
HTP_ARGSORT_FN(128, dsc, GGML_SORT_ORDER_DESC, sort128_f32_hvx)
369+
HTP_ARGSORT_FN(256, asc, GGML_SORT_ORDER_ASC, sort256_f32_hvx)
370+
HTP_ARGSORT_FN(256, dsc, GGML_SORT_ORDER_DESC, sort256_f32_hvx)
371+
HTP_ARGSORT_FN(512, asc, GGML_SORT_ORDER_ASC, sort512_f32_hvx)
372+
HTP_ARGSORT_FN(512, dsc, GGML_SORT_ORDER_DESC, sort512_f32_hvx)
373+
HTP_ARGSORT_FN(1024, asc, GGML_SORT_ORDER_ASC, sort1024_f32_hvx)
374+
HTP_ARGSORT_FN(1024, dsc, GGML_SORT_ORDER_DESC, sort1024_f32_hvx)
375+
376+
static void htp_argsort_f32_fallback(unsigned int n, unsigned int i, void * data) {
174377
struct htp_argsort_context * actx = (struct htp_argsort_context *)data;
175378
struct htp_ops_context * octx = actx->octx;
176379

@@ -179,7 +382,7 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
179382
const struct htp_tensor * dst = octx->dst;
180383

181384
// Scratchpad memory
182-
uint8_t * spad = octx->src0_spad.data + octx->src0_spad.size_per_thread * i;
385+
uint8_t * spad = actx->vtcm_base + actx->vtcm_per_thread * i;
183386

184387
// Dimensions
185388
uint32_t ne00 = src0->ne[0];
@@ -188,12 +391,8 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
188391
uint32_t ne03 = src0->ne[3];
189392

190393
uint32_t nb01 = src0->nb[1];
191-
//uint32_t nb02 = src0->nb[2];
192-
//uint32_t nb03 = src0->nb[3];
193394

194395
uint32_t nb1 = dst->nb[1];
195-
//uint32_t nb2 = dst->nb[2];
196-
//uint32_t nb3 = dst->nb[3];
197396

198397
// Sort order
199398
enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
@@ -204,20 +403,17 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
204403
uint32_t start_row = rows_per_thread * i;
205404
uint32_t end_row = MIN(start_row + rows_per_thread, total_rows);
206405

207-
// Scratchpad layout:
208-
// We need space for one row of float data (values) and one row of int32 indices.
209-
// values: ne00 * sizeof(float)
210-
// indices: ne00 * sizeof(int32_t)
211-
// Padded to 128 bytes.
212-
213406
size_t values_size = hex_round_up(ne00 * sizeof(float), 128);
214-
size_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
407+
uint32_t num_vec_ind_values = hmx_ceil_div(ne00, VLEN/(sizeof(int32_t)));
215408
float * values_buf = (float *) spad;
216409
int32_t * indices_buf = (int32_t *) (spad + values_size);
217410
HVX_Vector * indices_buf_vec = (HVX_Vector *) (spad + values_size);
218411
const HVX_Vector ind_init_vec = *(HVX_Vector *)argosrt_ramp_lut;
219412
const HVX_Vector ind_diff_vec = Q6_V_vsplat_R(32);
220413

414+
struct htp_thread_trace * tr = octx->ctx ? &octx->ctx->trace[i] : NULL;
415+
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
416+
221417
for (uint32_t r = start_row; r < end_row; r++) {
222418
uint32_t src_offset = r * nb01;
223419
uint32_t dst_offset = r * nb1;
@@ -245,6 +441,8 @@ static void htp_argsort_f32(unsigned int n, unsigned int i, void * data) {
245441
// Copy indices back to DDR
246442
hvx_copy_f32_ua(dst_ptr, (const uint8_t *) indices_buf, ne00);
247443
}
444+
445+
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, start_row);
248446
}
249447

250448
int op_argsort(struct htp_ops_context * octx) {
@@ -273,11 +471,6 @@ int op_argsort(struct htp_ops_context * octx) {
273471
return HTP_STATUS_VTCM_TOO_SMALL;
274472
}
275473

276-
octx->src0_spad.data = octx->ctx->vtcm_base;
277-
octx->src0_spad.size = total_spad_size;
278-
octx->src0_spad.size_per_thread = spad_per_thread;
279-
octx->src0_spad.src = NULL;
280-
281474
FARF(HIGH, "argsort: %ux%ux%ux%u -> %ux%ux%ux%u (0x%x, 0x%x)",
282475
octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
283476
octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
@@ -286,9 +479,36 @@ int op_argsort(struct htp_ops_context * octx) {
286479
struct htp_argsort_context actx;
287480
actx.octx = octx;
288481
actx.nrows_per_thread = (total_rows + n_threads - 1) / n_threads;
482+
actx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
483+
actx.vtcm_per_thread = spad_per_thread;
484+
485+
enum ggml_sort_order order = (enum ggml_sort_order) octx->op_params[0];
486+
worker_callback_t job_func = htp_argsort_f32_fallback;
487+
488+
if (order == GGML_SORT_ORDER_ASC) {
489+
switch (ne00) {
490+
case 1024: job_func = htp_argsort_f32_1024_asc; break;
491+
case 512: job_func = htp_argsort_f32_512_asc; break;
492+
case 256: job_func = htp_argsort_f32_256_asc; break;
493+
case 128: job_func = htp_argsort_f32_128_asc; break;
494+
case 64: job_func = htp_argsort_f32_64_asc; break;
495+
case 32: job_func = htp_argsort_f32_32_asc; break;
496+
default: job_func = htp_argsort_f32_fallback; break;
497+
}
498+
} else {
499+
switch (ne00) {
500+
case 1024: job_func = htp_argsort_f32_1024_dsc; break;
501+
case 512: job_func = htp_argsort_f32_512_dsc; break;
502+
case 256: job_func = htp_argsort_f32_256_dsc; break;
503+
case 128: job_func = htp_argsort_f32_128_dsc; break;
504+
case 64: job_func = htp_argsort_f32_64_dsc; break;
505+
case 32: job_func = htp_argsort_f32_32_dsc; break;
506+
default: job_func = htp_argsort_f32_fallback; break;
507+
}
508+
}
289509

290510
// Run jobs
291-
worker_pool_run_func(octx->ctx->worker_pool, htp_argsort_f32, &actx, n_threads);
511+
worker_pool_run_func(octx->ctx->worker_pool, job_func, &actx, n_threads);
292512

293513
return HTP_STATUS_OK;
294514
}

0 commit comments

Comments
 (0)