extern/whisper.cpp
1
0
Fork 1
mirror of https://github.com/ggerganov/whisper.cpp.git synced 2025-03-22 03:56:46 +01:00
Code Issues Packages Projects Releases Wiki Activity

ggml-alloc : v3 (ggml/727)

Browse Source 
* ggml-alloc v3

ggml-ci

* fix ci

ggml-ci

* whisper : check for backend buffer allocation failures

* whisper : avoid leaks when initialization fails

* cleanup

ggml-ci

* style fixes

ggml-ci
This commit is contained in:
slaren 2024-02-11 13:37:58 +01:00 committed by Georgi Gerganov
parent a6fb6ab597
commit 1d3270cc8f
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
7 changed files with 1207 additions and 1226 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1235
ggml-alloc.c
View File

File diff suppressed because it is too large Load Diff

104
ggml-alloc.h
View File

@ -6,88 +6,62 @@
extern "C" { extern "C" {
#endif #endif
struct ggml_backend; typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t;
struct ggml_backend_buffer; typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
struct ggml_backend_buffer_type; typedef struct ggml_backend * ggml_backend_t;
//
// Legacy API
//
typedef struct ggml_allocr * ggml_allocr_t;
// initialize allocator for use with CPU backend only
GGML_API ggml_allocr_t ggml_allocr_new(void * data, size_t size, size_t alignment);
GGML_API ggml_allocr_t ggml_allocr_new_measure(size_t alignment);
// initialize allocator for use with ggml-backend
GGML_API ggml_allocr_t ggml_allocr_new_from_buffer(struct ggml_backend_buffer * buffer);
GGML_API ggml_allocr_t ggml_allocr_new_from_backend(struct ggml_backend * backend, size_t size); // allocates an owned buffer
GGML_API ggml_allocr_t ggml_allocr_new_measure_from_backend(struct ggml_backend * backend);
GGML_API struct ggml_backend_buffer * ggml_allocr_get_buffer(ggml_allocr_t alloc);
// tell the allocator to parse nodes following the order described in the list
// you should call this if your graph are optimized to execute out-of-order
GGML_API void ggml_allocr_set_parse_seq(ggml_allocr_t alloc, const int * list, int n);
GGML_API void ggml_allocr_free (ggml_allocr_t alloc);
GGML_API bool ggml_allocr_is_measure (ggml_allocr_t alloc);
GGML_API void ggml_allocr_reset (ggml_allocr_t alloc);
GGML_API void ggml_allocr_alloc (ggml_allocr_t alloc, struct ggml_tensor * tensor);
GGML_API size_t ggml_allocr_max_size (ggml_allocr_t alloc);
GGML_API size_t ggml_allocr_alloc_graph(ggml_allocr_t alloc, struct ggml_cgraph * graph);
//
// ggml-backend v2 API
//
// Separate tensor and graph allocator objects
// This is necessary for multi-backend allocation because the graph allocator needs to use multiple tensor allocators
// The original API is kept as a wrapper around the new API
// Tensor allocator // Tensor allocator
typedef struct ggml_tallocr * ggml_tallocr_t; typedef struct ggml_tallocr * ggml_tallocr_t;
GGML_API ggml_tallocr_t ggml_tallocr_new(void * data, size_t size, size_t alignment); GGML_API ggml_tallocr_t ggml_tallocr_new(ggml_backend_buffer_t buffer);
GGML_API ggml_tallocr_t ggml_tallocr_new_measure(size_t alignment);
GGML_API ggml_tallocr_t ggml_tallocr_new_from_buft(struct ggml_backend_buffer_type * buft, size_t size);
GGML_API ggml_tallocr_t ggml_tallocr_new_from_backend(struct ggml_backend * backend, size_t size); // allocates an owned buffer
GGML_API ggml_tallocr_t ggml_tallocr_new_from_buffer(struct ggml_backend_buffer * buffer);
GGML_API ggml_tallocr_t ggml_tallocr_new_measure_from_buft(struct ggml_backend_buffer_type * buft);
GGML_API ggml_tallocr_t ggml_tallocr_new_measure_from_backend(struct ggml_backend * backend);
GGML_API struct ggml_backend_buffer * ggml_tallocr_get_buffer(ggml_tallocr_t talloc);
GGML_API void ggml_tallocr_free(ggml_tallocr_t talloc); GGML_API void ggml_tallocr_free(ggml_tallocr_t talloc);
GGML_API bool ggml_tallocr_is_measure (ggml_tallocr_t talloc);
GGML_API void ggml_tallocr_reset (ggml_tallocr_t talloc);
GGML_API void ggml_tallocr_alloc(ggml_tallocr_t talloc, struct ggml_tensor * tensor); GGML_API void ggml_tallocr_alloc(ggml_tallocr_t talloc, struct ggml_tensor * tensor);
GGML_API size_t ggml_tallocr_max_size (ggml_tallocr_t talloc);
// Graph allocator // Graph allocator
/*
Example usage:
ggml_gallocr_t galloc = ggml_gallocr_new(ggml_bacckend_cpu_buffer_type());
// optional: create a worst-case graph and reserve the buffers to avoid reallocations
ggml_gallocr_reserve(galloc, build_graph(max_batch));
// allocate the graph
struct ggml_cgraph * graph = build_graph(batch);
ggml_gallocr_alloc_graph(galloc, graph);
printf("compute buffer size: %zu bytes\n", ggml_gallocr_get_buffer_size(galloc, 0));
// evaluate the graph
ggml_backend_graph_compute(backend, graph);
*/
// special tensor flags for use with the graph allocator:
// ggml_set_input(): all input tensors are allocated at the beginning of the graph in non-overlapping addresses
// ggml_set_output(): output tensors are never freed and never overwritten
typedef struct ggml_gallocr * ggml_gallocr_t; typedef struct ggml_gallocr * ggml_gallocr_t;
GGML_API ggml_gallocr_t ggml_gallocr_new(void); GGML_API ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft);
GGML_API ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs);
GGML_API void ggml_gallocr_free(ggml_gallocr_t galloc); GGML_API void ggml_gallocr_free(ggml_gallocr_t galloc);
GGML_API void ggml_gallocr_set_parse_seq(ggml_gallocr_t galloc, const int * list, int n); // pre-allocate buffers from a measure graph - does not allocate or modify the graph
GGML_API size_t ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, ggml_tallocr_t talloc, struct ggml_cgraph * graph); // call with a worst-case graph to avoid buffer reallocations
// not strictly required for single buffer usage: ggml_gallocr_alloc_graph will reallocate the buffers automatically if needed
// returns false if the buffer allocation failed
GGML_API bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
GGML_API bool ggml_gallocr_reserve_n(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids);
// Allocate tensors from the allocators given by the hash table // automatic reallocation if the topology changes when using a single buffer
GGML_API void ggml_gallocr_alloc_graph_n( // returns false if using multiple buffers and a re-allocation is needed (call ggml_gallocr_reserve_n first to set the node buffers)
ggml_gallocr_t galloc, GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
struct ggml_cgraph * graph,
struct ggml_hash_set hash_set,
ggml_tallocr_t * hash_node_talloc);
GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
// Utils // Utils
// Create a buffer and allocate all the tensors in a ggml_context // Create a buffer and allocate all the tensors in a ggml_context
GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, struct ggml_backend_buffer_type * buft); GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft);
GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, struct ggml_backend * backend); GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend);
#ifdef __cplusplus #ifdef __cplusplus
} }

479
ggml-backend.c
View File

