Skip to content

Commit 4633d93

Browse files
Xarbirusggerganov
authored andcommitted
ggml : add abort_callback for cpu backend (ggml/725)
* a way to use abort_callback with the cpu backend * whisper update
1 parent 4b7b38b commit 4633d93

File tree

4 files changed

+33
-9
lines changed

4 files changed

+33
-9
lines changed

ggml-backend.c

+22-4
Original file line numberDiff line numberDiff line change
@@ -653,6 +653,9 @@ struct ggml_backend_cpu_context {
653653
int n_threads;
654654
void * work_data;
655655
size_t work_size;
656+
657+
ggml_abort_callback abort_callback;
658+
void * abort_callback_data;
656659
};
657660

658661
GGML_CALL static const char * ggml_backend_cpu_name(ggml_backend_t backend) {
@@ -691,6 +694,9 @@ GGML_CALL static ggml_backend_graph_plan_t ggml_backend_cpu_graph_plan_create(gg
691694
cpu_plan->cplan.work_data = malloc(cpu_plan->cplan.work_size);
692695
}
693696

697+
cpu_plan->cplan.abort_callback = cpu_ctx->abort_callback;
698+
cpu_plan->cplan.abort_callback_data = cpu_ctx->abort_callback_data;
699+
694700
return cpu_plan;
695701
}
696702

@@ -721,9 +727,11 @@ GGML_CALL static bool ggml_backend_cpu_graph_compute(ggml_backend_t backend, str
721727
cpu_ctx->work_data = realloc(cpu_ctx->work_data, cplan.work_size);
722728
cpu_ctx->work_size = cplan.work_size;
723729
}
724-
725730
cplan.work_data = cpu_ctx->work_data;
726731

732+
cplan.abort_callback = cpu_ctx->abort_callback;
733+
cplan.abort_callback_data = cpu_ctx->abort_callback_data;
734+
727735
ggml_graph_compute(cgraph, &cplan);
728736
return true;
729737
}
@@ -759,9 +767,11 @@ static struct ggml_backend_i cpu_backend_i = {
759767
ggml_backend_t ggml_backend_cpu_init(void) {
760768
struct ggml_backend_cpu_context * ctx = malloc(sizeof(struct ggml_backend_cpu_context));
761769

762-
ctx->n_threads = GGML_DEFAULT_N_THREADS;
763-
ctx->work_data = NULL;
764-
ctx->work_size = 0;
770+
ctx->n_threads = GGML_DEFAULT_N_THREADS;
771+
ctx->work_data = NULL;
772+
ctx->work_size = 0;
773+
ctx->abort_callback = NULL;
774+
ctx->abort_callback_data = NULL;
765775

766776
ggml_backend_t cpu_backend = malloc(sizeof(struct ggml_backend));
767777

@@ -783,6 +793,14 @@ void ggml_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads) {
783793
ctx->n_threads = n_threads;
784794
}
785795

796+
void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data) {
797+
GGML_ASSERT(ggml_backend_is_cpu(backend_cpu));
798+
799+
struct ggml_backend_cpu_context * ctx = (struct ggml_backend_cpu_context *)backend_cpu->context;
800+
ctx->abort_callback = abort_callback;
801+
ctx->abort_callback_data = abort_callback_data;
802+
}
803+
786804
GGML_CALL ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) {
787805
return ggml_backend_buffer_init(ggml_backend_cpu_buffer_type(), cpu_backend_buffer_i_from_ptr, ptr, size);
788806
}

ggml-backend.h

+3-2
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,9 @@ extern "C" {
8383

8484
GGML_API ggml_backend_t ggml_backend_cpu_init(void);
8585

86-
GGML_API GGML_CALL bool ggml_backend_is_cpu (ggml_backend_t backend);
87-
GGML_API void ggml_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
86+
GGML_API GGML_CALL bool ggml_backend_is_cpu (ggml_backend_t backend);
87+
GGML_API void ggml_backend_cpu_set_n_threads (ggml_backend_t backend_cpu, int n_threads);
88+
GGML_API void ggml_backend_cpu_set_abort_callback(ggml_backend_t backend_cpu, ggml_abort_callback abort_callback, void * abort_callback_data);
8889

8990
// Create a backend buffer from an existing pointer
9091
GGML_API GGML_CALL ggml_backend_buffer_t ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size);

ggml.c

+1-1
Original file line numberDiff line numberDiff line change
@@ -16649,7 +16649,7 @@ struct ggml_compute_state_shared {
1664916649
atomic_int node_n; // active graph node
1665016650
atomic_int node_task; // active graph node task phase
1665116651

16652-
bool (*abort_callback)(void * data); // abort ggml_graph_compute when true
16652+
ggml_abort_callback abort_callback; // abort ggml_graph_compute when true
1665316653
void * abort_callback_data;
1665416654
};
1665516655

ggml.h

+7-2
Original file line numberDiff line numberDiff line change
@@ -567,6 +567,11 @@ extern "C" {
567567

568568
static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor);
569569

570+
// Abort callback
571+
// If not NULL, called before ggml computation
572+
// If it returns true, the computation is aborted
573+
typedef bool (*ggml_abort_callback)(void * data);
574+
570575
// the compute plan that needs to be prepared for ggml_graph_compute()
571576
// since https://github.com/ggerganov/ggml/issues/287
572577
struct ggml_cplan {
@@ -576,8 +581,8 @@ extern "C" {
576581
int n_threads;
577582

578583
// abort ggml_graph_compute when true
579-
bool (*abort_callback)(void * data);
580-
void * abort_callback_data;
584+
ggml_abort_callback abort_callback;
585+
void * abort_callback_data;
581586
};
582587

583588
enum ggml_cgraph_eval_order {

0 commit comments

Comments
 (0)