Skip to content

Commit

Permalink
threadpool: improve abort handling
Browse files Browse the repository at this point in the history
Do not use threadpool->ec (exit code) to decide whether to exit the compute loop.
threadpool->ec is not atomic which makes thread-sanitizer rightfully unhappy about it.

We already have an atomic threadpool->stop flag for this which is consitent with the
rest of the code.

While at it add an explicit atomic_load for n_threads_cur for consistency.
  • Loading branch information
max-krasnyansky committed Sep 16, 2024
1 parent b9763b3 commit e833ca4
Showing 1 changed file with 11 additions and 11 deletions.
22 changes: 11 additions & 11 deletions ggml/src/ggml.c
Original file line number Diff line number Diff line change
@@ -19928,34 +19928,33 @@ struct ggml_cplan ggml_graph_plan(

static thread_ret_t ggml_graph_compute_thread(void * data) {
struct ggml_compute_state * state = (struct ggml_compute_state *) data;
struct ggml_threadpool * tp = state->threadpool;

const struct ggml_cgraph * cgraph = state->threadpool->cgraph;
const struct ggml_cplan * cplan = state->threadpool->cplan;
const struct ggml_cgraph * cgraph = tp->cgraph;
const struct ggml_cplan * cplan = tp->cplan;

set_numa_thread_affinity(state->ith);

struct ggml_compute_params params = {
/*.ith =*/ state->ith,
/*.nth =*/ state->threadpool->n_threads_cur,
/*.nth =*/ atomic_load_explicit(&tp->n_threads_cur, memory_order_relaxed),
/*.wsize =*/ cplan->work_size,
/*.wdata =*/ cplan->work_data,
/*.threadpool=*/ state->threadpool,
/*.threadpool=*/ tp,
};

for (int node_n = 0; node_n < cgraph->n_nodes; node_n++) {
for (int node_n = 0; node_n < cgraph->n_nodes && !tp->stop; node_n++) {
struct ggml_tensor * node = cgraph->nodes[node_n];

ggml_compute_forward(&params, node);

if (state->ith == 0 && cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
state->threadpool->ec = GGML_STATUS_ABORTED;
if (state->ith == 0 && cplan->abort_callback &&
cplan->abort_callback(cplan->abort_callback_data)) {
tp->stop = true;
tp->ec = GGML_STATUS_ABORTED;
}

ggml_barrier(state->threadpool);

if (state->threadpool->ec != GGML_STATUS_SUCCESS) {
break;
}
}

return 0;
@@ -20220,6 +20219,7 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
threadpool->cgraph = cgraph;
threadpool->cplan = cplan;
threadpool->current_chunk = 0;
threadpool->stop = false;
threadpool->ec = GGML_STATUS_SUCCESS;
}

0 comments on commit e833ca4

Please sign in to comment.