diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 40cea024c3..56f0090cce 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -103,6 +103,16 @@ extern "C" { // Backend (stream) // + // passed to graph_optimize so the backend can add allocation dependencies: + // if the backend executes parts of the graph out of order (e.g. on concurrent streams), + // it must keep the affected tensors allocated until a node where execution is known to have joined + struct ggml_backend_graph_optimize_params { + // keep `tensor` allocated at least until `until` (a node of the same graph) has been computed + // can be called multiple times for the same tensor: the longest lifetime applies + void (*add_alloc_dep)(void * user_data, struct ggml_tensor * tensor, struct ggml_tensor * until); + void * user_data; + }; + struct ggml_backend_i { const char * (*get_name)(ggml_backend_t backend); @@ -137,7 +147,7 @@ extern "C" { void (*event_wait) (ggml_backend_t backend, ggml_backend_event_t event); // (optional) sort/optimize the nodes in the graph - void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph); + void (*graph_optimize) (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params); }; struct ggml_backend { diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index e519bdf50a..78eb10dfe9 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #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> 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> *) user_data; + std::vector & 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 & 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++) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 2456f7dcc6..bd9754c2ff 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4328,7 +4328,9 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev } } -static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; #ifdef USE_CUDA_GRAPH diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 53e8607559..e7dcdc3d55 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -4984,7 +4984,9 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vectorn_nodes; constexpr int MAX_FUSE = 16; diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050..4d58dc821c 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -558,7 +558,9 @@ static void ggml_backend_metal_event_wait(ggml_backend_t backend, ggml_backend_e ggml_metal_event_wait(ctx, ev); } -static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_metal_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); + ggml_metal_t ctx = (ggml_metal_t)backend->context; ggml_metal_graph_optimize(ctx, cgraph); diff --git a/ggml/src/ggml-virtgpu/ggml-backend.cpp b/ggml/src/ggml-virtgpu/ggml-backend.cpp index 12756c9282..996c57e358 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend.cpp @@ -17,7 +17,8 @@ static ggml_status ggml_backend_remoting_graph_compute(ggml_backend_t backend, g return apir_backend_graph_compute(gpu, cgraph); } -static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { +static void ggml_backend_remoting_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { + UNUSED(params); virtgpu * gpu = DEV_TO_GPU(backend->device); #if true UNUSED(gpu); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 39b4cd3598..8fbb1359f4 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17795,8 +17795,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } // Sort the graph for improved parallelism. -static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph) +static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, struct ggml_backend_graph_optimize_params * params) { + GGML_UNUSED(params); VK_LOG_DEBUG("ggml_vk_graph_optimize(" << graph->n_nodes << " nodes)"); ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp index 6d5428493e..8f1a98aa03 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -19,6 +19,8 @@ struct dummy_backend_context { size_t alignment = 8; ggml_backend_buffer_i buffer_interface; + ggml_backend_device device; + ggml_backend backend; std::vector buffers; size_t allocated_total() const { @@ -83,7 +85,27 @@ static void dummy_backend_buffer_get_tensor(ggml_backend_buffer_t, const ggml_te static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} -// dummy_backend (not really a full backend, just provides what gallocr needs) +// ggml_backend_device interface + +static enum ggml_backend_dev_type dummy_backend_device_get_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +static bool dummy_backend_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +static bool dummy_backend_device_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) { + return device->context == buft->context; +} + +// ggml_backend interface + +static const char * dummy_backend_get_name(ggml_backend_t) { + return "dummy_backend"; +} + +// dummy_backend struct dummy_backend { std::unique_ptr context; @@ -104,6 +126,16 @@ static dummy_backend dummy_backend_init(size_t max_buffer_size, size_t alignment b.context->buffer_interface.get_tensor = dummy_backend_buffer_get_tensor; b.context->buffer_interface.clear = dummy_backend_buffer_clear; + b.context->device.context = b.context.get(); + b.context->device.iface.get_type = dummy_backend_device_get_type; + b.context->device.iface.supports_op = dummy_backend_device_supports_op; + b.context->device.iface.supports_buft = dummy_backend_device_supports_buft; + + b.context->backend.context = b.context.get(); + b.context->backend.device = &b.context->device; + b.context->backend.iface.get_name = dummy_backend_get_name; + + b.buffer_type.device = &b.context->device; b.buffer_type.context = b.context.get(); b.buffer_type.iface.get_name = dummy_backend_buffer_type_get_name; b.buffer_type.iface.alloc_buffer = dummy_backend_buffer_type_alloc_buffer; @@ -583,6 +615,41 @@ static void test_reallocation() { } } +static void test_backend_graph_optimize(ggml_backend_t, ggml_cgraph * graph, ggml_backend_graph_optimize_params * params) { + GGML_ASSERT(graph->n_nodes == 3); + params->add_alloc_dep(params->user_data, graph->nodes[0], graph->nodes[2]); +} + +static bool graph_reuses_allocation(bool add_alloc_dep) { + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[4]; + x[0] = make_input_with_size(ctx, 16); + x[1] = ggml_scale(ctx, x[0], 2.0f); + x[2] = ggml_scale(ctx, x[1], 2.0f); + x[3] = ggml_scale(ctx, x[2], 2.0f); + + ggml_set_output(x[3]); + ggml_build_forward_expand(graph, x[3]); + + dummy_backend backend = dummy_backend_init(SIZE_MAX); + if (add_alloc_dep) { + backend.context->backend.iface.graph_optimize = test_backend_graph_optimize; + } + + ggml_backend_t backend_ptr = &backend.context->backend; + ggml_backend_buffer_type_t buft = &backend.buffer_type; + ggml_backend_sched_ptr sched(ggml_backend_sched_new(&backend_ptr, &buft, 1, 8, false, true)); + GGML_ASSERT(ggml_backend_sched_alloc_graph(sched.get(), graph)); + + return x[1]->data == x[2]->data; +} + +static void test_graph_optimize_alloc_dep() { + GGML_ASSERT(graph_reuses_allocation(false)); + GGML_ASSERT(!graph_reuses_allocation(true)); +} + static void run(const char * name, void (*f)()) { printf("%s ", name); fflush(stdout); @@ -604,5 +671,6 @@ int main() { run("test_multiple_buffer_types", test_multiple_buffer_types); run("test_buffer_size_zero", test_buffer_size_zero); run("test_reallocation", test_reallocation); + run("test_graph_optimize_alloc_dep", test_graph_optimize_alloc_dep); return 0; }