Skip to content

Commit e833ca4

Browse files
threadpool: improve abort handling
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.
1 parent b9763b3 commit e833ca4

File tree

1 file changed

+11
-11
lines changed

1 file changed

+11
-11
lines changed

ggml/src/ggml.c

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19928,34 +19928,33 @@ struct ggml_cplan ggml_graph_plan(
1992819928

1992919929
static thread_ret_t ggml_graph_compute_thread(void * data) {
1993019930
struct ggml_compute_state * state = (struct ggml_compute_state *) data;
19931+
struct ggml_threadpool * tp = state->threadpool;
1993119932

19932-
const struct ggml_cgraph * cgraph = state->threadpool->cgraph;
19933-
const struct ggml_cplan * cplan = state->threadpool->cplan;
19933+
const struct ggml_cgraph * cgraph = tp->cgraph;
19934+
const struct ggml_cplan * cplan = tp->cplan;
1993419935

1993519936
set_numa_thread_affinity(state->ith);
1993619937

1993719938
struct ggml_compute_params params = {
1993819939
/*.ith =*/ state->ith,
19939-
/*.nth =*/ state->threadpool->n_threads_cur,
19940+
/*.nth =*/ atomic_load_explicit(&tp->n_threads_cur, memory_order_relaxed),
1994019941
/*.wsize =*/ cplan->work_size,
1994119942
/*.wdata =*/ cplan->work_data,
19942-
/*.threadpool=*/ state->threadpool,
19943+
/*.threadpool=*/ tp,
1994319944
};
1994419945

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

1994819949
ggml_compute_forward(&params, node);
1994919950

19950-
if (state->ith == 0 && cplan->abort_callback && cplan->abort_callback(cplan->abort_callback_data)) {
19951-
state->threadpool->ec = GGML_STATUS_ABORTED;
19951+
if (state->ith == 0 && cplan->abort_callback &&
19952+
cplan->abort_callback(cplan->abort_callback_data)) {
19953+
tp->stop = true;
19954+
tp->ec = GGML_STATUS_ABORTED;
1995219955
}
1995319956

1995419957
ggml_barrier(state->threadpool);
19955-
19956-
if (state->threadpool->ec != GGML_STATUS_SUCCESS) {
19957-
break;
19958-
}
1995919958
}
1996019959

1996119960
return 0;
@@ -20220,6 +20219,7 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
2022020219
threadpool->cgraph = cgraph;
2022120220
threadpool->cplan = cplan;
2022220221
threadpool->current_chunk = 0;
20222+
threadpool->stop = false;
2022320223
threadpool->ec = GGML_STATUS_SUCCESS;
2022420224
}
2022520225

0 commit comments

Comments
 (0)