@ -475,6 +475,8 @@ ggml_backend_buffer_t ggml_backend_reg_alloc_buffer(size_t i, size_t size) {
// backend CPU // backend CPU
static const size_t TENSOR_ALIGNMENT = 64; // should be enough for AVX 512
GGML_CALL static const char * ggml_backend_cpu_buffer_name(ggml_backend_buffer_t buffer) { GGML_CALL static const char * ggml_backend_cpu_buffer_name(ggml_backend_buffer_t buffer) {
return "CPU"; return "CPU";
@ -482,7 +484,14 @@ GGML_CALL static const char * ggml_backend_cpu_buffer_name(ggml_backend_buffer_t
} }
GGML_CALL static void * ggml_backend_cpu_buffer_get_base(ggml_backend_buffer_t buffer) { GGML_CALL static void * ggml_backend_cpu_buffer_get_base(ggml_backend_buffer_t buffer) {
return (void *)buffer->context; uintptr_t data = (uintptr_t)buffer->context;
// align the buffer
if (data % TENSOR_ALIGNMENT != 0) {
data = GGML_PAD(data, TENSOR_ALIGNMENT);
}
return (void *)data;
} }
GGML_CALL static void ggml_backend_cpu_buffer_free_buffer(ggml_backend_buffer_t buffer) { GGML_CALL static void ggml_backend_cpu_buffer_free_buffer(ggml_backend_buffer_t buffer) {
@ -540,8 +549,6 @@ static struct ggml_backend_buffer_i cpu_backend_buffer_i_from_ptr = {
/* .reset = */ NULL, /* .reset = */ NULL,
}; };
static const size_t TENSOR_ALIGNMENT = 64; // should be enough for AVX 512
GGML_CALL static const char * ggml_backend_cpu_buffer_type_get_name(ggml_backend_buffer_type_t buft) { GGML_CALL static const char * ggml_backend_cpu_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
return "CPU"; return "CPU";
@ -550,9 +557,11 @@ GGML_CALL static const char * ggml_backend_cpu_buffer_type_get_name(ggml_backend
GGML_CALL static ggml_backend_buffer_t ggml_backend_cpu_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { GGML_CALL static ggml_backend_buffer_t ggml_backend_cpu_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
size += TENSOR_ALIGNMENT; // malloc may return an address that is not aligned size += TENSOR_ALIGNMENT; // malloc may return an address that is not aligned
void * data = malloc(size); // TODO: maybe use GGML_ALIGNED_MALLOC? void * data = malloc(size); // TODO: use GGML_ALIGNED_MALLOC (move to ggml-impl.h)
if (data == NULL) {
GGML_ASSERT(data != NULL && "failed to allocate buffer"); fprintf(stderr, "%s: failed to allocate buffer of size %zu\n", __func__, size);
return NULL;
}
return ggml_backend_buffer_init(buft, cpu_backend_buffer_i, data, size); return ggml_backend_buffer_init(buft, cpu_backend_buffer_i, data, size);
} }
@ -766,6 +775,9 @@ static struct ggml_backend_i cpu_backend_i = {
ggml_backend_t ggml_backend_cpu_init(void) { ggml_backend_t ggml_backend_cpu_init(void) {
struct ggml_backend_cpu_context * ctx = malloc(sizeof(struct ggml_backend_cpu_context)); struct ggml_backend_cpu_context * ctx = malloc(sizeof(struct ggml_backend_cpu_context));
if (ctx == NULL) {
return NULL;
}
ctx->n_threads = GGML_DEFAULT_N_THREADS; ctx->n_threads = GGML_DEFAULT_N_THREADS;
ctx->work_data = NULL; ctx->work_data = NULL;
@ -774,6 +786,10 @@ ggml_backend_t ggml_backend_cpu_init(void) {
ctx->abort_callback_data = NULL; ctx->abort_callback_data = NULL;
ggml_backend_t cpu_backend = malloc(sizeof(struct ggml_backend)); ggml_backend_t cpu_backend = malloc(sizeof(struct ggml_backend));
if (cpu_backend == NULL) {
free(ctx);
return NULL;
}
*cpu_backend = (struct ggml_backend) { *cpu_backend = (struct ggml_backend) {
/* .interface = */ cpu_backend_i, /* .interface = */ cpu_backend_i,
@ -865,6 +881,8 @@ GGML_CALL ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_back
ctx->n_buffers = n_buffers; ctx->n_buffers = n_buffers;
ctx->buffers = (ggml_backend_buffer_t *) malloc(n_buffers * sizeof(ggml_backend_buffer_t)); ctx->buffers = (ggml_backend_buffer_t *) malloc(n_buffers * sizeof(ggml_backend_buffer_t));
GGML_ASSERT(ctx->buffers != NULL);
size_t total_size = 0; size_t total_size = 0;
for (size_t i = 0; i < n_buffers; i++) { for (size_t i = 0; i < n_buffers; i++) {
ctx->buffers[i] = buffers[i]; ctx->buffers[i] = buffers[i];
@ -886,6 +904,18 @@ GGML_CALL void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer,
} }
} }
// creates a copy of the tensor with the same memory layout
static struct ggml_tensor * ggml_dup_tensor_layout(struct ggml_context * ctx, const struct ggml_tensor * tensor) {
struct ggml_tensor * dup = ggml_dup_tensor(ctx, tensor);
for (int i = 0; i < GGML_MAX_DIMS; i++) {
dup->nb[i] = tensor->nb[i];
}
return dup;
}
static bool ggml_is_view_op(enum ggml_op op) {
return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE;
}
// scheduler // scheduler
@ -894,7 +924,7 @@ GGML_CALL void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer,
#define GGML_MAX_SPLIT_INPUTS 16 #define GGML_MAX_SPLIT_INPUTS 16
struct ggml_backend_sched_split { struct ggml_backend_sched_split {
ggml_tallocr_t tallocr; int backend_id;
int i_start; int i_start;
int i_end; int i_end;
struct ggml_tensor * inputs[GGML_MAX_SPLIT_INPUTS]; struct ggml_tensor * inputs[GGML_MAX_SPLIT_INPUTS];
@ -909,15 +939,17 @@ struct ggml_backend_sched {
int n_backends; int n_backends;
ggml_backend_t backends[GGML_MAX_BACKENDS]; ggml_backend_t backends[GGML_MAX_BACKENDS];
ggml_backend_buffer_type_t bufts[GGML_MAX_BACKENDS]; ggml_backend_buffer_type_t bufts[GGML_MAX_BACKENDS];
ggml_tallocr_t tallocs[GGML_MAX_BACKENDS];
ggml_gallocr_t galloc; ggml_gallocr_t galloc;
// hash keys of the nodes in the graph // hash keys of the nodes in the graph
struct ggml_hash_set hash_set; struct ggml_hash_set hash_set;
// hash values (arrays of [hash_set.size]) // hash values
ggml_tallocr_t * node_talloc; // tallocr assigned to each node (indirectly this is the backend) int * tensor_backend_id;
struct ggml_tensor * (* node_copies)[GGML_MAX_BACKENDS]; // copies of each node for each destination backend struct ggml_tensor * (* tensor_copies)[GGML_MAX_BACKENDS];
int * node_backend_ids; // [n_nodes]
int n_nodes;
// copy of the graph with modified inputs // copy of the graph with modified inputs
struct ggml_cgraph * graph; struct ggml_cgraph * graph;
@ -927,77 +959,46 @@ struct ggml_backend_sched {
struct ggml_context * ctx; struct ggml_context * ctx;
ggml_backend_sched_eval_callback callback_eval;
void * callback_eval_user_data;
// align context_buffer to GGML_MEM_ALIGN // align context_buffer to GGML_MEM_ALIGN
#ifdef _MSC_VER #ifdef _MSC_VER
__declspec(align(GGML_MEM_ALIGN)) __declspec(align(GGML_MEM_ALIGN))
#else #else
__attribute__((aligned(GGML_MEM_ALIGN))) __attribute__((aligned(GGML_MEM_ALIGN)))
#endif #endif
char context_buffer[GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS*sizeof(struct ggml_tensor) + sizeof(struct ggml_cgraph)]; char context_buffer[GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS*2*sizeof(struct ggml_tensor) + sizeof(struct ggml_cgraph)];
ggml_backend_sched_eval_callback callback_eval;
void * callback_eval_user_data;
}; };
#define hash_id(node) ggml_hash_find_or_insert(sched->hash_set, node) #define hash_id(node) ggml_hash_find_or_insert(sched->hash_set, node)
#define node_allocr(node) sched->node_talloc[hash_id(node)] #define tensor_backend_id(node) sched->tensor_backend_id[hash_id(node)]
#define tensor_backend(node) (tensor_backend_id(node) == -1 ? NULL : sched->backends[tensor_backend_id(node)])
static bool ggml_is_view_op(enum ggml_op op) { // returns the priority of the backend, lower id is higher priority
return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) {
}
// returns the priority of the backend, lower is better
static int sched_backend_prio(ggml_backend_sched_t sched, ggml_backend_t backend) {
for (int i = 0; i < sched->n_backends; i++) { for (int i = 0; i < sched->n_backends; i++) {
if (sched->backends[i] == backend) { if (sched->backends[i] == backend) {
return i; return i;
} }
} }
return INT_MAX; return -1;
} }
static int sched_allocr_prio(ggml_backend_sched_t sched, ggml_tallocr_t allocr) { static int ggml_backend_sched_backend_from_buffer(ggml_backend_sched_t sched, ggml_backend_buffer_t buffer) {
for (int i = 0; i < sched->n_backends; i++) {
if (sched->tallocs[i] == allocr) {
return i;
}
}
return INT_MAX;
}
static ggml_tallocr_t sched_allocr_from_buffer(ggml_backend_sched_t sched, ggml_backend_buffer_t buffer) {
if (buffer == NULL) { if (buffer == NULL) {
return NULL; return -1;
}
// check if this is already allocate in a allocr buffer (from user manual allocations)
for (int i = 0; i < sched->n_backends; i++) {
if (ggml_tallocr_get_buffer(sched->tallocs[i]) == buffer) {
return sched->tallocs[i];
}
} }
// find highest prio backend that supports the buffer type // find highest prio backend that supports the buffer type
for (int i = 0; i < sched->n_backends; i++) { for (int i = 0; i < sched->n_backends; i++) {
if (ggml_backend_buft_supports_backend(buffer->buft, sched->backends[i])) { if (ggml_backend_buft_supports_backend(buffer->buft, sched->backends[i])) {
return sched->tallocs[i]; return i;
} }
} }
GGML_ASSERT(false && "tensor buffer type not supported by any backend"); GGML_ASSERT(false && "tensor buffer type not supported by any backend");
} }
static ggml_backend_t get_allocr_backend(ggml_backend_sched_t sched, ggml_tallocr_t allocr) {
if (allocr == NULL) {
return NULL;
}
for (int i = 0; i < sched->n_backends; i++) {
if (sched->tallocs[i] == allocr) {
return sched->backends[i];
}
}
GGML_UNREACHABLE();
}
#if 0 #if 0
static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS][128]; // debug only static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS][128]; // debug only
#define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__) #define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__)
@ -1008,37 +1009,39 @@ static char causes[GGML_DEFAULT_GRAPH_SIZE*16 + GGML_MAX_SPLITS*GGML_MAX_SPLIT_I
#endif #endif
// returns the backend that should be used for the node based on the current locations // returns the backend that should be used for the node based on the current locations
static ggml_tallocr_t sched_allocr_from_cur(ggml_backend_sched_t sched, struct ggml_tensor * node) { static int ggml_backend_sched_backend_id_from_cur(ggml_backend_sched_t sched, struct ggml_tensor * tensor) {
// TODO: use supports_op to check if the backend supports the op
// assign pre-allocated nodes to their backend // assign pre-allocated nodes to their backend
// dst // dst
ggml_tallocr_t cur_allocr = sched_allocr_from_buffer(sched, node->buffer); int cur_backend = ggml_backend_sched_backend_from_buffer(sched, tensor->buffer);
if (cur_allocr != NULL) { if (cur_backend != -1) {
SET_CAUSE(node, "1.dst"); SET_CAUSE(node, "1.dst");
return cur_allocr; return cur_backend;
} }
// view_src // view_src
if (node->view_src != NULL) { if (tensor->view_src != NULL) {
cur_allocr = sched_allocr_from_buffer(sched, node->view_src->buffer); cur_backend = ggml_backend_sched_backend_from_buffer(sched, tensor->view_src->buffer);
if (cur_allocr != NULL) { if (cur_backend != -1) {
SET_CAUSE(node, "1.vsrc"); SET_CAUSE(node, "1.vsrc");
return cur_allocr; return cur_backend;
} }
} }
// assign nodes that use weights to the backend of the weights // assign nodes that use weights to the backend of the weights
for (int i = 0; i < GGML_MAX_SRC; i++) { for (int i = 0; i < GGML_MAX_SRC; i++) {
const struct ggml_tensor * src = node->src[i]; const struct ggml_tensor * src = tensor->src[i];
if (src == NULL) { if (src == NULL) {
break; break;
} }
if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { if (src->buffer != NULL && src->buffer->usage == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
ggml_tallocr_t src_allocr = sched_allocr_from_buffer(sched, src->buffer); int src_backend = ggml_backend_sched_backend_from_buffer(sched, src->buffer);
// operations with weights are always run on the same backend as the weights // operations with weights are always run on the same backend as the weights
SET_CAUSE(node, "1.wgt%d", i); SET_CAUSE(node, "1.wgt%d", i);
return src_allocr; return src_backend;
} }
} }
return NULL; return -1;
} }
static char * fmt_size(size_t size) { static char * fmt_size(size_t size) {
@ -1051,11 +1054,11 @@ static char * fmt_size(size_t size) {
return buffer; return buffer;
} }
static void sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgraph * graph) {
int cur_split = 0; int cur_split = 0;
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) { if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) {
ggml_backend_t split_backend = get_allocr_backend(sched, sched->splits[cur_split].tallocr); ggml_backend_t split_backend = sched->backends[sched->splits[cur_split].backend_id];
fprintf(stderr, "\n## SPLIT #%d: %s # %d inputs: ", cur_split, ggml_backend_name(split_backend), fprintf(stderr, "\n## SPLIT #%d: %s # %d inputs: ", cur_split, ggml_backend_name(split_backend),
sched->splits[cur_split].n_inputs); sched->splits[cur_split].n_inputs);
for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) { for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) {
@ -1069,17 +1072,15 @@ static void sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgra
if (ggml_is_view_op(node->op)) { if (ggml_is_view_op(node->op)) {
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); ggml_backend_t tensor_backend = tensor_backend(node);
ggml_backend_t node_backend = node_allocr ? get_allocr_backend(sched, node_allocr) : NULL; // FIXME:
fprintf(stderr, "node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s]:", i, ggml_op_name(node->op), node->name, fprintf(stderr, "node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s]:", i, ggml_op_name(node->op), node->name,
fmt_size(ggml_nbytes(node)), node_allocr ? ggml_backend_name(node_backend) : "NULL", GET_CAUSE(node)); fmt_size(ggml_nbytes(node)), tensor_backend ? ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node));
for (int j = 0; j < GGML_MAX_SRC; j++) { for (int j = 0; j < GGML_MAX_SRC; j++) {
struct ggml_tensor * src = node->src[j]; struct ggml_tensor * src = node->src[j];
if (src == NULL) { if (src == NULL) {
break; break;
} }
ggml_tallocr_t src_allocr = node_allocr(src); ggml_backend_t src_backend = tensor_backend(src);
ggml_backend_t src_backend = src_allocr ? get_allocr_backend(sched, src_allocr) : NULL;
fprintf(stderr, " %20.20s (%5.5s) [%5.5s %8.8s]", src->name, fprintf(stderr, " %20.20s (%5.5s) [%5.5s %8.8s]", src->name,
fmt_size(ggml_nbytes(src)), src_backend ? ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src)); fmt_size(ggml_nbytes(src)), src_backend ? ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src));
} }
@ -1087,23 +1088,13 @@ static void sched_print_assignments(ggml_backend_sched_t sched, struct ggml_cgra
} }
} }
// creates a copy of the tensor with the same memory layout
static struct ggml_tensor * ggml_dup_tensor_layout(struct ggml_context * ctx, const struct ggml_tensor * tensor) {
struct ggml_tensor * dup = ggml_dup_tensor(ctx, tensor);
for (int i = 0; i < GGML_MAX_DIMS; i++) {
dup->nb[i] = tensor->nb[i];
}
return dup;
}
//#define DEBUG_PASS1 //#define DEBUG_PASS1
//#define DEBUG_PASS2 //#define DEBUG_PASS2
//#define DEBUG_PASS3 //#define DEBUG_PASS3
//#define DEBUG_PASS4 //#define DEBUG_PASS4
// assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend // assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend
static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { static void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) {
// reset splits // reset splits
sched->n_splits = 0; sched->n_splits = 0;
sched->is_reset = false; sched->is_reset = false;
@ -1125,28 +1116,28 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 1: assign backends to ops with pre-allocated inputs // pass 1: assign backends to ops with pre-allocated inputs
for (int i = 0; i < graph->n_leafs; i++) { for (int i = 0; i < graph->n_leafs; i++) {
struct ggml_tensor * leaf = graph->leafs[i]; struct ggml_tensor * leaf = graph->leafs[i];
if (node_allocr(leaf) != NULL) { if (tensor_backend_id(leaf) != -1) {
// do not overwrite user assignments // do not overwrite user assignments
continue; continue;
} }
node_allocr(leaf) = sched_allocr_from_cur(sched, leaf); tensor_backend_id(leaf) = ggml_backend_sched_backend_id_from_cur(sched, leaf);
} }
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (node_allocr(node) != NULL) { if (tensor_backend_id(node) != -1) {
// do not overwrite user assignments // do not overwrite user assignments
continue; continue;
} }
node_allocr(node) = sched_allocr_from_cur(sched, node); tensor_backend_id(node) = ggml_backend_sched_backend_id_from_cur(sched, node);
// src // src
for (int j = 0; j < GGML_MAX_SRC; j++) { for (int j = 0; j < GGML_MAX_SRC; j++) {
struct ggml_tensor * src = node->src[j]; struct ggml_tensor * src = node->src[j];
if (src == NULL) { if (src == NULL) {
break; break;
} }
if (node_allocr(src) == NULL) { if (tensor_backend_id(src) == -1) {
node_allocr(src) = sched_allocr_from_cur(sched, src); tensor_backend_id(src) = ggml_backend_sched_backend_id_from_cur(sched, src);
} }
} }
} }
@ -1161,22 +1152,22 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 2.1 expand gpu up // pass 2.1 expand gpu up
{ {
ggml_tallocr_t cur_allocr = NULL; int cur_backend_id = -1;
for (int i = graph->n_nodes - 1; i >= 0; i--) { for (int i = graph->n_nodes - 1; i >= 0; i--) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (ggml_is_view_op(node->op)) { if (ggml_is_view_op(node->op)) {
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); int tensor_backend_id = tensor_backend_id(node);
if (node_allocr != NULL) { if (tensor_backend_id != -1) {
if (sched_allocr_prio(sched, node_allocr) == sched->n_backends - 1) { if (tensor_backend_id == sched->n_backends - 1) {
// skip cpu (lowest prio backend) // skip cpu (lowest prio backend)
cur_allocr = NULL; cur_backend_id = -1;
} else { } else {
cur_allocr = node_allocr; cur_backend_id = tensor_backend_id;
} }
} else { } else {
node_allocr(node) = cur_allocr; tensor_backend_id(node) = cur_backend_id;
SET_CAUSE(node, "2.1"); SET_CAUSE(node, "2.1");
} }
} }
@ -1184,22 +1175,22 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 2.2 expand gpu down // pass 2.2 expand gpu down
{ {
ggml_tallocr_t cur_allocr = NULL; int cur_backend_id = -1;
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (ggml_is_view_op(node->op)) { if (ggml_is_view_op(node->op)) {
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); int tensor_backend_id = tensor_backend_id(node);
if (node_allocr != NULL) { if (tensor_backend_id != -1) {
if (sched_allocr_prio(sched, node_allocr) == sched->n_backends - 1) { if (tensor_backend_id == sched->n_backends - 1) {
// skip cpu (lowest prio backend) // skip cpu (lowest prio backend)
cur_allocr = NULL; cur_backend_id = -1;
} else { } else {
cur_allocr = node_allocr; cur_backend_id = tensor_backend_id;
} }
} else { } else {
node_allocr(node) = cur_allocr; tensor_backend_id(node) = cur_backend_id;
SET_CAUSE(node, "2.2"); SET_CAUSE(node, "2.2");
} }
} }
@ -1207,17 +1198,17 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 2.3 expand rest up // pass 2.3 expand rest up
{ {
ggml_tallocr_t cur_allocr = NULL; int cur_backend_id = -1;
for (int i = graph->n_nodes - 1; i >= 0; i--) { for (int i = graph->n_nodes - 1; i >= 0; i--) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (ggml_is_view_op(node->op)) { if (ggml_is_view_op(node->op)) {
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); int tensor_backend_id = tensor_backend_id(node);
if (node_allocr != NULL) { if (tensor_backend_id != -1) {
cur_allocr = node_allocr; cur_backend_id = tensor_backend_id;
} else { } else {
node_allocr(node) = cur_allocr; tensor_backend_id(node) = cur_backend_id;
SET_CAUSE(node, "2.3"); SET_CAUSE(node, "2.3");
} }
} }
@ -1225,17 +1216,17 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 2.4 expand rest down // pass 2.4 expand rest down
{ {
ggml_tallocr_t cur_allocr = NULL; int cur_backend_id = -1;
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (ggml_is_view_op(node->op)) { if (ggml_is_view_op(node->op)) {
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); int tensor_backend_id = tensor_backend_id(node);
if (node_allocr != NULL) { if (tensor_backend_id != -1) {
cur_allocr = node_allocr; cur_backend_id = tensor_backend_id;
} else { } else {
node_allocr(node) = cur_allocr; tensor_backend_id(node) = cur_backend_id;
SET_CAUSE(node, "2.4"); SET_CAUSE(node, "2.4");
} }
} }
@ -1247,9 +1238,9 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// pass 3: assign backends to remaining src from dst and view_src // pass 3: assign backends to remaining src from dst and view_src
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
ggml_tallocr_t cur_allocr = node_allocr(node); int cur_backend_id = tensor_backend_id(node);
if (node->view_src != NULL && cur_allocr == NULL) { if (node->view_src != NULL && cur_backend_id == -1) {
cur_allocr = node_allocr(node) = node_allocr(node->view_src); cur_backend_id = tensor_backend_id(node) = tensor_backend_id(node->view_src);
SET_CAUSE(node, "3.vsrc"); SET_CAUSE(node, "3.vsrc");
} }
for (int j = 0; j < GGML_MAX_SRC; j++) { for (int j = 0; j < GGML_MAX_SRC; j++) {
@ -1257,14 +1248,14 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
if (src == NULL) { if (src == NULL) {
break; break;
} }
ggml_tallocr_t src_allocr = node_allocr(src); int src_backend_id = tensor_backend_id(src);
if (src_allocr == NULL) { if (src_backend_id == -1) {
if (src->view_src != NULL) { if (src->view_src != NULL) {
// views are always on the same backend as the source // views are always on the same backend as the source
node_allocr(src) = node_allocr(src->view_src); tensor_backend_id(src) = tensor_backend_id(src->view_src);
SET_CAUSE(src, "3.vsrc"); SET_CAUSE(src, "3.vsrc");
} else { } else {
node_allocr(src) = cur_allocr; tensor_backend_id(src) = cur_backend_id;
SET_CAUSE(src, "3.cur"); SET_CAUSE(src, "3.cur");
} }
} }
@ -1281,15 +1272,14 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
if (!ggml_is_view_op(node->op)) { if (!ggml_is_view_op(node->op)) {
sched->splits[0].tallocr = node_allocr(node); sched->splits[0].backend_id = tensor_backend_id(node);
break; break;
} }
} }
sched->splits[0].i_start = 0; sched->splits[0].i_start = 0;
sched->splits[0].n_inputs = 0; sched->splits[0].n_inputs = 0;
memset(sched->splits[0].inputs, 0, sizeof(sched->splits[0].inputs)); //HACK memset(sched->splits[0].inputs, 0, sizeof(sched->splits[0].inputs)); //HACK
ggml_tallocr_t cur_allocr = sched->splits[0].tallocr; int cur_backend_id = sched->splits[0].backend_id;
size_t cur_backend_id = sched_allocr_prio(sched, cur_allocr);
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
@ -1297,19 +1287,18 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
continue; continue;
} }
ggml_tallocr_t node_allocr = node_allocr(node); int tensor_backend_id = tensor_backend_id(node);
GGML_ASSERT(node_allocr != NULL); // all nodes should be assigned by now GGML_ASSERT(tensor_backend_id != -1); // all nodes should be assigned by now
if (node_allocr != cur_allocr) { if (tensor_backend_id != cur_backend_id) {
sched->splits[cur_split].i_end = i; sched->splits[cur_split].i_end = i;
cur_split++; cur_split++;
GGML_ASSERT(cur_split < GGML_MAX_SPLITS); GGML_ASSERT(cur_split < GGML_MAX_SPLITS);
sched->splits[cur_split].tallocr = node_allocr; sched->splits[cur_split].backend_id = tensor_backend_id;
sched->splits[cur_split].i_start = i; sched->splits[cur_split].i_start = i;
sched->splits[cur_split].n_inputs = 0; sched->splits[cur_split].n_inputs = 0;
cur_allocr = node_allocr; cur_backend_id = tensor_backend_id;
cur_backend_id = sched_allocr_prio(sched, cur_allocr);
} }
// find inputs that are not on the same backend // find inputs that are not on the same backend
@ -1318,43 +1307,25 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
if (src == NULL) { if (src == NULL) {
break; break;
} }
ggml_tallocr_t src_allocr = node_allocr(src); int src_backend_id = tensor_backend_id(src);
GGML_ASSERT(src_allocr != NULL); // all inputs should be assigned by now assert(src_backend_id != -1); // all inputs should be assigned by now
if (src_allocr != node_allocr) { if (src_backend_id != tensor_backend_id) {
// create a copy of the input in the split's backend // create a copy of the input in the split's backend
size_t id = hash_id(src); size_t id = hash_id(src);
if (sched->node_copies[id][cur_backend_id] == NULL) { if (sched->tensor_copies[id][cur_backend_id] == NULL) {
ggml_backend_t backend = get_allocr_backend(sched, cur_allocr); ggml_backend_t backend = sched->backends[cur_backend_id];
struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src);
ggml_format_name(tensor_copy, "%s#%s", ggml_backend_name(backend), src->name); ggml_format_name(tensor_copy, "%s#%s", ggml_backend_name(backend), src->name);
sched->node_copies[id][cur_backend_id] = tensor_copy; sched->tensor_copies[id][cur_backend_id] = tensor_copy;
node_allocr(tensor_copy) = cur_allocr; tensor_backend_id(tensor_copy) = cur_backend_id;
SET_CAUSE(tensor_copy, "4.cpy"); SET_CAUSE(tensor_copy, "4.cpy");
int n_inputs = sched->splits[cur_split].n_inputs++; int n_inputs = sched->splits[cur_split].n_inputs++;
GGML_ASSERT(n_inputs < GGML_MAX_SPLIT_INPUTS); GGML_ASSERT(n_inputs < GGML_MAX_SPLIT_INPUTS);
sched->splits[cur_split].inputs[n_inputs] = src; sched->splits[cur_split].inputs[n_inputs] = src;
} }
node->src[j] = sched->node_copies[id][cur_backend_id]; node->src[j] = sched->tensor_copies[id][cur_backend_id];
#if 0
// check if the input is already in the split
bool found = false;
for (int k = 0; k < sched->splits[cur_split].n_inputs; k++) {
if (sched->splits[cur_split].inputs[k] == src) {
found = true;
break;
}
}
if (!found) {
int n_inputs = sched->splits[cur_split].n_inputs++;
//printf("split %d input %d: %s (%s)\n", cur_split, n_inputs, src->name, ggml_backend_name(get_allocr_backend(sched, src_allocr)));
GGML_ASSERT(n_inputs < GGML_MAX_SPLIT_INPUTS);
sched->splits[cur_split].inputs[n_inputs] = src;
}
#endif
} }
} }
} }
@ -1369,30 +1340,30 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
// sanity check: all sources should have the same backend as the node // sanity check: all sources should have the same backend as the node
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
ggml_tallocr_t node_allocr = node_allocr(node); ggml_backend_t tensor_backend = tensor_backend(node);
if (node_allocr == NULL) { if (tensor_backend == NULL) {
fprintf(stderr, "!!!!!!! %s has no backend\n", node->name); fprintf(stderr, "!!!!!!! %s has no backend\n", node->name);
} }
if (node->view_src != NULL && node_allocr != node_allocr(node->view_src)) { if (node->view_src != NULL && tensor_backend != tensor_backend(node->view_src)) {
fprintf(stderr, "!!!!!!! %s has backend %s, view_src %s has backend %s\n", fprintf(stderr, "!!!!!!! %s has backend %s, view_src %s has backend %s\n",
node->name, node_allocr ? ggml_backend_name(get_allocr_backend(sched, node_allocr)) : "NULL", node->name, tensor_backend ? ggml_backend_name(tensor_backend) : "NULL",
node->view_src->name, node_allocr(node->view_src) ? ggml_backend_name(get_allocr_backend(sched, node_allocr(node->view_src))) : "NULL"); node->view_src->name, tensor_backend(node->view_src) ? ggml_backend_name(tensor_backend(node->view_src)) : "NULL");
} }
for (int j = 0; j < GGML_MAX_SRC; j++) { for (int j = 0; j < GGML_MAX_SRC; j++) {
struct ggml_tensor * src = node->src[j]; struct ggml_tensor * src = node->src[j];
if (src == NULL) { if (src == NULL) {
break; break;
} }
ggml_tallocr_t src_allocr = node_allocr(src); ggml_backend_t src_backend = tensor_backend(src);
if (src_allocr != node_allocr /* && src_backend != NULL */) { // ignore nulls for now if (src_backend != tensor_backend /* && src_backend != NULL */) {
fprintf(stderr, "!!!! %s has backend %s, src %d (%s) has backend %s\n", fprintf(stderr, "!!!! %s has backend %s, src %d (%s) has backend %s\n",
node->name, node_allocr ? ggml_backend_name(get_allocr_backend(sched, node_allocr)) : "NULL", node->name, tensor_backend ? ggml_backend_name(tensor_backend) : "NULL",
j, src->name, src_allocr ? ggml_backend_name(get_allocr_backend(sched, src_allocr)) : "NULL"); j, src->name, src_backend ? ggml_backend_name(src_backend) : "NULL");
} }
if (src->view_src != NULL && src_allocr != node_allocr(src->view_src)) { if (src->view_src != NULL && src_backend != tensor_backend(src->view_src)) {
fprintf(stderr, "!!!!!!! [src] %s has backend %s, view_src %s has backend %s\n", fprintf(stderr, "!!!!!!! [src] %s has backend %s, view_src %s has backend %s\n",
src->name, src_allocr ? ggml_backend_name(get_allocr_backend(sched, src_allocr)) : "NULL", src->name, src_backend ? ggml_backend_name(src_backend) : "NULL",
src->view_src->name, node_allocr(src->view_src) ? ggml_backend_name(get_allocr_backend(sched, node_allocr(src->view_src))) : "NULL"); src->view_src->name, tensor_backend(src->view_src) ? ggml_backend_name(tensor_backend(src->view_src)) : "NULL");
} }
} }
} }
@ -1406,32 +1377,43 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
struct ggml_backend_sched_split * split = &sched->splits[i]; struct ggml_backend_sched_split * split = &sched->splits[i];
split->graph = ggml_graph_view(graph, split->i_start, split->i_end); split->graph = ggml_graph_view(graph, split->i_start, split->i_end);
// add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split
for (int j = 0; j < split->n_inputs; j++) { for (int j = 0; j < split->n_inputs; j++) {
struct ggml_tensor * input = split->inputs[j]; struct ggml_tensor * input = split->inputs[j];
struct ggml_tensor * input_cpy = sched->node_copies[hash_id(input)][sched_allocr_prio(sched, split->tallocr)]; struct ggml_tensor * input_cpy = sched->tensor_copies[hash_id(input)][split->backend_id];
// add a dependency to the input source so that it is not freed before the copy is done // add a dependency to the input source so that it is not freed before the copy is done
GGML_ASSERT(input_cpy->src[0] == NULL || input_cpy->src[0] == input); struct ggml_tensor * input_dep = ggml_view_tensor(sched->ctx, input);
input_cpy->src[0] = input; sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(input);
graph_copy->nodes[graph_copy->n_nodes++] = input_dep;
// add a dependency to the input copy so that it is allocated at the start of the split
sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id;
graph_copy->nodes[graph_copy->n_nodes++] = input_cpy; graph_copy->nodes[graph_copy->n_nodes++] = input_cpy;
} }
for (int j = split->i_start; j < split->i_end; j++) { for (int j = split->i_start; j < split->i_end; j++) {
sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]);
graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j]; graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j];
} }
} }
sched->graph = graph_copy; sched->graph = graph_copy;
} }
static void sched_alloc_splits(ggml_backend_sched_t sched) { static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
ggml_gallocr_alloc_graph_n( // ggml_gallocr_reserve_n(sched->galloc, sched->graph, sched->node_backend_ids);
sched->galloc, if (!ggml_gallocr_alloc_graph(sched->galloc, sched->graph)) {
sched->graph, #ifndef NDEBUG
sched->hash_set, fprintf(stderr, "ggml_backend_sched: failed to allocate graph, reserving\n");
sched->node_talloc); #endif
ggml_gallocr_reserve_n(sched->galloc, sched->graph, sched->node_backend_ids);
if (!ggml_gallocr_alloc_graph(sched->galloc, sched->graph)) {
fprintf(stderr, "ggml_backend_sched: failed to allocate graph\n");
return false;
}
}
} }
static void sched_compute_splits(ggml_backend_sched_t sched) { static bool ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
uint64_t copy_us[GGML_MAX_BACKENDS] = {0}; uint64_t copy_us[GGML_MAX_BACKENDS] = {0};
uint64_t compute_us[GGML_MAX_BACKENDS] = {0}; uint64_t compute_us[GGML_MAX_BACKENDS] = {0};
@ -1439,20 +1421,18 @@ static void sched_compute_splits(ggml_backend_sched_t sched) {
for (int i = 0; i < sched->n_splits; i++) { for (int i = 0; i < sched->n_splits; i++) {
struct ggml_backend_sched_split * split = &splits[i]; struct ggml_backend_sched_split * split = &splits[i];
ggml_backend_t split_backend = get_allocr_backend(sched, split->tallocr); int split_backend_id = split->backend_id;
int split_backend_id = sched_backend_prio(sched, split_backend); ggml_backend_t split_backend = sched->backends[split_backend_id];
// copy the input tensors to the split backend // copy the input tensors to the split backend
uint64_t copy_start_us = ggml_time_us(); uint64_t copy_start_us = ggml_time_us();
for (int j = 0; j < split->n_inputs; j++) { for (int j = 0; j < split->n_inputs; j++) {
struct ggml_tensor * input = split->inputs[j]; struct ggml_tensor * input = split->inputs[j];
struct ggml_tensor * input_cpy = sched->node_copies[hash_id(input)][split_backend_id]; struct ggml_tensor * input_cpy = sched->tensor_copies[hash_id(input)][split_backend_id];
GGML_ASSERT(input->buffer != NULL); GGML_ASSERT(input->buffer != NULL);
GGML_ASSERT(input_cpy->buffer != NULL); GGML_ASSERT(input_cpy->buffer != NULL);
// TODO: avoid this copy if it was already copied in a previous split, and the input didn't change
// this is important to avoid copying constants such as KQ_mask and inp_pos multiple times
ggml_backend_tensor_copy_async(split_backend, input, input_cpy); ggml_backend_tensor_copy_async(split_backend, input, input_cpy);
} }
//ggml_backend_synchronize(split_backend); // necessary to measure copy time //ggml_backend_synchronize(split_backend); // necessary to measure copy time
@ -1468,7 +1448,9 @@ static void sched_compute_splits(ggml_backend_sched_t sched) {
uint64_t compute_start_us = ggml_time_us(); uint64_t compute_start_us = ggml_time_us();
if (!sched->callback_eval) { if (!sched->callback_eval) {
ggml_backend_graph_compute(split_backend, &split->graph); if (!ggml_backend_graph_compute(split_backend, &split->graph)) {
return false;
}
//ggml_backend_synchronize(split_backend); // necessary to measure compute time //ggml_backend_synchronize(split_backend); // necessary to measure compute time
} else { } else {
// similar to ggml_backend_compare_graph_backend // similar to ggml_backend_compare_graph_backend
@ -1488,7 +1470,9 @@ static void sched_compute_splits(ggml_backend_sched_t sched) {
struct ggml_cgraph gv = ggml_graph_view(&split->graph, j0, j1 + 1); struct ggml_cgraph gv = ggml_graph_view(&split->graph, j0, j1 + 1);
ggml_backend_graph_compute(split_backend, &gv); if (!ggml_backend_graph_compute(split_backend, &gv)) {
return false;
}
if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) { if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) {
break; break;
@ -1510,19 +1494,8 @@ static void sched_compute_splits(ggml_backend_sched_t sched) {
} }
} }
#endif #endif
}
static void sched_reset(ggml_backend_sched_t sched) { return true;
for (int i = 0; i < sched->n_backends; i++) {
ggml_tallocr_reset(sched->tallocs[i]);
}
// reset state for the next run
size_t hash_size = sched->hash_set.size;
memset(sched->hash_set.keys, 0, sizeof(sched->hash_set.keys[0]) * hash_size);
memset(sched->node_talloc, 0, sizeof(sched->node_talloc[0]) * hash_size);
memset(sched->node_copies, 0, sizeof(sched->node_copies[0]) * hash_size);
sched->is_reset = true;
} }
ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size) { ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size) {
@ -1533,8 +1506,9 @@ ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_back
// initialize hash table // initialize hash table
sched->hash_set = ggml_hash_set_new(graph_size + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS); sched->hash_set = ggml_hash_set_new(graph_size + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS);
sched->node_talloc = calloc(sizeof(sched->node_talloc[0]) * sched->hash_set.size, 1); sched->tensor_backend_id = calloc(sizeof(sched->tensor_backend_id[0]), sched->hash_set.size);
sched->node_copies = calloc(sizeof(sched->node_copies[0]) * sched->hash_set.size, 1); sched->tensor_copies = calloc(sizeof(sched->tensor_copies[0]), sched->hash_set.size);
sched->node_backend_ids = calloc(sizeof(sched->node_backend_ids[0]), graph_size);
sched->n_backends = n_backends; sched->n_backends = n_backends;
for (int i = 0; i < n_backends; i++) { for (int i = 0; i < n_backends; i++) {
@ -1542,14 +1516,9 @@ ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_back
sched->bufts[i] = bufts ? bufts[i] : ggml_backend_get_default_buffer_type(backends[i]); sched->bufts[i] = bufts ? bufts[i] : ggml_backend_get_default_buffer_type(backends[i]);
} }
sched->galloc = ggml_gallocr_new(); sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
// init measure allocs for each backend ggml_backend_sched_reset(sched);
for (int i = 0; i < n_backends; i++) {
sched->tallocs[i] = ggml_tallocr_new_measure_from_buft(sched->bufts[i]);
}
sched_reset(sched);
return sched; return sched;
} }
@ -1558,49 +1527,54 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
if (sched == NULL) { if (sched == NULL) {
return; return;
} }
for (int i = 0; i < sched->n_backends; i++) {
ggml_tallocr_free(sched->tallocs[i]);
}
ggml_gallocr_free(sched->galloc); ggml_gallocr_free(sched->galloc);
ggml_free(sched->ctx); ggml_free(sched->ctx);
free(sched->hash_set.keys); free(sched->hash_set.keys);
free(sched->node_talloc); free(sched->tensor_backend_id);
free(sched->node_copies); free(sched->tensor_copies);
free(sched->node_backend_ids);
free(sched); free(sched);
} }
void ggml_backend_sched_init_measure(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph) { void ggml_backend_sched_reset(ggml_backend_sched_t sched) {
GGML_ASSERT(ggml_tallocr_is_measure(sched->tallocs[0])); // can only be initialized once // reset state for the next run
size_t hash_size = sched->hash_set.size;
memset(sched->hash_set.keys, 0, sizeof(sched->hash_set.keys[0]) * hash_size); // NOLINT
memset(sched->tensor_backend_id, -1, sizeof(sched->tensor_backend_id[0]) * hash_size);
memset(sched->tensor_copies, 0, sizeof(sched->tensor_copies[0]) * hash_size);
sched_split_graph(sched, measure_graph); sched->is_reset = true;
sched_alloc_splits(sched);
// allocate buffers and reset allocators
for (int i = 0; i < sched->n_backends; i++) {
size_t size = ggml_tallocr_max_size(sched->tallocs[i]);
ggml_tallocr_free(sched->tallocs[i]);
sched->tallocs[i] = ggml_tallocr_new_from_buft(sched->bufts[i], size);
} }
sched_reset(sched); bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph) {
ggml_backend_sched_split_graph(sched, measure_graph);
if (!ggml_gallocr_reserve_n(sched->galloc, sched->graph, sched->node_backend_ids)) {
return false;
} }
void ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { ggml_backend_sched_reset(sched);
return true;
}
bool ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) {
GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS); GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + GGML_MAX_SPLITS*GGML_MAX_SPLIT_INPUTS);
if (!sched->is_reset) { if (!sched->is_reset) {
sched_reset(sched); ggml_backend_sched_reset(sched);
} }
sched_split_graph(sched, graph); ggml_backend_sched_split_graph(sched, graph);
sched_alloc_splits(sched); if (!ggml_backend_sched_alloc_splits(sched)) {
sched_compute_splits(sched); return false;
} }
void ggml_backend_sched_reset(ggml_backend_sched_t sched) { if (!ggml_backend_sched_compute_splits(sched)) {
sched_reset(sched); return false;
} }
return true;
}
void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data) { void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data) {
sched->callback_eval = callback; sched->callback_eval = callback;
@ -1611,37 +1585,30 @@ int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) {
return sched->n_splits; return sched->n_splits;
} }
ggml_tallocr_t ggml_backend_sched_get_tallocr(ggml_backend_sched_t sched, ggml_backend_t backend) { size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend) {
int backend_index = sched_backend_prio(sched, backend); int backend_index = ggml_backend_sched_backend_id(sched, backend);
GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
return sched->tallocs[backend_index]; return ggml_gallocr_get_buffer_size(sched->galloc, backend_index);
}
ggml_backend_buffer_t ggml_backend_sched_get_buffer(ggml_backend_sched_t sched, ggml_backend_t backend) {
int backend_index = sched_backend_prio(sched, backend);
GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
return ggml_tallocr_get_buffer(sched->tallocs[backend_index]);
} }
void ggml_backend_sched_set_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) { void ggml_backend_sched_set_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend) {
int backend_index = sched_backend_prio(sched, backend); int backend_index = ggml_backend_sched_backend_id(sched, backend);
GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends); GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
node_allocr(node) = sched->tallocs[backend_index]; tensor_backend_id(node) = backend_index;
} }
ggml_backend_t ggml_backend_sched_get_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) { ggml_backend_t ggml_backend_sched_get_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node) {
ggml_tallocr_t allocr = node_allocr(node); int backend_index = tensor_backend_id(node);
if (allocr == NULL) { if (backend_index == -1) {
return NULL; return NULL;
} }
return get_allocr_backend(sched, allocr); return sched->backends[backend_index];
} }
// utils // utils
void ggml_backend_view_init(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { void ggml_backend_view_init(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) {
GGML_ASSERT(tensor->buffer == NULL); GGML_ASSERT(tensor->buffer == NULL);
//GGML_ASSERT(tensor->data == NULL); // views of pre-allocated tensors may have the data set in ggml_new_tensor, but still need to be initialized by the backend
GGML_ASSERT(tensor->view_src != NULL); GGML_ASSERT(tensor->view_src != NULL);
GGML_ASSERT(tensor->view_src->buffer != NULL); GGML_ASSERT(tensor->view_src->buffer != NULL);
GGML_ASSERT(tensor->view_src->data != NULL); GGML_ASSERT(tensor->view_src->data != NULL);
@ -1665,7 +1632,7 @@ void ggml_backend_tensor_alloc(ggml_backend_buffer_t buffer, struct ggml_tensor
ggml_backend_buffer_init_tensor(buffer, tensor); ggml_backend_buffer_init_tensor(buffer, tensor);
} }
static struct ggml_tensor * graph_dup_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies, static struct ggml_tensor * graph_copy_dup_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies,
struct ggml_context * ctx_allocated, struct ggml_context * ctx_unallocated, struct ggml_tensor * src) { struct ggml_context * ctx_allocated, struct ggml_context * ctx_unallocated, struct ggml_tensor * src) {
GGML_ASSERT(src != NULL); GGML_ASSERT(src != NULL);
@ -1678,7 +1645,7 @@ static struct ggml_tensor * graph_dup_tensor(struct ggml_hash_set hash_set, stru
struct ggml_tensor * dst = ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src); struct ggml_tensor * dst = ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src);
if (src->view_src != NULL) { if (src->view_src != NULL) {
dst->view_src = graph_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src); dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src);
dst->view_offs = src->view_offs; dst->view_offs = src->view_offs;
} }
dst->op = src->op; dst->op = src->op;
@ -1691,14 +1658,14 @@ static struct ggml_tensor * graph_dup_tensor(struct ggml_hash_set hash_set, stru
if (s == NULL) { if (s == NULL) {
break; break;
} }
dst->src[i] = graph_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s); dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s);
} }
node_copies[id] = dst; node_copies[id] = dst;
return dst; return dst;
} }
static void graph_init_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies, bool * node_init, struct ggml_tensor * src) { static void graph_copy_init_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies, bool * node_init, struct ggml_tensor * src) {
size_t id = ggml_hash_find(hash_set, src); size_t id = ggml_hash_find(hash_set, src);
if (node_init[id]) { if (node_init[id]) {
return; return;
@ -1707,7 +1674,7 @@ static void graph_init_tensor(struct ggml_hash_set hash_set, struct ggml_tensor
struct ggml_tensor * dst = node_copies[id]; struct ggml_tensor * dst = node_copies[id];
if (dst->view_src != NULL) { if (dst->view_src != NULL) {
graph_init_tensor(hash_set, node_copies, node_init, src->view_src); graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src);
ggml_backend_view_init(dst->view_src->buffer, dst); ggml_backend_view_init(dst->view_src->buffer, dst);
} }
else { else {
@ -1720,17 +1687,17 @@ static void graph_init_tensor(struct ggml_hash_set hash_set, struct ggml_tensor
if (s == NULL) { if (s == NULL) {
break; break;
} }
graph_init_tensor(hash_set, node_copies, node_init, s); graph_copy_init_tensor(hash_set, node_copies, node_init, s);
} }
} }
struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph) { struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph) {
struct ggml_hash_set hash_set = { struct ggml_hash_set hash_set = {
/* .size = */ graph->visited_hash_table.size, /* .size = */ graph->visited_hash_table.size,
/* .keys = */ calloc(sizeof(hash_set.keys[0]) * graph->visited_hash_table.size, 1) /* .keys = */ calloc(sizeof(hash_set.keys[0]), graph->visited_hash_table.size) // NOLINT
}; };
struct ggml_tensor ** node_copies = calloc(sizeof(node_copies[0]) * hash_set.size, 1); struct ggml_tensor ** node_copies = calloc(sizeof(node_copies[0]), hash_set.size); // NOLINT
bool * node_init = calloc(sizeof(node_init[0]) * hash_set.size, 1); bool * node_init = calloc(sizeof(node_init[0]), hash_set.size);
struct ggml_init_params params = { struct ggml_init_params params = {
/* .mem_size = */ ggml_tensor_overhead()*hash_set.size + ggml_graph_overhead_custom(graph->size, false), /* .mem_size = */ ggml_tensor_overhead()*hash_set.size + ggml_graph_overhead_custom(graph->size, false),
@ -1759,7 +1726,7 @@ struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, s
// dup nodes // dup nodes
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
graph_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node); graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node);
} }
// allocate nodes // allocate nodes
@ -1784,7 +1751,7 @@ struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, s
// copy data and init views // copy data and init views
for (int i = 0; i < graph->n_nodes; i++) { for (int i = 0; i < graph->n_nodes; i++) {
struct ggml_tensor * node = graph->nodes[i]; struct ggml_tensor * node = graph->nodes[i];
graph_init_tensor(hash_set, node_copies, node_init, node); graph_copy_init_tensor(hash_set, node_copies, node_init, node);
} }
// build graph copy // build graph copy

