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

src/core/segment_weight_pipeline.cpp

@@ -0,0 +1,205 @@
+#include "core/segment_weight_pipeline.h"
+
+#include <utility>
+
+#include "core/ggml_graph_cut.h"
+#include "device_residency_manager.h"
+
+namespace sd {
+    static ggml_tensor* canonical_param(
+        ggml_tensor* tensor,
+        const std::unordered_set<const ggml_tensor*>& params) {
+        for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
+            if (params.find(current) != params.end()) {
+                return current;
+            }
+        }
+        return nullptr;
+    }
+
+    SegmentWeightPipeline::SegmentWeightPipeline(
+        const std::shared_ptr<DeviceResidencyManager>& residency_manager,
+        ggml_backend_t compute_backend,
+        uintptr_t owner_id,
+        ggml_cgraph* graph,
+        const ggml_graph_cut::Plan& plan,
+        const std::unordered_set<const ggml_tensor*>& params,
+        bool enabled)
+        : residency_manager_(residency_manager),
+          compute_backend_(compute_backend),
+          owner_id_(owner_id),
+          enabled_(enabled && residency_manager != nullptr) {
+        segment_params_.resize(plan.segments.size());
+        for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
+            std::unordered_set<ggml_tensor*> seen;
+            for (ggml_tensor* tensor :
+                 ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
+                ggml_tensor* param = canonical_param(tensor, params);
+                if (param != nullptr && seen.insert(param).second) {
+                    segment_params_[segment_index].push_back(param);
+                }
+            }
+        }
+    }
+
+    SegmentWeightPipeline::~SegmentWeightPipeline() {
+        segment_end();
+        clear();
+    }
+
+    size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const {
+        for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
+            if (!segment_params_[next].empty()) {
+                return next;
+            }
+        }
+        return SIZE_MAX;
+    }
+
+    std::vector<std::vector<ggml_tensor*>> SegmentWeightPipeline::preferred_eviction_order() const {
+        return {segment_params_.rbegin(), segment_params_.rend()};
+    }
+
+    void SegmentWeightPipeline::disable() {
+        clear();
+        enabled_ = false;
+    }
+
+    void SegmentWeightPipeline::activate(size_t segment_index) {
+        if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) {
+            return;
+        }
+
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr ||
+            !manager->activate_prefetched_params(owner_id_, queued_params_)) {
+            disable();
+            return;
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+    }
+
+    bool SegmentWeightPipeline::ensure_segment_capacity(
+        size_t segment_index,
+        const std::vector<DeviceMemoryRequest>& requests) {
+        if (segment_index >= segment_params_.size()) {
+            return false;
+        }
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            return segment_params_[segment_index].empty();
+        }
+        std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
+        protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end());
+        for (const auto& request : requests) {
+            if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index],
+                                                          preferred_eviction_order(), protected_params)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity) {
+        GGML_ASSERT(pinned_params_.empty());
+        activate(segment_index);
+        if (!ensure_capacity()) {
+            return false;
+        }
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            return segment_params_[segment_index].empty();
+        }
+        if (!manager->prepare_params(segment_params_[segment_index])) {
+            return false;
+        }
+        pinned_params_ = segment_params_[segment_index];
+        return true;
+    }
+
+    void SegmentWeightPipeline::segment_end() {
+        if (auto manager = residency_manager_.lock()) {
+            manager->release_compute_backend_params(pinned_params_);
+        }
+        pinned_params_.clear();
+    }
+
+    void SegmentWeightPipeline::enqueue_next(
+        size_t segment_index,
+        const DeviceMemoryRequest& request) {
+        if (!enabled_ || queued_segment_ != SIZE_MAX) {
+            return;
+        }
+
+        const size_t next_segment = next_parameter_segment(segment_index);
+        if (next_segment == SIZE_MAX) {
+            return;
+        }
+
+        std::unordered_set<ggml_tensor*> active_params(
+            segment_params_[segment_index].begin(),
+            segment_params_[segment_index].end());
+        std::vector<ggml_tensor*> params;
+        params.reserve(segment_params_[next_segment].size());
+        for (ggml_tensor* param : segment_params_[next_segment]) {
+            if (active_params.find(param) == active_params.end()) {
+                params.push_back(param);
+            }
+        }
+        if (params.empty()) {
+            return;
+        }
+
+        auto manager = residency_manager_.lock();
+        if (manager == nullptr) {
+            disable();
+            return;
+        }
+        const WeightResidencyInfo residency =
+            manager->inspect_compute_backend_params(params);
+        if (residency.missing_bytes == 0) {
+            return;
+        }
+        if (!residency.async_prefetch_supported) {
+            disable();
+            return;
+        }
+
+        DeviceMemoryRequest backend_request        = request;
+        backend_request.compute_backend            = compute_backend_;
+        backend_request.owner_id                   = owner_id_;
+        std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
+        protected_params.insert(protected_params.end(), params.begin(), params.end());
+        if (!manager->ensure_compute_backend_capacity(backend_request,
+                                                      params,
+                                                      preferred_eviction_order(),
+                                                      protected_params)) {
+            disable();
+            return;
+        }
+        switch (manager->prefetch_params(owner_id_, params)) {
+            case WeightPrefetchResult::Scheduled:
+                queued_params_  = std::move(params);
+                queued_segment_ = next_segment;
+                return;
+            case WeightPrefetchResult::AlreadyResident:
+                return;
+            case WeightPrefetchResult::Unsupported:
+                disable();
+                return;
+            case WeightPrefetchResult::Failed:
+                disable();
+                return;
+        }
+        disable();
+    }
+
+    void SegmentWeightPipeline::clear() {
+        if (auto manager = residency_manager_.lock()) {
+            manager->clear_prefetched_params(owner_id_);
+        }
+        queued_params_.clear();
+        queued_segment_ = SIZE_MAX;
+    }
+}