Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml: allow passing alloc dependencies in graph_optimize (#27301)

Browse Source 
* ggml: allow passing alloc dependencies in graph_optimize

* add alloc dep tests

* add TODO about using flat array
This commit is contained in:
Aman Gupta
2026-08-30 11:34:20 +08:00
committed by GitHub
parent c589f0ed10
commit 57291f2644
8 changed files with 151 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-backend.cpp
+58 -8
View File
@@ -20,6 +20,7 @@
#include <stdlib.h>
#include <string.h>
#include <algorithm>
#include <unordered_map>
#include <vector>
#ifdef __APPLE__
@@ -558,10 +559,10 @@ void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event)
backend->iface.event_wait(backend, event);
}
static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params) {
GGML_ASSERT(backend);
if (backend->iface.graph_optimize != NULL) {
backend->iface.graph_optimize(backend, cgraph);
backend->iface.graph_optimize(backend, cgraph, params);
}
}
@@ -1441,11 +1442,40 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
sched->prev_leaf_backend_ids = tmp;
}
// optimize the split graphs and collect the allocation dependencies added by the backends
// this needs to happen before we make graph_copy, so they are in sync
// TODO: this may create many small allocations in the scheduler, restructure to use a flat array
std::unordered_map<ggml_tensor *, std::vector<ggml_tensor *>> alloc_deps;
struct ggml_backend_graph_optimize_params opt_params = {
/* .add_alloc_dep = */ [](void * user_data, ggml_tensor * tensor, ggml_tensor * until) {
auto & deps = *(std::unordered_map<ggml_tensor *, std::vector<ggml_tensor *>> *) user_data;
std::vector<ggml_tensor *> & keep = deps[until];
if (std::find(keep.begin(), keep.end(), tensor) == keep.end()) {
keep.push_back(tensor);
}
},
/* .user_data = */ &alloc_deps,
};
for (int i = 0; i < sched->n_splits; i++) {
struct ggml_backend_sched_split * split = &sched->splits[i];
split->graph = ggml_graph_view(graph, split->i_start, split->i_end);
ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph, &opt_params);
}
// each dep is added to graph_copy as a GGML_OP_NONE node with the kept tensors as srcs
int n_dep_nodes = 0;
for (const auto & it : alloc_deps) {
n_dep_nodes += (it.second.size() + GGML_MAX_SRC - 1) / GGML_MAX_SRC;
}
int total_inputs = sched->n_graph_inputs;
for (int i = 0; i < sched->n_splits; i++) {
total_inputs += sched->splits[i].n_inputs;
}
int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies;
int graph_size = std::max(graph->n_nodes, graph->n_leafs) + total_inputs * 2 * sched->n_copies + n_dep_nodes;
// remember the actual graph_size for performing reallocation checks later [GGML_SCHED_DEBUG_REALLOC]
sched->debug_prev_graph_size = sched->debug_graph_size;
@@ -1463,13 +1493,10 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
struct ggml_cgraph * graph_copy = &sched->graph;
int n_dep_nodes_added = 0;
for (int i = 0; i < sched->n_splits; i++) {
struct ggml_backend_sched_split * split = &sched->splits[i];
split->graph = ggml_graph_view(graph, split->i_start, split->i_end);
// Optimize this split of the graph. This needs to happen before we make graph_copy,
// so they are in sync.
ggml_backend_graph_optimize(sched->backends[split->backend_id], &split->graph);
// 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++) {
@@ -1494,9 +1521,32 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
assert(graph_copy->size > graph_copy->n_nodes);
sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]);
graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j];
if (alloc_deps.empty()) {
continue;
}
// add a dependency node so that the kept tensors are not freed before this node is computed
auto it = alloc_deps.find(graph->nodes[j]);
if (it != alloc_deps.end()) {
const std::vector<ggml_tensor *> & keep = it->second;
for (size_t k = 0; k < keep.size(); k += GGML_MAX_SRC) {
struct ggml_tensor * dep = ggml_view_tensor(sched->ctx, keep[k]);
for (size_t s = 0; s < GGML_MAX_SRC && k + s < keep.size(); s++) {
dep->src[s] = keep[k + s];
}
assert(graph_copy->size > graph_copy->n_nodes);
sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id;
graph_copy->nodes[graph_copy->n_nodes++] = dep;
n_dep_nodes_added++;
}
}
}
}
// a mismatch means a backend added a dep with an `until` tensor that is not a node of the optimized graph
GGML_ASSERT(n_dep_nodes_added == n_dep_nodes);
if (sched->n_copies > 1) {
// add input copies as leafs so that they are allocated first
for (int i = 0; i < sched->n_graph_inputs; i++) {