15
ggml-backend.h
View File

@ -130,11 +130,7 @@ extern "C" {
// in build_graph: // in build_graph:
build_graph(...) { build_graph(...) {
// allocating tensors in a specific backend (optional, recommended: pre-allocate inputs in a different buffer) // manually assign nodes to a backend (optional, should not be needed in most cases)
alloc_cpu = ggml_backend_sched_get_allocr(sched, backend_cpu);
ggml_allocr_alloc(alloc_cpu, tensor);
// manually assigning nodes to a backend (optional, shouldn't be needed in most cases)
struct ggml_tensor * node = ggml_mul_mat(ctx, ...); struct ggml_tensor * node = ggml_mul_mat(ctx, ...);
ggml_backend_sched_set_node_backend(sched, node, backend_gpu); ggml_backend_sched_set_node_backend(sched, node, backend_gpu);
} }
@ -164,20 +160,19 @@ extern "C" {
GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size); GGML_API ggml_backend_sched_t ggml_backend_sched_new(ggml_backend_t * backends, ggml_backend_buffer_type_t * bufts, int n_backends, size_t graph_size);
GGML_API void ggml_backend_sched_free(ggml_backend_sched_t sched); GGML_API void ggml_backend_sched_free(ggml_backend_sched_t sched);
// Initialize backend buffers from a measure graph // Initialize backend buffers from a measure graph
GGML_API void ggml_backend_sched_init_measure(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph); GGML_API bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph);
// Get the number of splits of the last graph // Get the number of splits of the last graph
GGML_API int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched); GGML_API int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched);
GGML_API ggml_tallocr_t ggml_backend_sched_get_tallocr(ggml_backend_sched_t sched, ggml_backend_t backend); GGML_API size_t ggml_backend_sched_get_buffer_size(ggml_backend_sched_t sched, ggml_backend_t backend);
GGML_API ggml_backend_buffer_t ggml_backend_sched_get_buffer (ggml_backend_sched_t sched, ggml_backend_t backend);
GGML_API void ggml_backend_sched_set_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend); GGML_API void ggml_backend_sched_set_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node, ggml_backend_t backend);
GGML_API ggml_backend_t ggml_backend_sched_get_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node); GGML_API ggml_backend_t ggml_backend_sched_get_node_backend(ggml_backend_sched_t sched, struct ggml_tensor * node);
// Allocate and compute graph on the backend scheduler // Allocate and compute graph on the backend scheduler
GGML_API void ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph); GGML_API bool ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph);
// Reset all assignments and allocators - must be called before using the sched allocators to allocate inputs // Reset all assignments and allocators - must be called before changing the node backends
GGML_API void ggml_backend_sched_reset(ggml_backend_sched_t sched); GGML_API void ggml_backend_sched_reset(ggml_backend_sched_t sched);
// Set a callback to be called for each resulting node during graph compute // Set a callback to be called for each resulting node during graph compute

