refactor: unify runner lifecycles and weight residency (#1940)

src/core/compute_workspace.cpp

@@ -0,0 +1,257 @@
+#include "core/compute_workspace.h"
+
+#include <algorithm>
+#include <cstring>
+#include <map>
+#include <unordered_map>
+#include <unordered_set>
+
+#include "core/ggml_extend_backend.h"
+#include "core/ggml_graph_cut.h"
+#include "ggml-cpu.h"
+#include "ggml/src/ggml-impl.h"
+
+namespace sd {
+    ComputeWorkspace::~ComputeWorkspace() {
+        segment_end();
+        release();
+        ggml_backend_free(cpu_backend_);
+    }
+
+    void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) {
+        if (extra_backends_ != backends) {
+            GGML_ASSERT(!active_);
+            release();
+            extra_backends_ = backends;
+        }
+    }
+
+    bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
+        if (!extra_backends_.empty()) {
+            return true;
+        }
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
+        std::vector<ggml_backend_t> backends{backend_};
+        backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
+        if (!sd_backend_is_cpu(backend_)) {
+            if (cpu_backend_ == nullptr) {
+                cpu_backend_ = sd_backend_cpu_init();
+            }
+            if (cpu_backend_ == nullptr) {
+                return nullptr;
+            }
+            backends.push_back(cpu_backend_);
+        }
+        std::vector<ggml_backend_buffer_type_t> bufts;
+        for (auto backend : backends) {
+            auto buft = backend == cpu_backend_
+                            ? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_))
+                            : nullptr;
+            bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend));
+        }
+        return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast<int>(backends.size()),
+                                      graph_size, false, false);
+    }
+
+    bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
+        return measurement.scheduler == needs_scheduler(graph);
+    }
+
+    bool ComputeWorkspace::prepare(const Measurement& measurement) {
+        GGML_ASSERT(!active_);
+        if (measurement.buffers.empty()) {
+            return false;
+        }
+        const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
+                                       [&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
+        if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
+            release();
+        }
+        return true;
+    }
+
+    bool ComputeWorkspace::release_excess(const Measurement& measurement) {
+        return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
+                           [&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
+               release();
+    }
+
+    bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
+        GGML_ASSERT(!active_);
+        const bool use_scheduler = needs_scheduler(graph);
+        if (use_scheduler) {
+            if (allocator_ != nullptr) {
+                release();
+            }
+            const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+            if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
+                release();
+                scheduler_          = make_scheduler(capacity);
+                scheduler_capacity_ = capacity;
+            }
+            if (scheduler_ == nullptr) {
+                return false;
+            }
+            ggml_backend_sched_reset(scheduler_);
+            assign_nodes(scheduler_, graph);
+            // Scheduler allocation rewrites sources. Split the execution graph only once.
+            if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
+                release();
+                return false;
+            }
+        } else {
+            if (scheduler_ != nullptr) {
+                release();
+            }
+            if (allocator_ == nullptr) {
+                allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
+            }
+            auto signature = ggml_graph_cut::graph_layout(graph, true);
+            if (signature != reservation_) {
+                if (!ggml_gallocr_reserve(allocator_, graph)) {
+                    release();
+                    return false;
+                }
+                reservation_ = std::move(signature);
+                ++reservations_;
+            }
+            if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
+                release();
+                return false;
+            }
+        }
+        active_ = true;
+        return true;
+    }
+
+    ComputeWorkspace::Measurement ComputeWorkspace::measure(
+        ggml_cgraph* graph,
+        size_t direct_bytes,
+        const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
+        const AssignNodes& assign_nodes) {
+        if (!needs_scheduler(graph)) {
+            return {{{backend_, direct_bytes}}, false};
+        }
+        std::vector<const ggml_tensor*> tensors;
+        std::unordered_set<const ggml_tensor*> seen;
+        auto visit = [&](const ggml_tensor* tensor) {
+            if (tensor != nullptr && seen.insert(tensor).second) {
+                tensors.push_back(tensor);
+            }
+        };
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            visit(graph->nodes[i]);
+        }
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            visit(graph->leafs[i]);
+        }
+        for (size_t i = 0; i < tensors.size(); ++i) {
+            visit(tensors[i]->view_src);
+            for (auto source : tensors[i]->src) {
+                visit(source);
+            }
+        }
+        const size_t graph_size = tensors.size() + 8;
+        auto context            = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
+        if (context == nullptr) {
+            return {};
+        }
+        std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
+        std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
+        for (auto tensor : tensors) {
+            auto copy      = ggml_dup_tensor(context, tensor);
+            *copy          = *tensor;
+            copies[tensor] = copy;
+        }
+        for (const auto& entry : copies) {
+            auto source    = entry.first;
+            auto copy      = entry.second;
+            copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src);
+            for (int i = 0; i < GGML_MAX_SRC; ++i) {
+                copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]);
+            }
+            auto external = external_backend(source);
+            if (external != nullptr && source->view_src == nullptr) {
+                auto& buffer = external_buffers[external];
+                if (buffer == nullptr) {
+                    buffer = ggml_backend_alloc_buffer(external, 0);
+                    GGML_ASSERT(buffer != nullptr);
+                    ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+                }
+                copy->buffer = buffer;
+                copy->data   = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+                copy->extra  = nullptr;
+            }
+        }
+        auto copy_graph     = ggml_new_graph_custom(context, graph_size, false);
+        copy_graph->n_nodes = graph->n_nodes;
+        copy_graph->n_leafs = graph->n_leafs;
+        for (int i = 0; i < graph->n_nodes; ++i) {
+            copy_graph->nodes[i] = copies.at(graph->nodes[i]);
+        }
+        for (int i = 0; i < graph->n_leafs; ++i) {
+            copy_graph->leafs[i] = copies.at(graph->leafs[i]);
+        }
+        Measurement result;
+        result.scheduler = true;
+        auto scheduler   = make_scheduler(graph_size);
+        if (scheduler != nullptr) {
+            assign_nodes(scheduler, copy_graph);
+            std::vector<size_t> sizes(extra_backends_.size() + 2);
+            ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data());
+            result.buffers.push_back({backend_, sizes[0]});
+            for (size_t i = 0; i < extra_backends_.size(); ++i) {
+                result.buffers.push_back({extra_backends_[i], sizes[i + 1]});
+            }
+            ggml_backend_sched_free(scheduler);
+        }
+        for (const auto& entry : external_buffers) {
+            ggml_backend_buffer_free(entry.second);
+        }
+        ggml_free(context);
+        return result;
+    }
+
+    void ComputeWorkspace::synchronize() const {
+        if (scheduler_ != nullptr) {
+            ggml_backend_sched_synchronize(scheduler_);
+        } else {
+            ggml_backend_synchronize(backend_);
+        }
+    }
+
+    void ComputeWorkspace::segment_end() {
+        if (active_) {
+            synchronize();
+            active_ = false;
+        }
+    }
+
+    bool ComputeWorkspace::release() {
+        if (active_) {
+            return false;
+        }
+        ggml_gallocr_free(allocator_);
+        allocator_ = nullptr;
+        ggml_backend_sched_free(scheduler_);
+        scheduler_          = nullptr;
+        scheduler_capacity_ = 0;
+        reservation_.clear();
+        return true;
+    }
+
+    size_t ComputeWorkspace::bytes(ggml_backend_t backend) const {
+        if (scheduler_ != nullptr) {
+            return ggml_backend_sched_get_buffer_size(scheduler_, backend);
+        }
+        return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
+    }
+}