ClaudeBot/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-17 20:31:47 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
eafe15a5e3d87dd68ae33acf6a7cbd9415a0ac5e
llama.cpp/ggml/src/ggml-hexagon/htp/allreduce-ops.c
T
Max KrasnyanskyandAlexander Lu eafe15a5e3 hexagon: support for multi-device model split (aka row-split) (#28589) 
* hex-row-split: add support for multi-device row spliting

Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>

* hex-mdev: add work splitting to fused kernels

* hex-mdev: use mdev_ prefix for all multi-device state

* hex-mdev: make device configuration more expressive to support device groups

* hex-mdev: fix mdev session init

* hex-mdev: fused nx (2x,3x) matmuls must update row counts for each w/o

* hex-mdev: fix MUL_MAT work partitioning bugs introduced by mdev

* hex-cont: fix crashes with new tests due to wrong striding

* hex-mdev: move fences after l2flushes

* hex-cont: fix work splitting for mnpu -- align chunks to cachelines

* hex-mdev: fix CPY tests with multi-dev

* hex-mmid: fix work partitioning with mnpu

* hex-mm: fix test failures with mdev

* hex-binary: fix work partitioning for mdev

* hex-argsort: fix mdev partitioning

* hex-mdev: fix work partitioning and general updates for all simple ops

* hex-fa: fix mdev work splitting issues

* hex-mdev: fixing more failing ops test

* hex-mdev: update the rest of the ops

* hex-mdev: refactor all mdev splitting logic to be contained within if (mdev_count > 1) {...}

* hex-mdev: fix macros

* hex-mdev: simplify session flush logic

* hex-sync: fix recursion in session flush

* hex-mdev: factor out fence buffer and allocator

* hex-fence: make fence allocation more robust with reserved slots for mdev

* hex-mdev: keep all mdev state in htp_mdev_group

* hex-mdev: further cleanup mdev group handling at the host

* hex-mdev: update group idx in the opbatch before serializing

* hex-batch: remove separate op_pending and use batch_req/rsp_seq

* hex-async: workaround another missing tensor_init in ggml-meta

* hex-fence: cleanup and robustify fences and error handling in multi-device scenarios

* hex-ar: improve ALLREDUCE error handling

* hex-async: robust error handling for op_cpy_fence

* hex-async: use seq0 from allreduce context to allocate fence_seq

* hex-mdev: fix remaining issues with fence and barrier clearing in CPY_FENCE

* hex-misc: realign macros and fix misplaces trace events

* hex-misc: align macros

* hex-mdev: fix unclone buffer re-entrancy

* hex-glu: fix mdev partitioning logic

* hex-mdev: make buffer uncloning/cleanup work with tensor-split scenarios

* hex-mdev: tighten up the can_split check in act-ops

* hex-mdev: factor out common bits of the partitioning logic

* hex-mm: minor realignment of the macros

* hex-bufs: fix incorrectly placed assert for MAX_BUFS

* hex-pad: tighten up gating checks for PAD

* hex-kparams: make sure all kernels properly use kparams->n_threads

* hex-docs: update user and developer docs with new features and detailed guide for ops development

* hex-scripts: update run script to properly parse dev groups

* hex-misc: formatting

* hex-sess: minor cleanup for session init

* hex-ar: fix vtcm size calc in allreduce kparams

* hex-scripts: fix flake8 warnings

* hex-rope: update ROPE to support mdev work split

* hex-ops: remove redunant checks and minor reformat

* hex-dev-guide: update dev-guide to avoid redundant null checks

* hex-async: improve event_wait, event_sync and fence implementations

* hex-async: remove synchronous flush from event_sync

* hex-async: symplify fence recovery protocol and make sync more robust

* hex-async: futher simplify error recovery for fences

* hex-err: return status instead of just -1

* hex-async: print all seq nums in hex

* hex-async: make sure fences flush dirty ranges

* hex-async: add dirty ranges merging to reduce fence flushes

* hex-async: properly sync before freeing the event

* hex-async: make sure fence owner session is not overriden

* hex-async: more fence write order more robust

* hex-async: make sure not to fuse ALLREDUCE+ADD if their dsts overlap

* hex-fusion: cleanup redundant checks

---------

Co-authored-by: Alexander Lu <alexlu@qti.qualcomm.com>
2026-09-11 20:46:51 -07:00

458 lines
36 KiB
C
Raw Blame History

#pragma clang diagnostic ignored "-Wunused-variable"
#pragma clang diagnostic ignored "-Wunused-function"
#pragma clang diagnostic ignored "-Wunused-but-set-variable"
#include <HAP_farf.h>
#include <HAP_perf.h>
#include <stdatomic.h>
#include <math.h>
#include <string.h>
#define GGML_COMMON_DECL_C
#include "ggml-common.h"
#include "htp-ctx.h"
#include "htp-ops.h"
#include "hvx-utils.h"
#include "htp-tensor.h"
#include "hex-dma.h"
#include "hex-profile.h"
#include "allreduce-ops.h"
#include "htp-fence.h"
struct htp_allreduce_context {
struct htp_ops_context * octx;
uint32_t n_ranks;
uint32_t n_dsts;
uint32_t nelem;
uint32_t ne0;
uint32_t ne1;
uint32_t row_size_aligned;
uint32_t rank_elem_start;
uint32_t rank_nelem;
uint32_t elems_per_thread;
uint32_t block_elems;
uint32_t vtcm_size_per_thread;
bool is_row_bcast;
uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];
uint8_t * dst_spad_base;
uint8_t * res_spad_base;
};
#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD) \
static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \
struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \
struct htp_ops_context * octx = actx->octx; \
\
const uint32_t n_ranks = actx->n_ranks; \
const uint32_t n_dsts = actx->n_dsts; \
const uint32_t block_elems = actx->block_elems; \
\
const uint32_t dr = actx->elems_per_thread; \
const uint32_t ir0 = actx->rank_elem_start + dr * ith; \
const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem); \
if (ir0 >= ir1) return; \
\
struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
dma_queue * q = octx->ctx->dma[ith]; \
\
uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \
for (uint32_t s = 0; s < n_ranks; s++) { \
src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \
} \
uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \
uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL; \
\
const size_t spad_half = actx->vtcm_size_per_thread / 2; \
uint32_t ir_prefetch = ir0; \
int spad_idx = 0; \
\
for (int k = 0; k < 2 && ir_prefetch < ir1; k++) { \
uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch); \
size_t cur_bytes = cur_elems * sizeof(TYPE); \
uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \
for (uint32_t d = 0; d < n_dsts; d++) { \
uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0); \
} \
for (uint32_t s = 0; s < n_ranks; s++) { \
uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \
const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \
} \
if (HAS_ADD) { \
uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \
const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \
} \
ir_prefetch += cur_elems; \
spad_idx ^= 1; \
} \
\
for (uint32_t ir = ir0; ir < ir1; ) { \
uint32_t cur_elems = MIN(block_elems, ir1 - ir); \
size_t cur_bytes = cur_elems * sizeof(TYPE); \
uint8_t * d_spad = NULL; \
for (uint32_t d = 0; d < n_dsts; d++) { \
d_spad = (uint8_t *) dma_queue_pop(q).src; \
} \
uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \
for (uint32_t s = 0; s < n_ranks; s++) { \
s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \
} \
uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL; \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \
HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems); \
for (uint32_t s = 2; s < n_ranks; s++) { \
HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems); \
} \
if (HAS_ADD) { \
HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems); \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \
for (uint32_t d = 0; d < n_dsts; d++) { \
uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1); \
} \
if (ir_prefetch < ir1) { \
uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch); \
size_t next_bytes = next_elems * sizeof(TYPE); \
for (uint32_t s = 0; s < n_ranks; s++) { \
const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \
} \
if (HAS_ADD) { \
const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \
dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1); \
} \
ir_prefetch += next_elems; \
} \
ir += cur_elems; \
} \
dma_queue_flush(q); \
}
DEFINE_ALLREDUCE_THREAD_DMA_1D(f16, __fp16, hvx_add_f16_aaa, 0)
DEFINE_ALLREDUCE_THREAD_DMA_1D(f32, float, hvx_add_f32_aaa, 0)
DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1)
DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float, hvx_add_f32_aaa, 1)
#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST) \
static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \
struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \
struct htp_ops_context * octx = actx->octx; \
\
const uint32_t n_ranks = actx->n_ranks; \
const uint32_t n_dsts = actx->n_dsts; \
const uint32_t ne0 = actx->ne0; \
const uint32_t block_rows = actx->block_elems; \
const uint32_t row_size_aligned = actx->row_size_aligned; \
const uint32_t row_bytes = ne0 * sizeof(TYPE); \
\
const uint32_t dr = actx->elems_per_thread; \
const uint32_t r0 = actx->rank_elem_start + dr * ith; \
const uint32_t r1 = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem); \
if (r0 >= r1) return; \
\
struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \
dma_queue * q = octx->ctx->dma[ith]; \
\
uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \
for (uint32_t s = 0; s < n_ranks; s++) { \
src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \
} \
uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \
uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \
\
const size_t spad_half = actx->vtcm_size_per_thread / 2; \
uint32_t r_prefetch = r0; \
int spad_idx = 0; \
\
for (int k = 0; k < 2 && r_prefetch < r1; k++) { \
uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch); \
uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \
for (uint32_t d = 0; d < n_dsts; d++) { \
uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1]; \
dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0); \
} \
for (uint32_t s = 0; s < n_ranks; s++) { \
uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \
const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \
dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows); \
} \
if (HAS_ADD && !IS_ROW_BCAST) { \
uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \
const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \
dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows); \
} \
r_prefetch += cur_rows; \
spad_idx ^= 1; \
} \
\
for (uint32_t r = r0; r < r1; ) { \
uint32_t cur_rows = MIN(block_rows, r1 - r); \
uint8_t * d_spad = NULL; \
for (uint32_t d = 0; d < n_dsts; d++) { \
d_spad = (uint8_t *) dma_queue_pop(q).src; \
} \
uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \
for (uint32_t s = 0; s < n_ranks; s++) { \
s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \
} \
uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL; \
htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \
for (uint32_t row = 0; row < cur_rows; row++) { \
uint8_t * d_row = d_spad + row * row_size_aligned; \
const uint8_t * s0_row = s_spad[0] + row * row_size_aligned; \
const uint8_t * s1_row = s_spad[1] + row * row_size_aligned; \
HVX_ADD_FN(d_row, s0_row, s1_row, ne0); \
for (uint32_t s = 2; s < n_ranks; s++) { \
const uint8_t * ss_row = s_spad[s] + row * row_size_aligned; \
HVX_ADD_FN(d_row, d_row, ss_row, ne0); \
} \
if (HAS_ADD) { \
const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned); \
HVX_ADD_FN(d_row, d_row, res_row, ne0); \
} \
} \
htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \
for (uint32_t d = 0; d < n_dsts; d++) { \
uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1]; \
dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows); \
} \
if (r_prefetch < r1) { \
uint32_t next_rows = MIN(block_rows, r1 - r_prefetch); \
for (uint32_t s = 0; s < n_ranks; s++) { \
const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \
dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows); \
} \
if (HAS_ADD && !IS_ROW_BCAST) { \
const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \
dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \
} \
r_prefetch += next_rows; \
} \
r += cur_rows; \
} \
dma_queue_flush(q); \
}
DEFINE_ALLREDUCE_THREAD_DMA_2D(f16, __fp16, hvx_add_f16_aaa, 0, 0)
DEFINE_ALLREDUCE_THREAD_DMA_2D(f32, float, hvx_add_f32_aaa, 0, 0)
DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16, __fp16, hvx_add_f16_aaa, 1, 0)
DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32, float, hvx_add_f32_aaa, 1, 0)
DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1)
DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float, hvx_add_f32_aaa, 1, 1)
static int validate_allreduce(
struct htp_ops_context * octx,
const struct htp_allreduce_kernel_params * kparams,
uint32_t n_ranks
) {
if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
return HTP_STATUS_INVAL_PARAMS;
}
if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
return HTP_STATUS_INVAL_PARAMS;
}
const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
const size_t n_vtcm_buffers = htp_allreduce_vtcm_buffer_count(
n_ranks, octx->n_threads, has_add, kparams->is_row_bcast != 0);
const size_t vtcm_size = n_vtcm_buffers * (size_t) kparams->vtcm_size_per_thread;
if (vtcm_size != (size_t) kparams->vtcm_size) {
return HTP_STATUS_INVAL_PARAMS;
}
if (vtcm_size > octx->ctx->vtcm_size) {
return HTP_STATUS_VTCM_TOO_SMALL;
}
if (octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_F32) {
return HTP_STATUS_NO_SUPPORT;
}
return HTP_STATUS_OK;
}
int op_allreduce(struct htp_ops_context * octx) {
if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
return HTP_STATUS_OK;
}
const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params;
const struct htp_tensor * dst = octx->dst;
const uint32_t rank = (uint32_t) kparams->rank;
const uint32_t n_ranks = (uint32_t) kparams->n_ranks;
if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) {
return HTP_STATUS_INVAL_PARAMS;
}
const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0];
const uint32_t fence_seq_exit = (uint32_t) octx->op_params[1];
const struct htp_tensor * my_sync = octx->src[n_ranks + rank];
atomic_uint * my_fence = (atomic_uint *) (uintptr_t) my_sync->data;
const int status = validate_allreduce(octx, kparams, n_ranks);
if (status != HTP_STATUS_OK) {
if (status == HTP_STATUS_NO_SUPPORT) {
FARF(ERROR, "ggml-hex: allreduce unsupported type %d : rank %u\n", dst->type, rank);
}
htp_fence_write(my_fence, fence_seq_exit, status);
return status;
}
const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
// 1. Entry Barrier: Synchronize all ranks before reading
struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
htp_fence_write(my_fence, fence_seq_entry, octx->status);
for (uint32_t j = 0; j < n_ranks; j++) {
if (j == rank) continue;
const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
uint64_t spins = 0;
while (1) {
uint32_t peer_seq;
uint32_t peer_status;
htp_fence_read(peer_fence, &peer_seq, &peer_status);
if ((int32_t)(peer_seq - fence_seq_entry) >= 0) {
if (peer_status > HTP_STATUS_OK) {
FARF(ERROR, "ggml-hex: allreduce entry peer %u failed with status %u\n", j, peer_status);
htp_fence_write(my_fence, fence_seq_exit, peer_status);
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
return peer_status;
}
break;
}
if (++spins > HTP_FENCE_TIMEOUT) {
FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
rank, j, peer_fence, fence_seq_entry, peer_seq);
htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
return HTP_STATUS_INTERNAL_ERR;
}
hex_pause();
}
}
asm volatile ("syncht" : : : "memory");
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry);
// 2. Multi-threaded Reduction across assigned rank chunk
if (nelem > 0) {
const uint32_t n_threads = (uint32_t) kparams->n_threads;
const uint32_t block_elems = (uint32_t) kparams->block_elems;
const uint32_t elems_per_thread = (uint32_t) kparams->elems_per_thread;
const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread;
struct htp_allreduce_context actx;
actx.octx = octx;
actx.n_ranks = n_ranks;
actx.n_dsts = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks;
actx.nelem = nelem;
actx.ne0 = (uint32_t) kparams->ne0;
actx.ne1 = (uint32_t) kparams->ne1;
actx.row_size_aligned = (uint32_t) kparams->row_size_aligned;
actx.rank_elem_start = (uint32_t) kparams->rank_elem_start;
actx.rank_nelem = (uint32_t) kparams->rank_nelem;
actx.elems_per_thread = elems_per_thread;
actx.block_elems = block_elems;
actx.vtcm_size_per_thread = vtcm_size_per_thread;
actx.is_row_bcast = (kparams->is_row_bcast != 0);
work_queue_func_t reduce_fun = NULL;
switch (kparams->kernel_type) {
case HTP_ALLREDUCE_KERNEL_DMA_1D:
if (has_add) {
reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32;
} else {
reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32;
}
break;
case HTP_ALLREDUCE_KERNEL_DMA_2D:
if (has_add) {
if (kparams->is_row_bcast) {
reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32;
} else {
reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32;
}
} else {
reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32;
}
break;
default:
FARF(ERROR, "ggml-hex: allreduce unsupported kernel %d : rank %u\n", kparams->kernel_type, rank);
htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_NO_SUPPORT);
return HTP_STATUS_NO_SUPPORT;
}
uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base;
for (uint32_t s = 0; s < n_ranks; s++) {
actx.src_spad_base[s] = vtcm_ptr;
vtcm_ptr += n_threads * vtcm_size_per_thread;
}
actx.dst_spad_base = vtcm_ptr;
vtcm_ptr += n_threads * vtcm_size_per_thread;
if (has_add) {
actx.res_spad_base = vtcm_ptr;
vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread;
}
if (has_add && actx.is_row_bcast) {
const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data;
const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float));
dma_queue * q = octx->ctx->dma[0];
dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1);
dma_queue_pop(q);
}
work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads);
}
// 4. Exit Barrier: Synchronize all ranks after writing
htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
htp_fence_write(my_fence, fence_seq_exit, octx->status);
for (uint32_t j = 0; j < n_ranks; j++) {
if (j == rank) continue;
const struct htp_tensor * peer_sync = octx->src[n_ranks + j];
atomic_uint * peer_fence = (atomic_uint *) (uintptr_t) peer_sync->data;
uint64_t spins = 0;
while (1) {
uint32_t peer_seq;
uint32_t peer_status;
htp_fence_read(peer_fence, &peer_seq, &peer_status);
if ((int32_t)(peer_seq - fence_seq_exit) >= 0) {
if (peer_status > HTP_STATUS_OK) {
FARF(ERROR, "ggml-hex: allreduce exit peer %u failed with status %u\n", j, peer_status);
htp_fence_write(my_fence, fence_seq_exit, peer_status);
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
return peer_status;
}
break;
}
if (++spins > HTP_FENCE_TIMEOUT) {
FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT : rank %u waiting on %u fence %p seq 0x%x peer-seq 0x%x\n",
rank, j, peer_fence, fence_seq_exit, peer_seq);
htp_fence_write(my_fence, fence_seq_exit, HTP_STATUS_INTERNAL_ERR);
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
return HTP_STATUS_INTERNAL_ERR;
}
hex_pause();
}
}
asm volatile ("syncht" : : : "memory");
htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);
return octx->status;
}
Reference in New Issue View Git Blame Copy Permalink