28
ggml.c
View File

@ -2607,7 +2607,7 @@ static struct ggml_tensor * ggml_new_tensor_impl(
/*.nb =*/ { 0, 0, 0, 0 }, /*.nb =*/ { 0, 0, 0, 0 },
/*.op =*/ GGML_OP_NONE, /*.op =*/ GGML_OP_NONE,
/*.op_params =*/ { 0 }, /*.op_params =*/ { 0 },
/*.is_param =*/ false, /*.flags =*/ 0,
/*.grad =*/ NULL, /*.grad =*/ NULL,
/*.src =*/ { NULL }, /*.src =*/ { NULL },
/*.perf_runs =*/ 0, /*.perf_runs =*/ 0,
@ -6509,7 +6509,7 @@ struct ggml_tensor * ggml_cross_entropy_loss_back(
void ggml_set_param( void ggml_set_param(
struct ggml_context * ctx, struct ggml_context * ctx,
struct ggml_tensor * tensor) { struct ggml_tensor * tensor) {
tensor->is_param = true; tensor->flags |= GGML_TENSOR_FLAG_PARAM;
GGML_ASSERT(tensor->grad == NULL); GGML_ASSERT(tensor->grad == NULL);
tensor->grad = ggml_dup_tensor(ctx, tensor); tensor->grad = ggml_dup_tensor(ctx, tensor);
@ -15311,7 +15311,7 @@ static struct ggml_tensor * ggml_recompute_graph_node(
return NULL; return NULL;
} }
if (node->is_param) { if (node->flags & GGML_TENSOR_FLAG_PARAM) {
return node; return node;
} }
@ -15345,7 +15345,7 @@ static struct ggml_tensor * ggml_recompute_graph_node(
clone->op = node->op; clone->op = node->op;
clone->grad = node->grad; clone->grad = node->grad;
clone->is_param = node->is_param; clone->flags = node->flags;
clone->extra = node->extra; clone->extra = node->extra;
for (int k = 0; k < GGML_MAX_DIMS; ++k) { for (int k = 0; k < GGML_MAX_DIMS; ++k) {
clone->nb[k] = node->nb[k]; clone->nb[k] = node->nb[k];
@ -16377,7 +16377,7 @@ void ggml_build_backward_expand(struct ggml_context * ctx, struct ggml_cgraph *
for (int i = 0; i < gf->n_nodes; i++) { for (int i = 0; i < gf->n_nodes; i++) {
struct ggml_tensor * node = gf->nodes[i]; struct ggml_tensor * node = gf->nodes[i];
if (node->is_param) { if (node->flags & GGML_TENSOR_FLAG_PARAM) {
GGML_PRINT_DEBUG("%s: found root node %p\n", __func__, (void *) node); GGML_PRINT_DEBUG("%s: found root node %p\n", __func__, (void *) node);
ggml_build_forward_expand(gb, node->grad); ggml_build_forward_expand(gb, node->grad);
} }
@ -17862,7 +17862,7 @@ void ggml_graph_print(const struct ggml_cgraph * cgraph) {
GGML_PRINT(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s (%3d) cpu = %7.3f / %7.3f ms, wall = %7.3f / %7.3f ms\n", GGML_PRINT(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s (%3d) cpu = %7.3f / %7.3f ms, wall = %7.3f / %7.3f ms\n",
i, i,
node->ne[0], node->ne[1], node->ne[2], node->ne[0], node->ne[1], node->ne[2],
ggml_op_name(node->op), node->is_param ? "x" : node->grad ? "g" : " ", node->perf_runs, ggml_op_name(node->op), (node->flags & GGML_TENSOR_FLAG_PARAM) ? "x" : node->grad ? "g" : " ", node->perf_runs,
(double) node->perf_cycles / (double) ggml_cycles_per_ms(), (double) node->perf_cycles / (double) ggml_cycles_per_ms(),
(double) node->perf_cycles / (double) ggml_cycles_per_ms() / (double) node->perf_runs, (double) node->perf_cycles / (double) ggml_cycles_per_ms() / (double) node->perf_runs,
(double) node->perf_time_us / 1000.0, (double) node->perf_time_us / 1000.0,
@ -17955,7 +17955,7 @@ void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph
continue; continue;
} }
if (node->is_param) { if (node->flags & GGML_TENSOR_FLAG_PARAM) {
snprintf(color, sizeof(color), "yellow"); snprintf(color, sizeof(color), "yellow");
} else if (node->grad) { } else if (node->grad) {
if (ggml_graph_find(gf, node)) { if (ggml_graph_find(gf, node)) {
@ -18129,7 +18129,7 @@ static enum ggml_opt_result ggml_opt_adam(
int np = 0; int np = 0;
int64_t nx = 0; int64_t nx = 0;
for (int i = 0; i < gf->n_nodes; ++i) { for (int i = 0; i < gf->n_nodes; ++i) {
if (gf->nodes[i]->is_param) { if (gf->nodes[i]->flags & GGML_TENSOR_FLAG_PARAM) {
GGML_PRINT_DEBUG("found param %d: grad->op = %d\n", np, gf->nodes[i]->grad->op); GGML_PRINT_DEBUG("found param %d: grad->op = %d\n", np, gf->nodes[i]->grad->op);
GGML_ASSERT(np < GGML_MAX_PARAMS); GGML_ASSERT(np < GGML_MAX_PARAMS);
@ -18492,7 +18492,7 @@ static enum ggml_opt_result ggml_opt_lbfgs(
int np = 0; int np = 0;
int nx = 0; int nx = 0;
for (int i = 0; i < gf->n_nodes; ++i) { for (int i = 0; i < gf->n_nodes; ++i) {
if (gf->nodes[i]->is_param) { if (gf->nodes[i]->flags & GGML_TENSOR_FLAG_PARAM) {
GGML_PRINT_DEBUG("found param %d: grad->op = %d\n", np, gf->nodes[i]->grad->op); GGML_PRINT_DEBUG("found param %d: grad->op = %d\n", np, gf->nodes[i]->grad->op);
GGML_ASSERT(np < GGML_MAX_PARAMS); GGML_ASSERT(np < GGML_MAX_PARAMS);
@ -18967,6 +18967,16 @@ enum ggml_opt_result ggml_opt_resume_g(
//////////////////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////////
void ggml_set_input(struct ggml_tensor * tensor) {
tensor->flags |= GGML_TENSOR_FLAG_INPUT;
}
void ggml_set_output(struct ggml_tensor * tensor) {
tensor->flags |= GGML_TENSOR_FLAG_OUTPUT;
}
////////////////////////////////////////////////////////////////////////////////
void ggml_quantize_init(enum ggml_type type) { void ggml_quantize_init(enum ggml_type type) {
ggml_critical_section_start(); ggml_critical_section_start();

14
ggml.h
View File

@ -510,6 +510,12 @@ extern "C" {
GGML_LOG_LEVEL_DEBUG = 5 GGML_LOG_LEVEL_DEBUG = 5
}; };
enum ggml_tensor_flag {
GGML_TENSOR_FLAG_INPUT = 1,
GGML_TENSOR_FLAG_OUTPUT = 2,
GGML_TENSOR_FLAG_PARAM = 4,
};
// ggml object // ggml object
struct ggml_object { struct ggml_object {
size_t offs; size_t offs;
@ -543,7 +549,7 @@ extern "C" {
// op params - allocated as int32_t for alignment // op params - allocated as int32_t for alignment
int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)];
bool is_param; int32_t flags;
struct ggml_tensor * grad; struct ggml_tensor * grad;
struct ggml_tensor * src[GGML_MAX_SRC]; struct ggml_tensor * src[GGML_MAX_SRC];
@ -2092,6 +2098,12 @@ extern "C" {
ggml_opt_callback callback, ggml_opt_callback callback,
void * callback_data); void * callback_data);
//
// tensor flags
//
GGML_API void ggml_set_input(struct ggml_tensor * tensor);
GGML_API void ggml_set_output(struct ggml_tensor * tensor);
// //
// quantization // quantization
// //

392
whisper.cpp
View File

@ -471,52 +471,32 @@ struct whisper_pair {
// ggml_allocr wrapper for whisper usage // ggml_allocr wrapper for whisper usage
struct whisper_allocr { struct whisper_allocr {
ggml_allocr * alloc = nullptr; ggml_gallocr_t alloc = nullptr;
std::vector<uint8_t> meta; std::vector<uint8_t> meta;
ggml_backend_buffer_t buffer;
}; };
static size_t whisper_allocr_size(struct whisper_allocr & allocr) { static size_t whisper_allocr_size(struct whisper_allocr & allocr) {
return allocr.meta.size() + ggml_allocr_max_size(allocr.alloc); return allocr.meta.size() + ggml_gallocr_get_buffer_size(allocr.alloc, 0);
} }
// measure the memory usage of a graph and prepare the allocr's internal data buffer // measure the memory usage of a graph and prepare the allocr's internal data buffer
static void whisper_allocr_graph_init(struct whisper_allocr & allocr, ggml_backend_t backend, std::function<struct ggml_cgraph *()> && get_graph) { static bool whisper_allocr_graph_init(struct whisper_allocr & allocr, ggml_backend_t backend, std::function<struct ggml_cgraph *()> && get_graph) {
auto & alloc = allocr.alloc; auto & alloc = allocr.alloc;
auto & meta = allocr.meta; auto & meta = allocr.meta;
alloc = ggml_allocr_new_measure_from_backend(backend); alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
meta.resize(ggml_tensor_overhead()*WHISPER_MAX_NODES + ggml_graph_overhead()); meta.resize(ggml_tensor_overhead()*WHISPER_MAX_NODES + ggml_graph_overhead());
ggml_allocr_alloc_graph(alloc, get_graph()); // since there are dependencies between the different graphs,
} // we need to allocate them instead of only reserving to get the correct compute buffer size
if (!ggml_gallocr_alloc_graph(alloc, get_graph())) {
static void whisper_allocr_graph_realloc(struct whisper_allocr & allocr, ggml_backend_t backend) { // failed to allocate the compute buffer
if (allocr.alloc == nullptr) { WHISPER_LOG_ERROR("%s: failed to allocate the compute buffer\n", __func__);
// this can be null if we use external encoder like CoreML or OpenVINO return false;
return;
}
auto & alloc = allocr.alloc;
auto & buffer = allocr.buffer;
size_t size = ggml_allocr_max_size(alloc);
ggml_allocr_free(alloc);
buffer = ggml_backend_alloc_buffer(backend, size);
alloc = ggml_allocr_new_from_buffer(buffer);
}
static void whisper_allocr_free(struct whisper_allocr & allocr) {
if (allocr.alloc) {
ggml_allocr_free(allocr.alloc);
ggml_backend_buffer_free(allocr.buffer);
allocr.alloc = nullptr;
} }
return true;
} }
// medium // medium
@ -658,9 +638,9 @@ struct whisper_kv_cache {
struct ggml_tensor * k; struct ggml_tensor * k;
struct ggml_tensor * v; struct ggml_tensor * v;
struct ggml_context * ctx; struct ggml_context * ctx = nullptr;
ggml_backend_buffer_t buffer; ggml_backend_buffer_t buffer = nullptr;
}; };
struct whisper_model { struct whisper_model {
@ -698,10 +678,10 @@ struct whisper_model {
std::vector<whisper_layer_decoder> layers_decoder; std::vector<whisper_layer_decoder> layers_decoder;
// ggml context that contains all the meta information about the model tensors // ggml context that contains all the meta information about the model tensors
struct ggml_context * ctx; struct ggml_context * ctx = nullptr;
// the model backend data is read-only and can be shared between processors // the model backend data is read-only and can be shared between processors
std::vector<struct ggml_backend_buffer *> buffers; ggml_backend_buffer_t buffer = nullptr;
// tensors // tensors
int n_loaded; int n_loaded;
@ -903,37 +883,27 @@ static bool kv_cache_init(
cache.ctx = ggml_init(params); cache.ctx = ggml_init(params);
if (!cache.ctx) { if (!cache.ctx) {
WHISPER_LOG_ERROR("%s: failed to allocate memory for kv cache\n", __func__); WHISPER_LOG_ERROR("%s: failed to allocate memory for the kv cache context\n", __func__);
return false; return false;
} }
cache.k = ggml_new_tensor_1d(cache.ctx, wtype, n_elements); cache.k = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
cache.v = ggml_new_tensor_1d(cache.ctx, wtype, n_elements); cache.v = ggml_new_tensor_1d(cache.ctx, wtype, n_elements);
const size_t mem_bytes = ggml_nbytes(cache.k) + ggml_nbytes(cache.v); cache.buffer = ggml_backend_alloc_ctx_tensors(cache.ctx, backend);
if (!cache.buffer) {
cache.buffer = ggml_backend_alloc_buffer(backend, mem_bytes); WHISPER_LOG_ERROR("%s: failed to allocate memory for the kv cache\n", __func__);
return false;
// allocate the tensors into the backend buffer
{
ggml_allocr * alloc = ggml_allocr_new_from_buffer(cache.buffer);
ggml_allocr_alloc(alloc, cache.k);
ggml_allocr_alloc(alloc, cache.v);
ggml_allocr_free(alloc);
} }
return true; return true;
} }
static void kv_cache_free(struct whisper_kv_cache & cache) { static void kv_cache_free(struct whisper_kv_cache & cache) {
if (cache.ctx) {
ggml_free(cache.ctx); ggml_free(cache.ctx);
ggml_backend_buffer_free(cache.buffer); ggml_backend_buffer_free(cache.buffer);
cache.ctx = nullptr; cache.ctx = nullptr;
} }
}
static bool whisper_kv_cache_find_slot( static bool whisper_kv_cache_find_slot(
struct whisper_kv_cache & cache, struct whisper_kv_cache & cache,
@ -1513,68 +1483,21 @@ static bool whisper_model_load(struct whisper_model_loader * loader, whisper_con
} }
wctx.backend = whisper_backend_init(wctx.params); wctx.backend = whisper_backend_init(wctx.params);
if (!wctx.backend) {
// some devices have a limit on the maximum size of single memory buffer WHISPER_LOG_ERROR("%s: failed to initialize the backend\n", __func__);
// for example, iPhones are limited to 1GB per buffer return false;
// to workaround this, we will allocate multiple buffers of smaller size and will split the tensors with the
// model weights between them
//
// the map_t2b maps tensor names to buffer indices
// as we iterate over the tensors, we will allocate new buffers when the current one is full
//
// finally, we create a separate allocator for each buffer and use it to allocate the tensors
// we keep the allocators alive until all the tensors are loaded
GGML_ASSERT(model.buffers.empty());
std::map<std::string, int> map_t2b;
{
size_t size_main = 0;
size_t size_cur = 0;
static const size_t GB = 1024ull*1024ull*1024ull;
for (const auto & t : model.tensors) {
const size_t cur = ggml_nbytes(t.second) + ggml_tensor_overhead();
// adding the tensor to the current buffer will exceed the limit, so we need to allocate a new buffer
if (size_cur + cur > GB) {
GGML_ASSERT(size_cur > 0 && "A tensor is too large to fit in a single buffer");
model.buffers.emplace_back(ggml_backend_alloc_buffer(wctx.backend, size_cur));
size_cur = cur;
}
map_t2b[t.first] = model.buffers.size();
size_cur += cur;
size_main += cur;
}
// allocate the last buffer if needed
if (size_cur > 0) {
model.buffers.emplace_back(ggml_backend_alloc_buffer(wctx.backend, size_cur));
}
GGML_ASSERT(model.buffers.size() > 0);
WHISPER_LOG_INFO("%s: %8s total size = %8.2f MB (%d buffers)\n", __func__, ggml_backend_name(wctx.backend), size_main / 1e6, (int) model.buffers.size());
}
std::vector<ggml_allocr *> allocs(model.buffers.size());
for (size_t i = 0; i < allocs.size(); ++i) {
allocs[i] = ggml_allocr_new_from_buffer(model.buffers[i]);
} }
// allocate tensors in the backend buffers // allocate tensors in the backend buffers
{ model.buffer = ggml_backend_alloc_ctx_tensors(model.ctx, wctx.backend);
for (const auto & t : model.tensors) { if (!model.buffer) {
ggml_allocr_alloc(allocs[map_t2b[t.first]], t.second); WHISPER_LOG_ERROR("%s: failed to allocate memory for the model\n", __func__);
} return false;
} }
size_t size_main = ggml_backend_buffer_get_size(model.buffer);
WHISPER_LOG_INFO("%s: %8s total size = %8.2f MB\n", __func__, ggml_backend_name(wctx.backend), size_main / 1e6);
// load weights // load weights
{ {
size_t total_size = 0; size_t total_size = 0;
@ -1636,15 +1559,11 @@ static bool whisper_model_load(struct whisper_model_loader * loader, whisper_con
return false; return false;
} }
ggml_backend_t backend = wctx.backend; //ggml_backend_t backend = wctx.backend;
//printf("%s: [%5.5s] %s\n", __func__, ggml_backend_name(backend), name.c_str()); //printf("%s: [%5.5s] %s\n", __func__, ggml_backend_name(backend), name.c_str());
if ((ggml_backend_is_cpu(backend) if (ggml_backend_buffer_is_host(model.buffer)) {
#ifdef GGML_USE_METAL
|| ggml_backend_is_metal(backend)
#endif
)) {
// for the CPU and Metal backend, we can read directly into the tensor // for the CPU and Metal backend, we can read directly into the tensor
loader->read(loader->context, tensor->data, ggml_nbytes(tensor)); loader->read(loader->context, tensor->data, ggml_nbytes(tensor));
BYTESWAP_TENSOR(tensor); BYTESWAP_TENSOR(tensor);
@ -1672,10 +1591,6 @@ static bool whisper_model_load(struct whisper_model_loader * loader, whisper_con
} }
} }
for (auto & alloc : allocs) {
ggml_allocr_free(alloc);
}
wctx.t_load_us = ggml_time_us() - t_start_us; wctx.t_load_us = ggml_time_us() - t_start_us;
return true; return true;
@ -1704,7 +1619,6 @@ static struct ggml_cgraph * whisper_build_graph_conv(
whisper_state & wstate, whisper_state & wstate,
const int mel_offset) { const int mel_offset) {
const auto & model = wctx.model; const auto & model = wctx.model;
const auto & mel_inp = wstate.mel;
const auto & hparams = model.hparams; const auto & hparams = model.hparams;
const int n_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : hparams.n_audio_ctx; const int n_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : hparams.n_audio_ctx;
@ -1722,31 +1636,9 @@ static struct ggml_cgraph * whisper_build_graph_conv(
ggml_cgraph * gf = ggml_new_graph(ctx0); ggml_cgraph * gf = ggml_new_graph(ctx0);
ggml_allocr * alloc = wstate.alloc_conv.alloc;
struct ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 2*n_ctx, n_mels); struct ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 2*n_ctx, n_mels);
ggml_allocr_alloc(alloc, mel); ggml_set_name(mel, "mel");
ggml_set_input(mel);
assert(mel->type == GGML_TYPE_F32);
if (!ggml_allocr_is_measure(alloc)) {
assert(mel_inp.n_mel == n_mels);
wstate.inp_mel.resize(ggml_nelements(mel));
float * dst = wstate.inp_mel.data();
memset(dst, 0, ggml_nbytes(mel));
const int i0 = std::min(mel_offset, mel_inp.n_len);
const int i1 = std::min(mel_offset + 2*n_ctx, mel_inp.n_len);
for (int j = 0; j < mel_inp.n_mel; ++j) {
for (int i = i0; i < i1; ++i) {
dst[j*2*n_ctx + (i - i0)] = mel_inp.data[j*mel_inp.n_len + i];
}
}
ggml_backend_tensor_set(mel, wstate.inp_mel.data(), 0, ggml_nelements(mel)*sizeof(float));
}
struct ggml_tensor * cur = nullptr; struct ggml_tensor * cur = nullptr;
@ -2138,11 +2030,39 @@ static bool whisper_encode_internal(
{ {
auto & alloc = wstate.alloc_conv.alloc; auto & alloc = wstate.alloc_conv.alloc;
ggml_allocr_reset(alloc);
ggml_cgraph * gf = whisper_build_graph_conv(wctx, wstate, mel_offset); ggml_cgraph * gf = whisper_build_graph_conv(wctx, wstate, mel_offset);
ggml_allocr_alloc_graph(alloc, gf); if (!ggml_gallocr_alloc_graph(alloc, gf)) {
// should never happen as we pre-allocate the memory
return false;
}
// set the input
{
const auto & mel_inp = wstate.mel;
const int n_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : wctx.model.hparams.n_audio_ctx;
struct ggml_tensor * mel = ggml_graph_get_tensor(gf, "mel");
assert(mel->type == GGML_TYPE_F32);
assert(mel_inp.n_mel == wctx.model.hparams.n_mels);
wstate.inp_mel.resize(ggml_nelements(mel));
float * dst = wstate.inp_mel.data();
memset(dst, 0, ggml_nbytes(mel));
const int i0 = std::min(mel_offset, mel_inp.n_len);
const int i1 = std::min(mel_offset + 2*n_ctx, mel_inp.n_len);
for (int j = 0; j < mel_inp.n_mel; ++j) {
for (int i = i0; i < i1; ++i) {
dst[j*2*n_ctx + (i - i0)] = mel_inp.data[j*mel_inp.n_len + i];
}
}
ggml_backend_tensor_set(mel, wstate.inp_mel.data(), 0, ggml_nelements(mel)*sizeof(float));
}
if (!whisper_encode_external(wstate)) { if (!whisper_encode_external(wstate)) {
if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) { if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) {
@ -2155,11 +2075,12 @@ static bool whisper_encode_internal(
if (!whisper_encode_external(wstate)) { if (!whisper_encode_external(wstate)) {
auto & alloc = wstate.alloc_encode.alloc; auto & alloc = wstate.alloc_encode.alloc;
ggml_allocr_reset(alloc);
ggml_cgraph * gf = whisper_build_graph_encoder(wctx, wstate); ggml_cgraph * gf = whisper_build_graph_encoder(wctx, wstate);
ggml_allocr_alloc_graph(alloc, gf); if (!ggml_gallocr_alloc_graph(alloc, gf)) {
// should never happen as we pre-allocate the memory
return false;
}
if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) { if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) {
return false; return false;
@ -2170,11 +2091,12 @@ static bool whisper_encode_internal(
{ {
auto & alloc = wstate.alloc_cross.alloc; auto & alloc = wstate.alloc_cross.alloc;
ggml_allocr_reset(alloc);
ggml_cgraph * gf = whisper_build_graph_cross(wctx, wstate); ggml_cgraph * gf = whisper_build_graph_cross(wctx, wstate);
ggml_allocr_alloc_graph(alloc, gf); if (!ggml_gallocr_alloc_graph(alloc, gf)) {
// should never happen as we pre-allocate the memory
return false;
}
if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) { if (!ggml_graph_compute_helper(wstate.backend, gf, n_threads)) {
return false; return false;
@ -2190,7 +2112,8 @@ static bool whisper_encode_internal(
static struct ggml_cgraph * whisper_build_graph_decoder( static struct ggml_cgraph * whisper_build_graph_decoder(
whisper_context & wctx, whisper_context & wctx,
whisper_state & wstate, whisper_state & wstate,
const whisper_batch & batch) { const whisper_batch & batch,
bool worst_case) {
const auto & model = wctx.model; const auto & model = wctx.model;
const auto & hparams = model.hparams; const auto & hparams = model.hparams;
@ -2198,8 +2121,6 @@ static struct ggml_cgraph * whisper_build_graph_decoder(
WHISPER_ASSERT(!!kv_self.ctx); WHISPER_ASSERT(!!kv_self.ctx);
ggml_allocr * alloc = wstate.alloc_decode.alloc;
const int n_ctx = kv_self.size; const int n_ctx = kv_self.size;
const int n_state = hparams.n_text_state; const int n_state = hparams.n_text_state;
const int n_head = hparams.n_text_head; const int n_head = hparams.n_text_head;
@ -2208,8 +2129,8 @@ static struct ggml_cgraph * whisper_build_graph_decoder(
const int n_tokens = batch.n_tokens; const int n_tokens = batch.n_tokens;
const int n_audio_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : hparams.n_audio_ctx; const int n_audio_ctx = wstate.exp_n_audio_ctx > 0 ? wstate.exp_n_audio_ctx : hparams.n_audio_ctx;
const int32_t n_kv = ggml_allocr_is_measure(alloc) ? n_ctx : kv_self.n; const int32_t n_kv = worst_case ? n_ctx : kv_self.n;
const int32_t kv_head = ggml_allocr_is_measure(alloc) ? n_ctx - n_tokens : kv_self.head; const int32_t kv_head = worst_case ? n_ctx - n_tokens : kv_self.head;
//WHISPER_LOG_DEBUG("%s: n_past = %d, n_tokens = %d, n_audio_ctx = %d, n_ctx = %d\n", __func__, n_past, n_tokens, n_audio_ctx, n_ctx); //WHISPER_LOG_DEBUG("%s: n_past = %d, n_tokens = %d, n_audio_ctx = %d, n_ctx = %d\n", __func__, n_past, n_tokens, n_audio_ctx, n_ctx);
@ -2224,48 +2145,18 @@ static struct ggml_cgraph * whisper_build_graph_decoder(
ggml_cgraph * gf = ggml_new_graph_custom(ctx0, WHISPER_MAX_NODES, false); ggml_cgraph * gf = ggml_new_graph_custom(ctx0, WHISPER_MAX_NODES, false);
struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_allocr_alloc(alloc, embd); ggml_set_name(embd, "embd");
ggml_set_input(embd);
if (!ggml_allocr_is_measure(alloc)) {
ggml_backend_tensor_set(embd, batch.token, 0, n_tokens*ggml_element_size(embd));
}
struct ggml_tensor * position = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); struct ggml_tensor * position = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_allocr_alloc(alloc, position); ggml_set_name(position, "position");
ggml_set_input(position);
if (!ggml_allocr_is_measure(alloc)) {
for (int i = 0; i < n_tokens; ++i) {
const int32_t val = batch.pos[i];
ggml_backend_tensor_set(position, &val, i*sizeof(int32_t), sizeof(int32_t));
}
}
const float KQscale = pow(float(n_state)/n_head, -0.25); const float KQscale = pow(float(n_state)/n_head, -0.25);
struct ggml_tensor * KQ_mask = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tokens, 1); struct ggml_tensor * KQ_mask = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tokens, 1);
ggml_allocr_alloc(alloc, KQ_mask); ggml_set_name(KQ_mask, "KQ_mask");
ggml_set_input(KQ_mask);
if (!ggml_allocr_is_measure(alloc)) {
wstate.inp_mask.resize(n_kv*n_tokens);
float * data = wstate.inp_mask.data();
memset(data, 0, ggml_nbytes(KQ_mask));
for (int h = 0; h < 1; ++h) {
for (int j = 0; j < n_tokens; ++j) {
const whisper_pos pos = batch.pos[j];
const whisper_seq_id seq_id = batch.seq_id[j][0];
for (int i = 0; i < n_kv; ++i) {
if (!kv_self.cells[i].has_seq_id(seq_id) || kv_self.cells[i].pos > pos) {
data[h*(n_kv*n_tokens) + j*n_kv + i] = -INFINITY;
}
}
}
}
ggml_backend_tensor_set(KQ_mask, wstate.inp_mask.data(), 0, ggml_nelements(KQ_mask)*sizeof(float));
}
// token encoding + position encoding // token encoding + position encoding
struct ggml_tensor * cur = struct ggml_tensor * cur =
@ -2592,11 +2483,53 @@ static bool whisper_decode_internal(
{ {
auto & alloc = wstate.alloc_decode.alloc; auto & alloc = wstate.alloc_decode.alloc;
ggml_allocr_reset(alloc); ggml_cgraph * gf = whisper_build_graph_decoder(wctx, wstate, batch, false);
ggml_cgraph * gf = whisper_build_graph_decoder(wctx, wstate, batch); if (!ggml_gallocr_alloc_graph(alloc, gf)) {
// should never happen as we pre-allocate the memory
return false;
}
ggml_allocr_alloc_graph(alloc, gf); // set the inputs
{
struct ggml_tensor * embd = ggml_graph_get_tensor(gf, "embd");
ggml_backend_tensor_set(embd, batch.token, 0, n_tokens*ggml_element_size(embd));
}
{
struct ggml_tensor * position = ggml_graph_get_tensor(gf, "position");
for (int i = 0; i < n_tokens; ++i) {
const int32_t val = batch.pos[i];
ggml_backend_tensor_set(position, &val, i*sizeof(int32_t), sizeof(int32_t));
}
}
{
struct ggml_tensor * KQ_mask = ggml_graph_get_tensor(gf, "KQ_mask");
auto & kv_self = wstate.kv_self;
const int32_t n_kv = kv_self.n;
wstate.inp_mask.resize(n_kv*n_tokens);
float * data = wstate.inp_mask.data();
memset(data, 0, ggml_nbytes(KQ_mask));
for (int h = 0; h < 1; ++h) {
for (int j = 0; j < n_tokens; ++j) {
const whisper_pos pos = batch.pos[j];
const whisper_seq_id seq_id = batch.seq_id[j][0];
for (int i = 0; i < n_kv; ++i) {
if (!kv_self.cells[i].has_seq_id(seq_id) || kv_self.cells[i].pos > pos) {
data[h*(n_kv*n_tokens) + j*n_kv + i] = -INFINITY;
}
}
}
}
ggml_backend_tensor_set(KQ_mask, wstate.inp_mask.data(), 0, ggml_nelements(KQ_mask)*sizeof(float));
}
logits = gf->nodes[gf->n_nodes - 1]; logits = gf->nodes[gf->n_nodes - 1];
@ -3046,6 +2979,11 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
whisper_state * state = new whisper_state; whisper_state * state = new whisper_state;
state->backend = whisper_backend_init(ctx->params); state->backend = whisper_backend_init(ctx->params);
if (!state->backend) {
WHISPER_LOG_ERROR("%s: whisper_backend_init() failed\n", __func__);
whisper_free_state(state);
return nullptr;
}
// at this point, we don't know yet how many decoders will be used, so we overallocate 3x ctx // at this point, we don't know yet how many decoders will be used, so we overallocate 3x ctx
// in theory, there can be a case where this is not enough, but in practice it should always be enough // in theory, there can be a case where this is not enough, but in practice it should always be enough
@ -3053,7 +2991,7 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
if (!kv_cache_init(ctx->model.hparams, state->kv_self, ctx->backend, ctx->itype, factor*ctx->model.hparams.n_text_ctx)) { if (!kv_cache_init(ctx->model.hparams, state->kv_self, ctx->backend, ctx->itype, factor*ctx->model.hparams.n_text_ctx)) {
WHISPER_LOG_ERROR("%s: kv_cache_init() failed for self-attention cache\n", __func__); WHISPER_LOG_ERROR("%s: kv_cache_init() failed for self-attention cache\n", __func__);
delete state; whisper_free_state(state);
return nullptr; return nullptr;
} }
@ -3064,7 +3002,7 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
if (!kv_cache_init(ctx->model.hparams, state->kv_cross, ctx->backend, ctx->itype, ctx->model.hparams.n_audio_ctx)) { if (!kv_cache_init(ctx->model.hparams, state->kv_cross, ctx->backend, ctx->itype, ctx->model.hparams.n_audio_ctx)) {
WHISPER_LOG_ERROR("%s: kv_cache_init() failed for cross-attention cache\n", __func__); WHISPER_LOG_ERROR("%s: kv_cache_init() failed for cross-attention cache\n", __func__);
delete state; whisper_free_state(state);
return nullptr; return nullptr;
} }
@ -3083,7 +3021,7 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
if (!state->ctx_coreml) { if (!state->ctx_coreml) {
WHISPER_LOG_ERROR("%s: failed to load Core ML model from '%s'\n", __func__, path_coreml.c_str()); WHISPER_LOG_ERROR("%s: failed to load Core ML model from '%s'\n", __func__, path_coreml.c_str());
#ifndef WHISPER_COREML_ALLOW_FALLBACK #ifndef WHISPER_COREML_ALLOW_FALLBACK
delete state; whisper_free_state(state);
return nullptr; return nullptr;
#endif #endif
} else { } else {
@ -3107,37 +3045,55 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
// conv allocator // conv allocator
{ {
whisper_allocr_graph_init(state->alloc_conv, ctx->backend, bool ok = whisper_allocr_graph_init(state->alloc_conv, ctx->backend,
[&]() { [&]() {
return whisper_build_graph_conv(*ctx, *state, 0); return whisper_build_graph_conv(*ctx, *state, 0);
}); });
if (!ok) {
WHISPER_LOG_ERROR("%s: failed to init conv allocator\n", __func__);
whisper_free_state(state);
return nullptr;
}
WHISPER_LOG_INFO("%s: compute buffer (conv) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_conv) / 1e6); WHISPER_LOG_INFO("%s: compute buffer (conv) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_conv) / 1e6);
} }
// encoder allocator // encoder allocator
if (!whisper_encode_external(*state)) { if (!whisper_encode_external(*state)) {
whisper_allocr_graph_init(state->alloc_encode, ctx->backend, bool ok = whisper_allocr_graph_init(state->alloc_encode, ctx->backend,
[&]() { [&]() {
return whisper_build_graph_encoder(*ctx, *state); return whisper_build_graph_encoder(*ctx, *state);
}); });
if (!ok) {
WHISPER_LOG_ERROR("%s: failed to init encoder allocator\n", __func__);
whisper_free_state(state);
return nullptr;
}
WHISPER_LOG_INFO("%s: compute buffer (encode) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_encode) / 1e6); WHISPER_LOG_INFO("%s: compute buffer (encode) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_encode) / 1e6);
} }
// cross allocator // cross allocator
{ {
whisper_allocr_graph_init(state->alloc_cross, ctx->backend, bool ok = whisper_allocr_graph_init(state->alloc_cross, ctx->backend,
[&]() { [&]() {
return whisper_build_graph_cross(*ctx, *state); return whisper_build_graph_cross(*ctx, *state);
}); });
if (!ok) {
WHISPER_LOG_ERROR("%s: failed to init cross allocator\n", __func__);
whisper_free_state(state);
return nullptr;
}
WHISPER_LOG_INFO("%s: compute buffer (cross) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_cross) / 1e6); WHISPER_LOG_INFO("%s: compute buffer (cross) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_cross) / 1e6);
} }
// decoder allocator // decoder allocator
{ {
whisper_allocr_graph_init(state->alloc_decode, ctx->backend, bool ok = whisper_allocr_graph_init(state->alloc_decode, ctx->backend,
[&]() { [&]() {
const auto & hparams = ctx->model.hparams; const auto & hparams = ctx->model.hparams;
@ -3147,17 +3103,18 @@ struct whisper_state * whisper_init_state(whisper_context * ctx) {
whisper_batch_prep_legacy(state->batch, nullptr, n_tokens, n_past, 0); whisper_batch_prep_legacy(state->batch, nullptr, n_tokens, n_past, 0);
return whisper_build_graph_decoder(*ctx, *state, state->batch); return whisper_build_graph_decoder(*ctx, *state, state->batch, true);
}); });
if (!ok) {
WHISPER_LOG_ERROR("%s: failed to init decoder allocator\n", __func__);
whisper_free_state(state);
return nullptr;
}
WHISPER_LOG_INFO("%s: compute buffer (decode) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_decode) / 1e6); WHISPER_LOG_INFO("%s: compute buffer (decode) = %7.2f MB\n", __func__, whisper_allocr_size(state->alloc_decode) / 1e6);
} }
whisper_allocr_graph_realloc(state->alloc_conv, ctx->backend);
whisper_allocr_graph_realloc(state->alloc_encode, ctx->backend);
whisper_allocr_graph_realloc(state->alloc_cross, ctx->backend);
whisper_allocr_graph_realloc(state->alloc_decode, ctx->backend);
return state; return state;
} }
@ -3380,8 +3337,7 @@ struct whisper_context * whisper_init_no_state(struct whisper_model_loader * loa
return whisper_init_with_params_no_state(loader, whisper_context_default_params()); return whisper_init_with_params_no_state(loader, whisper_context_default_params());
} }
void whisper_free_state(struct whisper_state * state) void whisper_free_state(struct whisper_state * state) {
{
if (state) { if (state) {
kv_cache_free(state->kv_self); kv_cache_free(state->kv_self);
kv_cache_free(state->kv_cross); kv_cache_free(state->kv_cross);
@ -3402,10 +3358,10 @@ void whisper_free_state(struct whisper_state * state)
whisper_batch_free(state->batch); whisper_batch_free(state->batch);
whisper_allocr_free(state->alloc_conv); ggml_gallocr_free(state->alloc_conv.alloc);
whisper_allocr_free(state->alloc_encode); ggml_gallocr_free(state->alloc_encode.alloc);
whisper_allocr_free(state->alloc_cross); ggml_gallocr_free(state->alloc_cross.alloc);
whisper_allocr_free(state->alloc_decode); ggml_gallocr_free(state->alloc_decode.alloc);
ggml_backend_free(state->backend); ggml_backend_free(state->backend);
@ -3415,15 +3371,9 @@ void whisper_free_state(struct whisper_state * state)
void whisper_free(struct whisper_context * ctx) { void whisper_free(struct whisper_context * ctx) {
if (ctx) { if (ctx) {
if (ctx->model.ctx) {
ggml_free(ctx->model.ctx); ggml_free(ctx->model.ctx);
}
for (auto & buffer : ctx->model.buffers) { ggml_backend_buffer_free(ctx->model.buffer);
if (buffer) {
ggml_backend_buffer_free(buffer);
}
}
whisper_free_state(ctx->state); whisper_free_state(ctx->state);