From ddc367ab5d502113e0df3b5d9415112ce89fac24 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 23:45:01 +0800 Subject: [PATCH] refactor: unify runner execution and weight residency (#2321) --- docs/backend.md | 24 +- src/core/execution_engine.cpp | 992 +++++++++++++++++++++++++ src/core/execution_engine.h | 107 +++ src/core/ggml_runner.cpp | 818 +------------------- src/core/ggml_runner.h | 76 +- src/device_residency_manager.h | 57 +- src/model/diffusion/qwen_image_2_1.hpp | 2 +- src/model_manager.cpp | 385 +++++----- src/model_manager.h | 30 +- src/model_manager_files.cpp | 42 +- src/pipeline/image.cpp | 22 +- src/pipeline/image_upscale.cpp | 138 ++++ src/pipeline/image_upscale.h | 57 ++ src/upscaler.cpp | 173 +---- src/upscaler.h | 9 +- 15 files changed, 1679 insertions(+), 1253 deletions(-) create mode 100644 src/core/execution_engine.cpp create mode 100644 src/core/execution_engine.h create mode 100644 src/pipeline/image_upscale.cpp create mode 100644 src/pipeline/image_upscale.h diff --git a/docs/backend.md b/docs/backend.md index da90e489..cc6545da 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -160,10 +160,12 @@ the runner's graph-cut capacity checks. Runtime capacity checks also leave 512 MiB of currently free device memory for backend scratch buffers and pipelines, including with explicit backend assignments. They cap free-memory reports by the device's total memory minus tracked -resident allocations. Vulkan reports exceeding total memory are rejected because -its heap-budget subtraction can underflow. Other backends use the cap instead of -treating such reports as zero free memory. Failed checks log the reported free and -total memory alongside tracked weight and runtime allocations. +resident allocations. Vulkan reports exceeding total memory are rejected on +discrete GPUs because its heap-budget subtraction can underflow. On integrated +GPUs, a wrapped report triggers reclamation before the tracked estimate can be +accepted, subject to the byte checks and configured budgets. Other backends use +the cap instead of treating such reports as zero free memory. Failed checks log +the reported free and total memory alongside tracked weight and runtime allocations. With `--mmap`, device-backed mappings count toward these budgets at their full mapped-file size, once per device buffer even when multiple parameter blocks share it. Mappings retained in the loader cache continue to count. @@ -181,8 +183,18 @@ GPU cache space follows the same component priority. Before a lower-priority component can become permanently resident, the planner leaves room for the full weights and estimated compute space of higher-priority offloaded components. If offloaded diffusion already needs the entire main GPU budget, TE and VAE also -use offloaded parameters. Their GPU copies can then be released after their -phases, leaving more room to reuse diffusion weights across sampling steps. +use offloaded parameters. GPU copies remain cached across sampling steps while +a component's execution is active and are released under memory pressure. +Ending a component's execution releases its staged GPU copies, disk-backed +weight buffers, temporary compute buffers and execution caches. Parameters +assigned permanently to a backend remain resident. Hires model upscaling shares +the runtime budget and can +reclaim idle weight copies from other components. +Reclamation releases idle compute buffers first, then cached weight allocations +until the current work fits. Weights used by active computation or pending +transfers stay protected. +For CPU computation with disk-backed parameters, cached weights are released +when execution ends so they do not accumulate in RAM across components. CPU parameter residency allows GPU weight caching; it does not force every weight to be copied again at every step. diff --git a/src/core/execution_engine.cpp b/src/core/execution_engine.cpp new file mode 100644 index 00000000..5ad241a1 --- /dev/null +++ b/src/core/execution_engine.cpp @@ -0,0 +1,992 @@ +#include "core/execution_engine.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" +#include "core/layer_split_partition.h" +#include "core/segment_graph_bindings.h" +#include "core/segment_weight_pipeline.h" +#include "core/tensor_ggml.hpp" +#include "core/util.h" + +std::shared_ptr DeviceResidencyManager::execution_engine() { + if (execution_engine_ == nullptr) { + execution_engine_ = std::make_shared(this); + } + return execution_engine_; +} + +namespace sd { + class GraphExecution { + ExecutionEngine& engine_; + ExecutionSession& session_; + const ComputeTask& task_; + std::shared_ptr manager_; + std::unordered_set params_tensor_set_; + ggml_status last_compute_status_ = GGML_STATUS_FAILED; + Tensor output_; + + using GraphCutPlan = ggml_graph_cut::Plan; + ggml_tensor* canonical_param_tensor(ggml_tensor* tensor); + std::vector collect_used_param_tensors(ggml_cgraph* graph); + ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const; + // Weightless ops follow the most recent weight device. Views must stay + // unpinned so the scheduler preserves cross-device copies for consumers. + void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* graph, ggml_cgraph* original_graph = nullptr); + Tensor read_graph_tensor(ggml_tensor* tensor, const char* label); + void copy_data_to_backend_tensor(ggml_cgraph* graph); + const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* graph); + bool assign_graph_cut_layer_split_backends(ggml_cgraph* graph); + ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set* external_outputs = nullptr); + std::vector memory_requests(const std::vector& sizes, size_t pending_cache_bytes) const; + bool fits(const std::vector& requests, const std::vector& params) const; + bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute); + bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); + bool is_multi_device() const { return !session_.extra_backends_.empty(); } + size_t retained_runtime_buffer_bytes(ggml_backend_t backend) const { return session_.resident_bytes(backend); } + void sync_runtime_residency() { engine_.sync_residency(session_); } + + public: + GraphExecution(ExecutionEngine& engine, ExecutionSession& session, const ComputeTask& task) + : engine_(engine), session_(session), task_(task), manager_(session.residency_manager_.lock()) { + params_tensor_set_.insert(task.weights.begin(), task.weights.end()); + if (manager_ != nullptr) { + for (int i = 0; i < ggml_graph_cut::leaf_count(task.graph); ++i) { + auto* parameter = manager_->resolve_param_tensor(ggml_graph_cut::leaf_tensor(task.graph, i)); + if (parameter != nullptr) { + params_tensor_set_.insert(parameter); + } + } + for (auto* tensor : params_tensor_set_) { + auto* parameter = manager_->resolve_param_tensor(const_cast(tensor)); + if (parameter != nullptr) { + session_.params_.insert(parameter); + } + } + } + } + + ComputeResult run() { + if (!execute_graph(task_.graph, task_.n_threads, task_.no_return, task_.read_outputs)) { + return {last_compute_status_, {}}; + } + return {GGML_STATUS_SUCCESS, std::move(output_)}; + } + }; + + static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { + // Vulkan reports a preferred allocation chunk size, not a hard tensor limit. + if (sd_backend_is(backend, "Vulkan")) { + return true; + } + auto buft = ggml_backend_get_default_buffer_type(backend); + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t max_size = ggml_backend_buft_get_max_size(buft); + max_size -= max_size % alignment; + if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { + return false; + } + // Cross-backend copies must also fit the destination's allocation limit. + for (auto src : node->src) { + if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { + return false; + } + } + return true; + } + + static size_t add_bytes(size_t a, size_t b) { + return b > SIZE_MAX - a ? SIZE_MAX : a + b; + } + + ggml_tensor* GraphExecution::canonical_param_tensor(ggml_tensor* tensor) { + for (auto* current = tensor; current != nullptr; current = current->view_src) { + if (params_tensor_set_.count(current) != 0) + return current; + } + return nullptr; + } + + std::vector GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) { + std::vector used_params; + if (gf == nullptr || params_tensor_set_.empty()) { + return used_params; + } + + std::unordered_set seen_params; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + seen_params.reserve(static_cast(n_leafs)); + for (int i = 0; i < n_leafs; ++i) { + ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); + ggml_tensor* param_leaf = canonical_param_tensor(leaf); + if (param_leaf != nullptr && + seen_params.insert(param_leaf).second) { + used_params.push_back(param_leaf); + } + } + return used_params; + } + + ggml_backend_t GraphExecution::backend_for_weight(const ggml_tensor* tensor) const { + if (tensor == nullptr || tensor->buffer == nullptr) { + return nullptr; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + ggml_backend_buffer_is_host(tensor->buffer)) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + if (dev == nullptr) { + return nullptr; + } + if (ggml_backend_get_device(session_.backend_) == dev) { + return session_.backend_; + } + for (ggml_backend_t backend : session_.extra_backends_) { + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + return nullptr; + } + + void GraphExecution::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { + if (sched == nullptr || gf == nullptr) { + return; + } + ggml_backend_t current = session_.backend_; + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); + if (node_assignment != session_.node_assignments_.end()) { + current = node_assignment->second; + } + for (int s = 0; s < GGML_MAX_SRC; s++) { + ggml_backend_t weight_backend = backend_for_weight(node->src[s]); + if (weight_backend != nullptr) { + if (node_assignment == session_.node_assignments_.end()) { + current = weight_backend; + } + } + } + if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (!backend_can_allocate_op(current, node)) { + auto cpu = session_.workspace_.cpu_backend(); + if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, cpu); + } + } else if (ggml_backend_supports_op(current, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, current); + } + } + } + + Tensor GraphExecution::read_graph_tensor(ggml_tensor* tensor, const char* label) { + if (tensor == nullptr) { + LOG_ERROR("%s %s tensor is null", task_.description.c_str(), label); + return {}; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_ERROR("%s %s tensor type mismatch: got %s", + task_.description.c_str(), + label, + ggml_type_name(tensor->type)); + return {}; + } + ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); + if (buf == nullptr) { + LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", + task_.description.c_str(), + label, + tensor->name[0] != '\0' ? tensor->name : "", + ggml_op_name(tensor->op), + tensor->buffer, + tensor->view_src, + tensor->view_src ? tensor->view_src->buffer : nullptr, + tensor->data); + return {}; + } + + return sd::make_sd_tensor_from_ggml(tensor); + } + + void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) { + GGML_ASSERT(gf != nullptr); + std::unordered_set graph_tensor_set; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_nodes = ggml_graph_n_nodes(gf); + graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); + for (int i = 0; i < n_leafs; ++i) { + graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_nodes; ++i) { + graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + + for (auto& kv : task_.inputs) { + auto tensor = kv.first; + auto data = kv.second; + if (tensor == nullptr || data == nullptr) { + continue; + } + const char* name = ggml_get_name(tensor); + if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { + continue; + } + if (tensor->buffer == nullptr) { + LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", + task_.description.c_str(), + name != nullptr ? name : "", + (long long)tensor->ne[0], + (long long)tensor->ne[1], + (long long)tensor->ne[2], + (long long)tensor->ne[3], + ggml_type_name(tensor->type)); + continue; + } + + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr) { + LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", + task_.description.c_str(), + tensor && tensor->name[0] != '\0' ? tensor->name : "", + tensor ? ggml_op_name(tensor->op) : "", + data, + tensor ? tensor->view_src : nullptr, + (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); + continue; + } + + ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); + } + } + + const GraphExecution::GraphCutPlan& GraphExecution::resolve_graph_cut_plan(ggml_cgraph* gf) { + GGML_ASSERT(gf != nullptr); + return sd::ggml_graph_cut::resolve_plan(session_.backend_, + gf, + &session_.plan_cache_, + params_tensor_set_, + task_.description.c_str()); + } + + bool GraphExecution::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { + session_.node_assignments_.clear(); + if (!session_.layer_split_enabled_) { + return true; + } + if (!is_multi_device()) { + LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", task_.description.c_str()); + return false; + } + + const auto& plan = resolve_graph_cut_plan(gf); + if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { + auto manager = manager_; + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str()); + return false; + } + std::vector graph_params = collect_used_param_tensors(gf); + if (!graph_params.empty() && + !manager->assign_compute_backend(graph_params, session_.backend_)) { + LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str()); + return false; + } + for (ggml_tensor* param : graph_params) { + if (param != nullptr) { + session_.weight_assignments_[param] = session_.backend_; + } + } + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (node != nullptr) { + session_.node_assignments_[node] = session_.backend_; + } + } + if (!session_.primary_notice_logged_) { + LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str(), + graph_params.size()); + session_.primary_notice_logged_ = true; + } else { + LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str(), + graph_params.size()); + } + return true; + } + + std::vector split_backends; + split_backends.reserve(session_.extra_backends_.size() + 1); + split_backends.push_back(session_.backend_); + for (ggml_backend_t backend : session_.extra_backends_) { + if (backend != nullptr) { + split_backends.push_back(backend); + } + } + + auto manager = manager_; + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str()); + return false; + } + + sd::GraphCutLayerSplitAssignment assignment; + auto canonicalize_param = [this](ggml_tensor* tensor) { + return canonical_param_tensor(tensor); + }; + if (!sd::partition_graph_cut_layer_split(task_.description.c_str(), + gf, + plan, + split_backends, + session_.backend_limits_, + session_.max_backend_bytes_, + session_.weight_assignments_, + canonicalize_param, + &assignment)) { + return false; + } + + for (size_t i = 0; i < split_backends.size(); i++) { + if (assignment.tensors_by_backend[i].empty()) { + continue; + } + if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { + LOG_ERROR("%s graph-cut layer split failed to assign params to %s", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); + return false; + } + } + + session_.node_assignments_ = std::move(assignment.node_assignments); + sd::log_graph_cut_layer_split_assignment(task_.description.c_str(), split_backends, assignment); + + return true; + } + + ComputeWorkspace::Measurement GraphExecution::measure(ggml_cgraph* graph, const std::unordered_set* external_outputs) { + auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t { + if (external_outputs != nullptr && external_outputs->count(tensor)) { + return session_.backend_; + } + if (!params_tensor_set_.count(tensor)) { + return nullptr; + } + auto placement = session_.weight_assignments_.find(tensor); + return placement == session_.weight_assignments_.end() ? session_.backend_ : placement->second; + }; + auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { + pin_multi_device_nodes(scheduler, copy, graph); + }; + return session_.workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device()); + } + + std::vector GraphExecution::memory_requests( + const std::vector& sizes, + size_t pending_cache_bytes) const { + std::vector requests; + for (const auto& size : sizes) { + const size_t retained = retained_runtime_buffer_bytes(size.backend); + const size_t reusable = session_.workspace_.bytes(size.backend); + const size_t cache_bytes = size.backend == session_.backend_ ? pending_cache_bytes : 0; + const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); + size_t limit = session_.max_backend_bytes_; + if (is_multi_device()) { + size_t index = 0; + if (size.backend != session_.backend_) { + auto position = std::find(session_.extra_backends_.begin(), session_.extra_backends_.end(), size.backend); + index = static_cast(position - session_.extra_backends_.begin()) + 1; + } + if (index < session_.backend_limits_.size()) { + limit = session_.backend_limits_[index]; + } + } + requests.push_back({size.backend, session_.owner_id(), pending, + retained, limit}); + } + return requests; + } + + bool GraphExecution::fits(const std::vector& requests, + const std::vector& params) const { + auto manager = manager_; + if (manager == nullptr) { + return params.empty(); + } + for (const auto& request : requests) { + if (!manager->fits_compute_backend_capacity(request, params)) { + return false; + } + } + return true; + } + + bool GraphExecution::execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute) { + if (sd_backend_is_cpu(session_.backend_)) { + sd_backend_cpu_set_n_threads(session_.backend_, n_threads); + } + if (session_.workspace_.cpu_backend() != nullptr) { + sd_backend_cpu_set_n_threads(session_.workspace_.cpu_backend(), n_threads); + } + auto scheduler = session_.workspace_.scheduler(); + GGML_ASSERT(scheduler != nullptr); + ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data()); +#ifndef SD_USE_UPSTREAM_GGML + struct BeforeComputeCallback { + ggml_backend_t backend; + const std::function& callback; + bool invoked = false; + } before_compute{session_.backend_, on_before_compute}; + ggml_backend_sched_set_compute_callback( + scheduler, [](ggml_backend_t backend, void* data) { + auto& hook = *static_cast(data); + if (backend == hook.backend && !hook.invoked) { + hook.invoked = true; + hook.callback(); + } + }, + &before_compute); + const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); + ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr); +#else + on_before_compute(); + const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); +#endif + session_.workspace_.synchronize(); + if (status != GGML_STATUS_SUCCESS) { + last_compute_status_ = status; + LOG_ERROR("%s compute failed: %s", task_.description.c_str(), ggml_status_to_string(status)); + return false; + } + const std::string description = task_.description; + if (!task_.debug_tensors.empty()) { + std::unordered_set graph_tensors; + const int leaf_count = ggml_graph_cut::leaf_count(graph); + const int node_count = ggml_graph_n_nodes(graph); + graph_tensors.reserve(static_cast(leaf_count + node_count)); + for (int index = 0; index < leaf_count; ++index) { + graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); + } + for (int index = 0; index < node_count; ++index) { + graph_tensors.insert(ggml_graph_node(graph, index)); + } + + for (const auto& entry : task_.debug_tensors) { + ggml_tensor* tensor = entry.first; + if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { + continue; + } + ggml_backend_buffer_t buffer = + tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + description.c_str(), + entry.second.c_str()); + continue; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", + description.c_str(), + entry.second.c_str(), + ggml_type_name(tensor->type)); + continue; + } + auto debug_tensor = make_sd_tensor_from_ggml(tensor); + print_sd_tensor(debug_tensor, false, entry.second.c_str()); + } + } + + return true; + } + + bool GraphExecution::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { + if (!assign_graph_cut_layer_split_backends(graph)) { + return false; + } + const auto params = collect_used_param_tensors(graph); + const auto& cached_plan = resolve_graph_cut_plan(graph); + auto manager = manager_; + const bool can_segment = !is_multi_device() && !sd_backend_is_cpu(session_.backend_) && + manager != nullptr && manager->segmented_compute_enabled() && + cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1; + // An empty workspace request is a lower bound on the full graph's memory demand. + bool segmented = can_segment && !fits(memory_requests({{session_.backend_, 0}}, 0), params); + ComputeWorkspace::Measurement full_measurement; + if (!segmented) { + full_measurement = measure(graph); + if (full_measurement.buffers.empty()) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return false; + } + if (can_segment) { + // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits. + constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL; + auto requests = memory_requests(full_measurement.buffers, session_.cache_.pending_bytes(graph)); + for (auto& request : requests) { + request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom); + } + segmented = !fits(requests, params); + } + } + ggml_graph_cut::Plan monolithic_plan; + if (!segmented) { + monolithic_plan.segments.emplace_back(); + auto& segment = monolithic_plan.segments.back(); + segment.group_name = "graph"; + segment.compute_buffer_size = cached_plan.compute_buffer_size; + segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph)); + segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph)); + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + segment.internal_node_indices.push_back(i); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + auto tensor = ggml_graph_cut::leaf_tensor(graph, i); + ggml_graph_cut::Segment::InputRef input; + input.leaf_index = i; + input.type = canonical_param_tensor(tensor) != nullptr + ? ggml_graph_cut::Segment::INPUT_PARAM + : ggml_graph_cut::Segment::INPUT_EXTERNAL; + segment.input_refs.push_back(input); + } + } + const auto& plan = segmented ? cached_plan : monolithic_plan; + const bool segments_changed = plan.segments.size() != session_.logged_segment_count_; + if (segments_changed && (segmented || session_.logged_segment_count_ > 1)) { + LOG_VERBOSE("%s using %zu segment%s", task_.description.c_str(), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + SegmentGraphBindings bindings(session_.cut_cache_, plan, graph); + SegmentWeightPipeline weights(manager, session_.backend_, session_.owner_id(), + graph, plan, params_tensor_set_, + segmented && manager != nullptr && manager->prefetch_enabled()); + + std::map peak_compute_bytes; + auto track_compute_buffer = [&](ggml_backend_t backend) { + if (backend != nullptr) { + auto& peak = peak_compute_bytes[backend]; + peak = std::max(peak, session_.workspace_.bytes(backend)); + } + }; + for (size_t index = 0; index < plan.segments.size(); ++index) { + const auto& segment = plan.segments[index]; + const bool last = index + 1 == plan.segments.size(); + auto fail_segment = [&](const char* phase) { + LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", task_.description.c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + return false; + }; + session_.cut_cache_.prune(segment.live_cut_names); + bindings.reset(segment); + if (!bindings.bind_cached_inputs(segment, task_.description.c_str())) { + return fail_segment("input binding"); + } + ggml_context* segment_context = nullptr; + auto segment_graph = segmented + ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) + : graph; + struct SegmentCleanup { + GraphExecution& runner; + SegmentWeightPipeline& weights; + SegmentGraphBindings& bindings; + ggml_context* context; + ~SegmentCleanup() { + if (!runner.session_.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } + bindings.restore(); + weights.segment_end(); + ggml_free(context); + runner.sync_runtime_residency(); + } + } segment_cleanup{*this, weights, bindings, segment_context}; + + std::unordered_set cache_outputs; + if (segmented) { + for (int output_index : segment.output_node_indices) { + auto tensor = ggml_graph_node(graph, output_index); + if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) && + GraphCutTensorCache::can_bind_output(session_.backend_, tensor)) { + cache_outputs.insert(tensor); + } + } + } + auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement; + if (!session_.workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + const size_t cut_bytes = last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment); + size_t new_cache_bytes = add_bytes(cut_bytes, session_.cache_.pending_bytes(segment_graph)); + auto ensure_capacity = [&]() { + sync_runtime_residency(); + auto requests = memory_requests(measurement.buffers, new_cache_bytes); + if (fits(requests, weights.params(index))) { + return true; + } + if (session_.workspace_.release_excess(measurement)) { + sync_runtime_residency(); + requests = memory_requests(measurement.buffers, new_cache_bytes); + } + const bool ready = weights.ensure_segment_capacity(index, requests); + if (!ready && manager != nullptr) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + } + return ready; + }; + if (!weights.segment_start(index, ensure_capacity)) { + return fail_segment("weight preparation"); + } + // Preparing weights can execute LoRA graphs and reclaim an idle workspace. + if (!session_.workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + if (!ensure_capacity()) { + return fail_segment("workspace capacity check"); + } + if (!session_.cut_cache_.bind_outputs(cache_outputs)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("output cache allocation"); + } + new_cache_bytes = add_bytes(last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph)); + if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { + pin_multi_device_nodes(scheduler, current); + })) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace allocation"); + } + for (const auto& size : measurement.buffers) { + track_compute_buffer(size.backend); + } + if (session_.workspace_.scheduler() != nullptr) { + track_compute_buffer(session_.workspace_.cpu_backend()); + } + if (!ensure_capacity()) { + return fail_segment("allocated capacity check"); + } + copy_data_to_backend_tensor(segment_graph); + auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!prefetch_requests.empty()) { + weights.enqueue_next(index, prefetch_requests.front()); + } + LOG_DEBUG("%s executing segment %zu/%zu: %s", task_.description.c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str()); + if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) { + return fail_segment("execution"); + } + auto cache_status = session_.cache_.capture(segment_graph); + if (cache_status == GGML_STATUS_SUCCESS) { + cache_status = session_.cut_cache_.capture(graph, segment, task_.description.c_str()); + } + if (cache_status != GGML_STATUS_SUCCESS) { + last_compute_status_ = cache_status; + return fail_segment("output caching"); + } + sync_runtime_residency(); + if (last) { + if (read_outputs && !read_outputs()) { + return fail_segment("output finalization"); + } + if (!no_return) { + output_ = read_graph_tensor(task_.output, "output"); + if (output_.empty()) { + return fail_segment("output readback"); + } + } + } + if (!session_.workspace_.segment_end()) { + last_compute_status_ = GGML_STATUS_FAILED; + return fail_segment("workspace synchronization"); + } + // Final outputs and their callbacks may still be views of consumed cuts. + session_.cut_cache_.prune(segment.future_cut_names); + } + if (segments_changed || peak_compute_bytes != session_.logged_compute_bytes_) { + for (const auto& entry : peak_compute_bytes) { + LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", + task_.description.c_str(), entry.second / (1024.0 * 1024.0), + sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + session_.logged_compute_bytes_ = std::move(peak_compute_bytes); + session_.logged_segment_count_ = plan.segments.size(); + } + return true; + } + + ExecutionSession::ExecutionSession(ggml_backend_t backend, const std::shared_ptr& manager) + : backend_(backend), residency_manager_(manager), cache_(backend), cut_cache_(backend), workspace_(backend) { + GGML_ASSERT(backend_ != nullptr); + } + + size_t ExecutionSession::resident_bytes(ggml_backend_t backend) const { + size_t bytes = workspace_.bytes(backend); + if (backend == backend_) { + const auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; }; + bytes = add(bytes, cache_.resident_bytes(ggml_backend_get_device(backend))); + bytes = add(bytes, cut_cache_.resident_bytes(ggml_backend_get_device(backend))); + } + return bytes; + } + + void ExecutionSession::set_layer_split_enabled(bool enabled) { + layer_split_enabled_ = enabled; + if (!enabled) { + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + } + + void ExecutionSession::set_backend_limits(const std::vector& limits) { + backend_limits_ = limits; + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + + void ExecutionSession::set_runtime_backends(const std::vector& backends) { + extra_backends_.clear(); + for (auto backend : backends) { + if (backend != nullptr && backend != backend_ && + std::find(extra_backends_.begin(), extra_backends_.end(), backend) == extra_backends_.end()) { + extra_backends_.push_back(backend); + } + } + workspace_.set_extra_backends(extra_backends_); + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + + ExecutionEngine::~ExecutionEngine() { + for (auto& entry : sessions_) { + end(*entry.second); + } + } + + ExecutionSession& ExecutionEngine::create_session(ggml_backend_t backend, const std::shared_ptr& manager) { + GGML_ASSERT(manager.get() == manager_); + auto session = std::make_unique(backend, manager); + auto& result = *session; + sessions_.emplace(session->owner_id(), std::move(session)); + return result; + } + + void ExecutionEngine::release_session(ExecutionSession& session) { + end(session); + sessions_.erase(session.owner_id()); + } + + bool ExecutionEngine::start(ExecutionSession& session) { + if (session.started_) { + return true; + } + sd_backend_set_buffer_cache_enabled(session.backend_, true); + for (auto backend : session.extra_backends_) { + sd_backend_set_buffer_cache_enabled(backend, true); + } + session.workspace_.set_extra_backends(session.extra_backends_); + if (auto manager = session.residency_manager_.lock()) { + manager->register_runtime_owner(session.owner_id()); + } + session.started_ = true; + return true; + } + + void ExecutionEngine::end(ExecutionSession& session) { + GGML_ASSERT(!session.executing_); + if (!session.started_) { + return; + } + session.workspace_.release(); + session.cache_.clear(); + session.cut_cache_.clear(); + session.logged_compute_bytes_.clear(); + session.logged_segment_count_ = 0; + if (auto manager = session.residency_manager_.lock()) { + manager->clear_prefetched_params(session.owner_id()); + std::vector params(session.params_.begin(), session.params_.end()); + manager->evict_compute_backend_params(params); + manager->remove_runtime_owner(session.owner_id()); + } + session.params_.clear(); + sd_backend_clear_buffer_cache(session.backend_); + for (auto backend : session.extra_backends_) { + sd_backend_clear_buffer_cache(backend); + } + session.started_ = false; + } + + void ExecutionEngine::sync_residency(ExecutionSession& session) { + if (auto manager = session.residency_manager_.lock()) { + manager->update_runtime_residency(session.owner_id(), session.backend_, session.resident_bytes(session.backend_)); + for (auto backend : session.extra_backends_) { + manager->update_runtime_residency(session.owner_id(), backend, session.resident_bytes(backend)); + } + } + } + + bool ExecutionEngine::ensure_capacity( + DeviceResidencyManager& manager, + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) { + GGML_ASSERT(&manager == manager_); + auto capacity = [&]() { return manager.inspect_compute_backend_capacity(request, required_params); }; + const auto initial = capacity(); + if (!initial.valid) { + return false; + } + if (initial.fits()) { + return true; + } + for (auto& entry : sessions_) { + auto& session = *entry.second; + if (session.owner_id() != request.owner_id && session.workspace_.release()) { + sync_residency(session); + if (capacity().fits(true)) { + return true; + } + } + } + auto protected_weights = protected_params; + protected_weights.insert(protected_weights.end(), required_params.begin(), required_params.end()); + auto candidates = manager.evictable_compute_backend_blocks(request.compute_backend, protected_weights); + std::unordered_map preferred_positions; + for (size_t index = 0; index < preferred_eviction_order.size(); ++index) { + for (auto* tensor : preferred_eviction_order[index]) { + auto* weight = manager.resolve_param_tensor(tensor); + preferred_positions.emplace(weight != nullptr ? weight : tensor, index); + } + } + auto preferred_position = [&](const WeightResidencyBlock& block) { + size_t position = 0; + // A shared allocation follows the last of its members in the eviction order. + for (auto* tensor : block.tensors) { + const auto found = preferred_positions.find(tensor); + if (found == preferred_positions.end()) { + return SIZE_MAX; + } + position = std::max(position, found->second); + } + return position; + }; + std::stable_sort(candidates.begin(), candidates.end(), [&](const auto& lhs, const auto& rhs) { + if (lhs.requires_demotion != rhs.requires_demotion) { + return !lhs.requires_demotion; + } + if (!lhs.requires_demotion) { + const auto left = preferred_position(lhs), right = preferred_position(rhs); + if (left != right) { + return left < right; + } + if (left != SIZE_MAX) { + return false; + } + } + return lhs.last_use_epoch < rhs.last_use_epoch; + }); + for (const auto& block : candidates) { + if (manager.evict_compute_backend_block(block) && capacity().fits(true)) { + return true; + } + } + auto available = capacity(); + if (available.fits(true)) { + return true; + } + if (sd_backend_is(request.compute_backend, "Vulkan") && + available.required_budget_bytes <= available.available_budget_bytes) { + ggml_backend_synchronize(request.compute_backend); + sd_backend_clear_buffer_cache(request.compute_backend); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); + do { + available = capacity(); + if (available.fits(true)) { + return true; + } + if (available.required_budget_bytes > available.available_budget_bytes) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } while (std::chrono::steady_clock::now() < deadline); + } + available = manager.inspect_compute_backend_capacity(request, required_params, true); + const auto device_bytes = available.available_device_bytes == SIZE_MAX + ? std::string("unknown") + : sd_format("%.2f MB", available.available_device_bytes / (1024.0 * 1024.0)); + const auto budget_bytes = available.available_budget_bytes == SIZE_MAX + ? std::string("unlimited") + : sd_format("%.2f MB", available.available_budget_bytes / (1024.0 * 1024.0)); + LOG_WARN("execution engine cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget", + ggml_backend_name(request.compute_backend), + available.required_device_bytes / (1024.0 * 1024.0), + available.required_budget_bytes / (1024.0 * 1024.0), device_bytes.c_str(), budget_bytes.c_str()); + return false; + } + + void ExecutionEngine::clear_cache(ExecutionSession& session) { + session.cache_.clear(); + sync_residency(session); + } + + void ExecutionEngine::discard_graph(ExecutionSession& session) { + GGML_ASSERT(!session.executing_); + session.cache_.graph_end(false); + session.cut_cache_.clear(); + session.workspace_.release(); + sync_residency(session); + } + + ComputeResult ExecutionEngine::execute(ExecutionSession& session, const ComputeTask& task) { + if (session.executing_ || task.graph == nullptr || !start(session)) { + return {}; + } + session.executing_ = true; + ComputeResult result; + try { + GraphExecution execution(*this, session, task); + result = execution.run(); + } catch (const std::bad_alloc&) { + result.status = GGML_STATUS_ALLOC_FAILED; + LOG_ERROR("%s graph allocation failed", task.description.c_str()); + } catch (const std::exception& error) { + LOG_ERROR("%s graph execution failed on %s: %s", task.description.c_str(), + ggml_backend_name(session.backend_), error.what()); + } catch (...) { + LOG_ERROR("%s graph execution failed on %s: unknown exception", task.description.c_str(), + ggml_backend_name(session.backend_)); + } + if (!session.workspace_.segment_end()) { + result.status = GGML_STATUS_FAILED; + } + const bool success = result.status == GGML_STATUS_SUCCESS; + session.cache_.graph_end(success); + session.cut_cache_.clear(); + if (!success) { + result.output = {}; + session.workspace_.release(); + } + session.executing_ = false; + sync_residency(session); + return result; + } +} diff --git a/src/core/execution_engine.h b/src/core/execution_engine.h new file mode 100644 index 00000000..79ced1d0 --- /dev/null +++ b/src/core/execution_engine.h @@ -0,0 +1,107 @@ +#ifndef __SD_CORE_EXECUTION_ENGINE_H__ +#define __SD_CORE_EXECUTION_ENGINE_H__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/compute_workspace.h" +#include "core/ggml_graph_cut.h" +#include "core/runner_cache.h" +#include "core/tensor.hpp" +#include "device_residency_manager.h" + +namespace sd { + struct ComputeTask { + std::string description; + ggml_cgraph* graph = nullptr; + ggml_tensor* output = nullptr; + std::vector weights; + std::map inputs; + std::vector> debug_tensors; + int n_threads = 1; + bool no_return = false; + std::function read_outputs; + }; + + struct ComputeResult { + ggml_status status = GGML_STATUS_FAILED; + Tensor output; + }; + + class ExecutionSession { + friend class ExecutionEngine; + friend class GraphExecution; + + ggml_backend_t backend_; + std::weak_ptr residency_manager_; + RunnerCache cache_; + GraphCutTensorCache cut_cache_; + ComputeWorkspace workspace_; + std::unordered_set params_; + std::vector extra_backends_; + size_t max_backend_bytes_ = 0; + bool layer_split_enabled_ = false; + std::vector backend_limits_; + ggml_graph_cut::PlanCache plan_cache_; + std::unordered_map weight_assignments_; + std::unordered_map node_assignments_; + bool primary_notice_logged_ = false; + std::map logged_compute_bytes_; + size_t logged_segment_count_ = 0; + bool started_ = false; + bool executing_ = false; + + size_t resident_bytes(ggml_backend_t backend) const; + uintptr_t owner_id() const { return reinterpret_cast(this); } + + public: + ExecutionSession(ggml_backend_t backend, const std::shared_ptr& manager); + ExecutionSession(const ExecutionSession&) = delete; + ExecutionSession& operator=(const ExecutionSession&) = delete; + + RunnerCache& cache() { return cache_; } + bool started() const { return started_; } + size_t compute_bytes() const { return workspace_.bytes(backend_); } + void clear_plan_cache() { plan_cache_.graph_cut_plans.clear(); } + void set_max_backend_bytes(size_t bytes) { max_backend_bytes_ = bytes; } + void set_layer_split_enabled(bool enabled); + void set_backend_limits(const std::vector& limits); + void set_runtime_backends(const std::vector& backends); + }; + + class ExecutionEngine { + friend class GraphExecution; + + std::map> sessions_; + DeviceResidencyManager* manager_; + + void sync_residency(ExecutionSession& session); + + public: + explicit ExecutionEngine(DeviceResidencyManager* manager = nullptr) + : manager_(manager) {} + ~ExecutionEngine(); + ExecutionSession& create_session(ggml_backend_t backend, const std::shared_ptr& manager); + void release_session(ExecutionSession& session); + bool start(ExecutionSession& session); + void end(ExecutionSession& session); + void clear_cache(ExecutionSession& session); + void discard_graph(ExecutionSession& session); + bool ensure_capacity(DeviceResidencyManager& manager, + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params); + // Graphs, input data and output callbacks remain borrowed until execute returns. + ComputeResult execute(ExecutionSession& session, const ComputeTask& task); + }; +} + +#endif // __SD_CORE_EXECUTION_ENGINE_H__ diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 39d54297..7a278119 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -1,15 +1,11 @@ #include #include -#include #include #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" #include "core/ggml_runner.h" #include "core/ggml_tensor_utils.h" -#include "core/layer_split_partition.h" -#include "core/segment_graph_bindings.h" -#include "core/segment_weight_pipeline.h" using namespace sd; @@ -101,16 +97,6 @@ std::vector GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf return used_params; } -void GGMLRunner::evict_compute_backend_param_tensors(const std::vector& tensors) { - if (tensors.empty()) { - return; - } - auto manager = residency_manager.lock(); - if (manager != nullptr) { - manager->evict_compute_backend_params(tensors); - } -} - void GGMLRunner::prepare_build_in_tensor_before() { one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1); ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one"); @@ -148,7 +134,7 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) { ggml_build_forward_expand(gf, entry.first); } } - for (const auto& entry : cache_.outputs()) { + for (const auto& entry : execution_session_.cache().outputs()) { if (entry.second != nullptr) { ggml_build_forward_expand(gf, entry.second); } @@ -172,365 +158,22 @@ bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph, return true; } -ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { - if (tensor == nullptr || tensor->buffer == nullptr) { - return nullptr; - } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || - ggml_backend_buffer_is_host(tensor->buffer)) { - return nullptr; - } - ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); - if (dev == nullptr) { - return nullptr; - } - if (ggml_backend_get_device(runtime_backend) == dev) { - return runtime_backend; - } - for (ggml_backend_t backend : extra_runtime_backends) { - if (ggml_backend_get_device(backend) == dev) { - return backend; - } - } - return nullptr; -} - -static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { - // Vulkan reports a preferred allocation chunk size, not a hard tensor limit. - if (sd_backend_is(backend, "Vulkan")) { - return true; - } - auto buft = ggml_backend_get_default_buffer_type(backend); - size_t alignment = ggml_backend_buft_get_alignment(buft); - size_t max_size = ggml_backend_buft_get_max_size(buft); - max_size -= max_size % alignment; - if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { - return false; - } - // Cross-backend copies must also fit the destination's allocation limit. - for (auto src : node->src) { - if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { - return false; - } - } - return true; -} - -void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { - if (sched == nullptr || gf == nullptr) { - return; - } - ggml_backend_t current = runtime_backend; - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); - if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { - current = node_assignment->second; - } - for (int s = 0; s < GGML_MAX_SRC; s++) { - ggml_backend_t weight_backend = backend_for_weight(node->src[s]); - if (weight_backend != nullptr) { - if (node_assignment == graph_cut_layer_split_node_assignments_.end()) { - current = weight_backend; - } - } - } - if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || - node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { - continue; - } - if (!backend_can_allocate_op(current, node)) { - auto cpu = workspace_.cpu_backend(); - if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { - ggml_backend_sched_set_tensor_backend(sched, node, cpu); - } - } else if (ggml_backend_supports_op(current, node)) { - ggml_backend_sched_set_tensor_backend(sched, node, current); - } - } -} - -size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const { - backend = backend == nullptr ? runtime_backend : backend; - size_t bytes = workspace_.bytes(backend); - if (backend == runtime_backend) { - const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend)); - bytes = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes; - const size_t cut_bytes = cut_cache_.resident_bytes(ggml_backend_get_device(backend)); - bytes = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes; - } - return bytes; -} - -void GGMLRunner::sync_runtime_residency() { - if (auto manager = residency_manager.lock()) { - manager->update_runtime_residency(reinterpret_cast(this), - runtime_backend, retained_runtime_buffer_bytes()); - for (auto backend : extra_runtime_backends) { - manager->update_runtime_residency(reinterpret_cast(this), - backend, retained_runtime_buffer_bytes(backend)); - } - } -} - -std::optional> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) { - if (tensor == nullptr) { - LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label); - return std::nullopt; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_ERROR("%s %s tensor type mismatch: got %s", - get_desc().c_str(), - label, - ggml_type_name(tensor->type)); - return std::nullopt; - } - ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); - if (buf == nullptr) { - LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", - get_desc().c_str(), - label, - tensor->name[0] != '\0' ? tensor->name : "", - ggml_op_name(tensor->op), - tensor->buffer, - tensor->view_src, - tensor->view_src ? tensor->view_src->buffer : nullptr, - tensor->data); - return std::nullopt; - } - - return sd::make_sd_tensor_from_ggml(tensor); -} - -void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) { - GGML_ASSERT(gf != nullptr); - std::unordered_set graph_tensor_set; - const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_nodes = ggml_graph_n_nodes(gf); - graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); - for (int i = 0; i < n_leafs; ++i) { - graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_nodes; ++i) { - graph_tensor_set.insert(ggml_graph_node(gf, i)); - } - - for (auto& kv : backend_tensor_data_map) { - auto tensor = kv.first; - auto data = kv.second; - if (tensor == nullptr || data == nullptr) { - continue; - } - const char* name = ggml_get_name(tensor); - if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { - continue; - } - if (tensor->buffer == nullptr) { - LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", - get_desc().c_str(), - name != nullptr ? name : "", - (long long)tensor->ne[0], - (long long)tensor->ne[1], - (long long)tensor->ne[2], - (long long)tensor->ne[3], - ggml_type_name(tensor->type)); - continue; - } - - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (buf == nullptr) { - LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", - get_desc().c_str(), - tensor && tensor->name[0] != '\0' ? tensor->name : "", - tensor ? ggml_op_name(tensor->op) : "", - data, - tensor ? tensor->view_src : nullptr, - (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); - continue; - } - - ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); - } - - if (clear_after_copy) { - backend_tensor_data_map.clear(); - } -} - -const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) { - GGML_ASSERT(gf != nullptr); - return sd::ggml_graph_cut::resolve_plan(runtime_backend, - gf, - &graph_cut_plan_cache_, - params_tensor_set_, - get_desc().c_str()); -} - -const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) { - return resolve_graph_cut_plan(gf); -} - -bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { - graph_cut_layer_split_node_assignments_.clear(); - if (!graph_cut_layer_split_enabled) { - return true; - } - if (!is_multi_device()) { - LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str()); - return false; - } - - const auto& plan = resolve_graph_cut_layer_split_plan(gf); - if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { - auto manager = residency_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - std::vector graph_params = collect_used_param_tensors(gf); - if (!graph_params.empty() && - !manager->assign_compute_backend(graph_params, runtime_backend)) { - LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str()); - return false; - } - for (ggml_tensor* param : graph_params) { - if (param != nullptr) { - graph_cut_layer_split_assignments_[param] = runtime_backend; - } - } - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - if (node != nullptr) { - graph_cut_layer_split_node_assignments_[node] = runtime_backend; - } - } - if (!graph_cut_layer_split_primary_notice_logged_) { - LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - graph_cut_layer_split_primary_notice_logged_ = true; - } else { - LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - } - return true; - } - - std::vector split_backends; - split_backends.reserve(extra_runtime_backends.size() + 1); - split_backends.push_back(runtime_backend); - for (ggml_backend_t backend : extra_runtime_backends) { - if (backend != nullptr) { - split_backends.push_back(backend); - } - } - - auto manager = residency_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - - sd::GraphCutLayerSplitAssignment assignment; - auto canonicalize_param = [this](ggml_tensor* tensor) { - return canonical_param_tensor(tensor); - }; - if (!sd::partition_graph_cut_layer_split(get_desc().c_str(), - gf, - plan, - split_backends, - graph_cut_layer_split_backend_vram_limits_, - max_graph_vram_bytes, - graph_cut_layer_split_assignments_, - canonicalize_param, - &assignment)) { - return false; - } - - for (size_t i = 0; i < split_backends.size(); i++) { - if (assignment.tensors_by_backend[i].empty()) { - continue; - } - if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { - LOG_ERROR("%s graph-cut layer split failed to assign params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); - return false; - } - } - - graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments); - sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment); - - return true; -} - bool GGMLRunner::runner_start() { - if (runner_started_) { - return true; - } - sd_backend_set_buffer_cache_enabled(runtime_backend, true); - for (auto backend : extra_runtime_backends) { - sd_backend_set_buffer_cache_enabled(backend, true); - } - cache_.clear(); - workspace_.set_extra_backends(extra_runtime_backends); - if (auto manager = residency_manager.lock()) { - manager->set_workspace_reclaimer(reinterpret_cast(this), [this]() { - if (!workspace_.release()) { - return false; - } - sync_runtime_residency(); - return true; - }); - } - runner_started_ = true; - return true; + return execution_engine_->start(execution_session_); } void GGMLRunner::runner_end() { GGML_ASSERT(!graph_active_); - if (!runner_started_) { - return; - } - workspace_.release(); + execution_engine_->end(execution_session_); free_compute_ctx(); compute_ctx_buffer_.reset(); - cache_.clear(); - cut_cache_.clear(); - logged_compute_bytes_.clear(); - logged_segment_count_ = 0; - if (auto manager = residency_manager.lock()) { - manager->clear_prefetched_params(reinterpret_cast(this)); - std::vector tensors; - for (auto tensor : params_tensor_set_) { - auto* parameter = manager->resolve_param_tensor(const_cast(tensor)); - if (parameter != nullptr) - tensors.push_back(parameter); - } - manager->evict_compute_backend_params(tensors); - manager->remove_runtime_owner(reinterpret_cast(this)); - } - sd_backend_clear_buffer_cache(runtime_backend); - for (auto backend : extra_runtime_backends) { - sd_backend_clear_buffer_cache(backend); - } - runner_started_ = false; } GGMLRunner::GGMLRunner(ggml_backend_t backend, std::shared_ptr manager) - : runtime_backend(backend), - cache_(backend), - cut_cache_(backend), - workspace_(backend), + : execution_engine_(manager != nullptr ? manager->execution_engine() : std::make_shared()), + execution_session_(execution_engine_->create_session(backend, manager)), + runtime_backend(backend), residency_manager(manager) { GGML_ASSERT(runtime_backend != nullptr); alloc_params_ctx(); @@ -538,6 +181,7 @@ GGMLRunner::GGMLRunner(ggml_backend_t backend, GGMLRunner::~GGMLRunner() { runner_end(); + execution_engine_->release_session(execution_session_); free_compute_ctx(); free_params_ctx(); } @@ -575,8 +219,7 @@ void GGMLRunner::reset_compute_ctx() { } void GGMLRunner::free_cache_ctx_and_buffer() { - cache_.clear(); - sync_runtime_residency(); + execution_engine_->clear_cache(execution_session_); } void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) { @@ -610,7 +253,7 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* if (tensor != nullptr) { ggml_set_output(tensor); } - cache_.stage(name, tensor); + execution_session_.cache().stage(name, tensor); if (graph != nullptr && tensor != nullptr) { // Schedule the cache output here so its source can be reused before graph end. ggml_build_forward_expand(graph, tensor); @@ -646,439 +289,54 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, GGMLRunner& runner; const bool& success; ~GraphEndGuard() { - if (!runner.workspace_.segment_end()) { - runner.last_compute_status_ = GGML_STATUS_FAILED; + if (!success) { + runner.execution_engine_->discard_graph(runner.execution_session_); } - runner.cache_.graph_end(false); - runner.cut_cache_.clear(); runner.free_compute_ctx(); runner.graph_active_ = false; - if (!success) { - runner.workspace_.release(); - } - runner.sync_runtime_residency(); } } graph_guard{*this, success}; - ggml_cgraph* graph = nullptr; - if (!prepare_compute_graph(get_graph, &graph)) { - return std::nullopt; - } - params_tensor_set_dirty_ = true; - rebuild_params_tensor_set(); - if (auto manager = residency_manager.lock()) { - for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) { - auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i)); - if (parameter != nullptr) - params_tensor_set_.insert(parameter); - } - } - std::optional> output; try { - output = execute_graph(graph, n_threads, no_return, read_outputs); + ggml_cgraph* graph = nullptr; + if (!prepare_compute_graph(get_graph, &graph)) { + return std::nullopt; + } + params_tensor_set_dirty_ = true; + rebuild_params_tensor_set(); + sd::ComputeTask task; + task.description = get_desc(); + task.graph = graph; + task.output = no_return ? nullptr : ggml_get_tensor(compute_ctx, final_result_name.c_str()); + task.weights = collect_used_param_tensors(graph); + task.inputs = backend_tensor_data_map; + task.debug_tensors = debug_tensors; + task.n_threads = n_threads; + task.no_return = no_return; + task.read_outputs = read_outputs; + auto result = execution_engine_->execute(execution_session_, task); + last_compute_status_ = result.status; + success = result.status == GGML_STATUS_SUCCESS; + if (success) { + return std::move(result.output); + } } catch (const std::bad_alloc&) { last_compute_status_ = GGML_STATUS_ALLOC_FAILED; LOG_ERROR("%s graph allocation failed", get_desc().c_str()); - return std::nullopt; } catch (const std::exception& error) { - last_compute_status_ = GGML_STATUS_FAILED; - LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), - ggml_backend_name(runtime_backend), error.what()); - return std::nullopt; + LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what()); } - success = output.has_value(); - if (success) { - cache_.graph_end(true); - last_compute_status_ = GGML_STATUS_SUCCESS; - } - return output; + return std::nullopt; } void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) { - graph_cut_layer_split_enabled = enabled; - if (!enabled) { - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; - } + execution_session_.set_layer_split_enabled(enabled); } void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) { - graph_cut_layer_split_backend_vram_limits_ = limits; - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; + execution_session_.set_backend_limits(limits); } void GGMLRunner::set_runtime_backends(const std::vector& backends) { - extra_runtime_backends.clear(); - for (ggml_backend_t backend : backends) { - if (backend == nullptr || backend == runtime_backend) { - continue; - } - if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == - extra_runtime_backends.end()) { - extra_runtime_backends.push_back(backend); - } - } - workspace_.set_extra_backends(extra_runtime_backends); - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; -} - -static size_t add_bytes(size_t a, size_t b) { - return b > SIZE_MAX - a ? SIZE_MAX : a + b; -} - -ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std::unordered_set* external_outputs) { - auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t { - if (external_outputs != nullptr && external_outputs->count(tensor)) { - return runtime_backend; - } - if (!params_tensor_set_.count(tensor)) { - return nullptr; - } - auto placement = graph_cut_layer_split_assignments_.find(tensor); - return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second; - }; - auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { - pin_multi_device_nodes(scheduler, copy, graph); - }; - return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device()); -} - -std::vector GGMLRunner::memory_requests( - const std::vector& sizes, - size_t pending_cache_bytes) const { - std::vector requests; - for (const auto& size : sizes) { - const size_t retained = retained_runtime_buffer_bytes(size.backend); - const size_t reusable = workspace_.bytes(size.backend); - const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0; - const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); - size_t limit = max_graph_vram_bytes; - if (is_multi_device()) { - size_t index = 0; - if (size.backend != runtime_backend) { - auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend); - index = static_cast(position - extra_runtime_backends.begin()) + 1; - } - if (index < graph_cut_layer_split_backend_vram_limits_.size()) { - limit = graph_cut_layer_split_backend_vram_limits_[index]; - } - } - requests.push_back({size.backend, reinterpret_cast(this), pending, - retained, limit}); - } - return requests; -} - -bool GGMLRunner::fits(const std::vector& requests, - const std::vector& params) const { - auto manager = residency_manager.lock(); - if (manager == nullptr) { - return params.empty(); - } - for (const auto& request : requests) { - if (!manager->fits_compute_backend_capacity(request, params)) { - return false; - } - } - return true; -} - -bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute) { - if (sd_backend_is_cpu(runtime_backend)) { - sd_backend_cpu_set_n_threads(runtime_backend, n_threads); - } - if (workspace_.cpu_backend() != nullptr) { - sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads); - } - auto scheduler = workspace_.scheduler(); - GGML_ASSERT(scheduler != nullptr); - ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data()); -#ifndef SD_USE_UPSTREAM_GGML - struct BeforeComputeCallback { - ggml_backend_t backend; - const std::function& callback; - bool invoked = false; - } before_compute{runtime_backend, on_before_compute}; - ggml_backend_sched_set_compute_callback( - scheduler, [](ggml_backend_t backend, void* data) { - auto& hook = *static_cast(data); - if (backend == hook.backend && !hook.invoked) { - hook.invoked = true; - hook.callback(); - } - }, - &before_compute); - const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); - ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr); -#else - on_before_compute(); - const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); -#endif - workspace_.synchronize(); - if (status != GGML_STATUS_SUCCESS) { - last_compute_status_ = status; - LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - return false; - } - const std::string description = get_desc(); - if (!debug_tensors.empty()) { - std::unordered_set graph_tensors; - const int leaf_count = ggml_graph_cut::leaf_count(graph); - const int node_count = ggml_graph_n_nodes(graph); - graph_tensors.reserve(static_cast(leaf_count + node_count)); - for (int index = 0; index < leaf_count; ++index) { - graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); - } - for (int index = 0; index < node_count; ++index) { - graph_tensors.insert(ggml_graph_node(graph, index)); - } - - for (const auto& entry : debug_tensors) { - ggml_tensor* tensor = entry.first; - if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { - continue; - } - ggml_backend_buffer_t buffer = - tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; - if (buffer == nullptr) { - LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", - description.c_str(), - entry.second.c_str()); - continue; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", - description.c_str(), - entry.second.c_str(), - ggml_type_name(tensor->type)); - continue; - } - auto debug_tensor = make_sd_tensor_from_ggml(tensor); - print_sd_tensor(debug_tensor, false, entry.second.c_str()); - } - } - - return true; -} - -std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { - if (!assign_graph_cut_layer_split_backends(graph)) { - return std::nullopt; - } - const auto params = collect_used_param_tensors(graph); - const auto& cached_plan = resolve_graph_cut_plan(graph); - auto manager = residency_manager.lock(); - const bool can_segment = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && - manager != nullptr && manager->segmented_compute_enabled() && - cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1; - // An empty workspace request is a lower bound on the full graph's memory demand. - bool segmented = can_segment && !fits(memory_requests({{runtime_backend, 0}}, 0), params); - ComputeWorkspace::Measurement full_measurement; - if (!segmented) { - full_measurement = measure(graph); - if (full_measurement.buffers.empty()) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return std::nullopt; - } - if (can_segment) { - // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits. - constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL; - auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)); - for (auto& request : requests) { - request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom); - } - segmented = !fits(requests, params); - } - } - ggml_graph_cut::Plan monolithic_plan; - if (!segmented) { - monolithic_plan.segments.emplace_back(); - auto& segment = monolithic_plan.segments.back(); - segment.group_name = "graph"; - segment.compute_buffer_size = cached_plan.compute_buffer_size; - segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph)); - segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph)); - for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { - segment.internal_node_indices.push_back(i); - } - for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { - auto tensor = ggml_graph_cut::leaf_tensor(graph, i); - ggml_graph_cut::Segment::InputRef input; - input.leaf_index = i; - input.type = canonical_param_tensor(tensor) != nullptr - ? ggml_graph_cut::Segment::INPUT_PARAM - : ggml_graph_cut::Segment::INPUT_EXTERNAL; - segment.input_refs.push_back(input); - } - } - const auto& plan = segmented ? cached_plan : monolithic_plan; - const bool segments_changed = plan.segments.size() != logged_segment_count_; - if (segments_changed && (segmented || logged_segment_count_ > 1)) { - LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(), - plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); - } - SegmentGraphBindings bindings(cut_cache_, plan, graph); - SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast(this), - graph, plan, params_tensor_set_, - segmented && manager != nullptr && manager->prefetch_enabled()); - - std::map peak_compute_bytes; - auto track_compute_buffer = [&](ggml_backend_t backend) { - if (backend != nullptr) { - auto& peak = peak_compute_bytes[backend]; - peak = std::max(peak, workspace_.bytes(backend)); - } - }; - std::optional> output = Tensor(); - for (size_t index = 0; index < plan.segments.size(); ++index) { - const auto& segment = plan.segments[index]; - const bool last = index + 1 == plan.segments.size(); - auto fail_segment = [&](const char* phase) { - LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(), - index + 1, plan.segments.size(), segment.group_name.c_str(), phase); - return std::nullopt; - }; - cut_cache_.prune(segment.live_cut_names); - bindings.reset(segment); - if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) { - return fail_segment("input binding"); - } - ggml_context* segment_context = nullptr; - auto segment_graph = segmented - ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) - : graph; - struct SegmentCleanup { - GGMLRunner& runner; - SegmentWeightPipeline& weights; - SegmentGraphBindings& bindings; - ggml_context* context; - ~SegmentCleanup() { - if (!runner.workspace_.segment_end()) { - runner.last_compute_status_ = GGML_STATUS_FAILED; - } - bindings.restore(); - weights.segment_end(); - ggml_free(context); - runner.sync_runtime_residency(); - } - } segment_cleanup{*this, weights, bindings, segment_context}; - - std::unordered_set cache_outputs; - if (segmented) { - for (int output_index : segment.output_node_indices) { - auto tensor = ggml_graph_node(graph, output_index); - if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) && - GraphCutTensorCache::can_bind_output(runtime_backend, tensor)) { - cache_outputs.insert(tensor); - } - } - } - auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement; - if (!workspace_.prepare(measurement)) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("workspace preparation"); - } - const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); - size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph)); - auto ensure_capacity = [&]() { - sync_runtime_residency(); - auto requests = memory_requests(measurement.buffers, new_cache_bytes); - if (fits(requests, weights.params(index))) { - return true; - } - if (workspace_.release_excess(measurement)) { - sync_runtime_residency(); - requests = memory_requests(measurement.buffers, new_cache_bytes); - } - const bool ready = weights.ensure_segment_capacity(index, requests); - if (!ready && manager != nullptr) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - } - return ready; - }; - if (!weights.segment_start(index, ensure_capacity)) { - return fail_segment("weight preparation"); - } - // Preparing weights can execute LoRA graphs and reclaim an idle workspace. - if (!workspace_.prepare(measurement)) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("workspace preparation"); - } - if (!ensure_capacity()) { - return fail_segment("workspace capacity check"); - } - if (!cut_cache_.bind_outputs(cache_outputs)) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("output cache allocation"); - } - new_cache_bytes = add_bytes(last ? 0 : cut_cache_.estimate_output_bytes(graph, segment), cache_.pending_bytes(segment_graph)); - if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { - pin_multi_device_nodes(scheduler, current); - })) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("workspace allocation"); - } - for (const auto& size : measurement.buffers) { - track_compute_buffer(size.backend); - } - if (workspace_.scheduler() != nullptr) { - track_compute_buffer(workspace_.cpu_backend()); - } - if (!ensure_capacity()) { - return fail_segment("allocated capacity check"); - } - copy_data_to_backend_tensor(segment_graph, false); - auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); - if (!prefetch_requests.empty()) { - weights.enqueue_next(index, prefetch_requests.front()); - } - LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), - index + 1, plan.segments.size(), segment.group_name.c_str()); - if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) { - return fail_segment("execution"); - } - auto cache_status = cache_.capture(segment_graph); - if (cache_status == GGML_STATUS_SUCCESS) { - cache_status = cut_cache_.capture(graph, segment, get_desc().c_str()); - } - if (cache_status != GGML_STATUS_SUCCESS) { - last_compute_status_ = cache_status; - return fail_segment("output caching"); - } - sync_runtime_residency(); - if (last) { - if (read_outputs && !read_outputs()) { - return fail_segment("output finalization"); - } - if (!no_return) { - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); - output = read_graph_tensor(result, "output"); - if (!output.has_value()) { - return fail_segment("output readback"); - } - } - } - if (!workspace_.segment_end()) { - last_compute_status_ = GGML_STATUS_FAILED; - return fail_segment("workspace synchronization"); - } - // Final outputs and their callbacks may still be views of consumed cuts. - cut_cache_.prune(segment.future_cut_names); - } - if (segments_changed || peak_compute_bytes != logged_compute_bytes_) { - for (const auto& entry : peak_compute_bytes) { - LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", - get_desc().c_str(), entry.second / (1024.0 * 1024.0), - sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), - plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); - } - logged_compute_bytes_ = std::move(peak_compute_bytes); - logged_segment_count_ = plan.segments.size(); - } - return output; + execution_session_.set_runtime_backends(backends); } diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index d300be76..95248242 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -12,9 +12,7 @@ #include #include -#include "core/compute_workspace.h" -#include "core/ggml_graph_cut.h" -#include "core/runner_cache.h" +#include "core/execution_engine.h" #include "core/tensor_ggml.hpp" #include "core/util.h" #include "device_residency_manager.h" @@ -132,43 +130,20 @@ struct GGMLRunnerContext { struct GGMLRunner { private: - std::map logged_compute_bytes_; - size_t logged_segment_count_ = 0; + std::shared_ptr execution_engine_; + sd::ExecutionSession& execution_session_; ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; - sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set* external_outputs = nullptr); - - std::vector memory_requests(const std::vector& sizes, - size_t pending_cache_bytes) const; - - bool fits(const std::vector& requests, - const std::vector& params) const; - - bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute); - - std::optional> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); - protected: typedef std::function get_graph_cb_t; - using GraphCutPlan = sd::ggml_graph_cut::Plan; ggml_backend_t runtime_backend = nullptr; ggml_context* params_ctx = nullptr; - sd::RunnerCache cache_; - sd::GraphCutTensorCache cut_cache_; - sd::ComputeWorkspace workspace_; ggml_context* compute_ctx = nullptr; std::unique_ptr compute_ctx_buffer_; - bool runner_started_ = false; - bool graph_active_ = false; - - size_t max_graph_vram_bytes = 0; - bool graph_cut_layer_split_enabled = false; - std::vector graph_cut_layer_split_backend_vram_limits_; - - std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) + bool graph_active_ = false; std::shared_ptr weight_adapter = nullptr; std::weak_ptr residency_manager; @@ -193,11 +168,7 @@ protected: bool circular_x_enabled = false; bool circular_y_enabled = false; - sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; std::unordered_set params_tensor_set_; - std::unordered_map graph_cut_layer_split_assignments_; - std::unordered_map graph_cut_layer_split_node_assignments_; - bool graph_cut_layer_split_primary_notice_logged_ = false; template static sd::Tensor take_or_empty(std::optional> tensor) { @@ -239,8 +210,6 @@ protected: std::vector collect_used_param_tensors(ggml_cgraph* gf); - void evict_compute_backend_param_tensors(const std::vector& tensors); - void prepare_build_in_tensor_before(); void prepare_build_in_tensor_after(ggml_cgraph* gf); @@ -252,39 +221,18 @@ protected: bool prepare_compute_graph(get_graph_cb_t get_graph, ggml_cgraph** gf_out); - ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const; - - // Weightless ops have no scheduler anchor, so pin them to the most recent - // weight device. Views must stay unpinned or cross-device copies can be - // skipped for their consumers. - void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr); - - bool is_multi_device() const { - return !extra_runtime_backends.empty(); - } - size_t reusable_compute_buffer_bytes() const { - return workspace_.bytes(runtime_backend); + return execution_session_.compute_bytes(); } - size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const; - - void sync_runtime_residency(); - - std::optional> read_graph_tensor(ggml_tensor* tensor, const char* label); - - void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true); - - const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* gf); - - const GraphCutPlan& resolve_graph_cut_layer_split_plan(ggml_cgraph* gf); - - bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf); + bool has_cache_tensors() const { + return !execution_session_.cache().empty(); + } public: bool runner_start(); - bool runner_started() const { return runner_started_; } + bool runner_started() const { return execution_session_.started(); } void runner_end(); @@ -335,7 +283,7 @@ public: void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr); ggml_tensor* get_cache_tensor_by_name(const std::string& name) { - return cache_.get(name); + return execution_session_.cache().get(name); } std::optional> compute(get_graph_cb_t get_graph, @@ -353,7 +301,7 @@ public: void set_sage_attention_enabled(bool enabled) { if (sage_attn_enabled != enabled) { free_cache_ctx_and_buffer(); - graph_cut_plan_cache_.graph_cut_plans.clear(); + execution_session_.clear_plan_cache(); sage_attn_enabled = enabled; } } @@ -381,7 +329,7 @@ public: } void set_max_graph_vram_bytes(size_t max_vram_bytes) { - max_graph_vram_bytes = max_vram_bytes; + execution_session_.set_max_backend_bytes(max_vram_bytes); } void set_graph_cut_layer_split_enabled(bool enabled); diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index faf7be27..73174101 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -3,12 +3,17 @@ #include #include +#include #include #include "ggml-backend.h" struct ggml_tensor; +namespace sd { + class ExecutionEngine; +} + enum class WeightPrefetchResult { Scheduled, AlreadyResident, @@ -36,13 +41,45 @@ struct DeviceMemoryRequest { } }; +struct DeviceMemoryCapacity { + size_t required_device_bytes = 0; + size_t required_budget_bytes = 0; + size_t available_device_bytes = SIZE_MAX; + size_t available_budget_bytes = SIZE_MAX; + bool requires_reclamation = false; + bool valid = true; + bool wrapped_device_report = false; + + bool fits(bool after_reclamation = false) const { + return valid && !requires_reclamation && + (after_reclamation || !wrapped_device_report) && + required_device_bytes <= available_device_bytes && + required_budget_bytes <= available_budget_bytes; + } +}; + +struct WeightResidencyBlock { + ggml_backend_buffer_t buffer = nullptr; + ggml_backend_t backend = nullptr; + std::vector tensors; + size_t bytes = 0; + uint64_t last_use_epoch = 0; + bool requires_demotion = false; +}; + struct DeviceResidencyManager { +private: + std::shared_ptr execution_engine_; + +public: virtual ~DeviceResidencyManager() = default; - virtual bool segmented_compute_enabled() const = 0; - virtual bool prefetch_enabled() const = 0; - virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) = 0; - virtual void remove_runtime_owner(uintptr_t owner_id) = 0; + std::shared_ptr execution_engine(); + + virtual bool segmented_compute_enabled() const = 0; + virtual bool prefetch_enabled() const = 0; + virtual void register_runtime_owner(uintptr_t owner_id) = 0; + virtual void remove_runtime_owner(uintptr_t owner_id) = 0; // Capacity requests select their backend's weights; protection spans all backends. virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, const std::vector& required_params) const = 0; @@ -62,6 +99,18 @@ struct DeviceResidencyManager { const std::vector& required_params, const std::vector>& preferred_eviction_order, const std::vector& protected_params) = 0; + virtual DeviceMemoryCapacity inspect_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& params, + bool log_details = false) const { + DeviceMemoryCapacity result; + result.valid = false; + return result; + } + virtual std::vector evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const { return {}; } + virtual bool evict_compute_backend_block(const WeightResidencyBlock& block) { return false; } virtual WeightPrefetchResult prefetch_params( uintptr_t owner_id, const std::vector& tensors) = 0; diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 7cd86fd3..00b88f4b 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -576,7 +576,7 @@ namespace Qwen { result = run(cache); } if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED && - (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) { + (cache.mode != QwenImage21PrefixCache::Mode::NONE || has_cache_tensors())) { // The failed graph has ended before persistent inputs are released. free_cache_ctx_and_buffer(); prefix_cache_disabled = true; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 8b618279..bfdf2625 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1,11 +1,9 @@ #include "model_manager.h" #include -#include #include #include #include -#include #include #include @@ -150,7 +148,7 @@ bool ModelManager::set_loras(std::vector loras, SDVersion version) { } if (lora_version_ == version && lora_specs_equal(loras_, loras)) return true; - if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { + if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { return state->pin_count != 0; })) { LOG_ERROR("cannot change LoRA configuration during execution"); @@ -166,7 +164,7 @@ bool ModelManager::set_loras(std::vector loras, SDVersion version) { std::set ModelManager::tensor_names() const { std::set names; for (const auto& state : tensor_states_) { - if (state != nullptr && state->component != ModelComponent::LoRA) { + if (state != nullptr && state->component != ModelComponent::LoRA && model_sources_.count(state->source_file) == 0) { names.insert(state->name); } } @@ -234,7 +232,7 @@ bool ModelManager::register_param_tensors(ModelComponent component, *registered_tensor_size += estimate_tensors_size(tensors); } - const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{}; + const auto scoped_sources = source_file != 0 ? source_loader(source_file).file_tensors(source_file, source_version) : String2TensorStorage{}; const auto& sources = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map(); std::unordered_set new_tensors; std::vector> new_states; @@ -424,6 +422,20 @@ bool ModelManager::validate_registered_tensors() { return ok; } +bool ModelManager::validate_param_tensors(const std::vector& tensors) { + std::vector states; + if (!resolve_required_tensor_states(tensors, states)) { + return false; + } + for (auto* state : states) { + if (!validate_tensor(*state)) { + return false; + } + state->metadata_validated = true; + } + return true; +} + bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { std::vector need_load; need_load.reserve(states.size()); @@ -500,7 +512,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { - std::map, std::vector> states_by_staging_target; + std::map, std::vector> states_by_staging_target; for (TensorState* state : states) { if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; @@ -526,12 +538,12 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorcompute_backend); } - states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state); + states_by_staging_target[{state->compute_backend, staging_buft, state->source_file}].push_back(state); } for (const auto& pair : states_by_staging_target) { - ggml_backend_t compute_backend = pair.first.first; - ggml_backend_buffer_type_t staging_buft = pair.first.second; + ggml_backend_t compute_backend = std::get<0>(pair.first); + ggml_backend_buffer_type_t staging_buft = std::get<1>(pair.first); const std::vector& target_states = pair.second; if (target_states.empty()) { continue; @@ -646,7 +658,7 @@ bool ModelManager::apply_loras_to_params(const std::vector& states std::map groups; for (TensorState* state : states) { if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA || - state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) { + state->component == ModelComponent::LatentUpsampler || model_sources_.count(state->source_file) != 0 || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; } if (state->applied_lora_epoch == current_lora_epoch_) { @@ -743,6 +755,9 @@ void ModelManager::reset_lora_applied_params() { } bool ModelManager::should_ignore(const TensorState& state) const { + if (model_sources_.count(state.source_file) != 0) { + return false; + } for (const auto& ignore_prefix : common_ignore_tensors_) { if (starts_with(state.name, ignore_prefix)) { return true; @@ -869,7 +884,7 @@ bool ModelManager::can_mmap_storage(const TensorState& state) const { bool ModelManager::alloc_params_buffers(const std::vector& states, std::vector& created_storage_blocks) { - std::map, std::vector> states_by_buffer_type; + std::map, std::vector> states_by_buffer_type; for (TensorState* state : states) { if (state == nullptr || state->tensor == nullptr) { continue; @@ -878,11 +893,11 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, if (params_buft == nullptr) { return false; } - states_by_buffer_type[{params_buft, static_cast(state->residency_mode)}].push_back(state); + states_by_buffer_type[{params_buft, static_cast(state->residency_mode), state->source_file}].push_back(state); } for (const auto& pair : states_by_buffer_type) { - ggml_backend_buffer_type_t params_buft = pair.first.first; + ggml_backend_buffer_type_t params_buft = std::get<0>(pair.first); const std::vector& states = pair.second; size_t alignment = ggml_backend_buft_get_alignment(params_buft); size_t max_size = ggml_backend_buft_get_max_size(params_buft); @@ -1010,7 +1025,7 @@ bool ModelManager::load_tensor_parts(TensorState& state) { }; const bool success = state.source_file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names, false) - : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_); + : source_loader(state.source_file).load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_); if (!success || loaded != names) { return false; } @@ -1098,7 +1113,7 @@ bool ModelManager::load_tensors(const std::vector& states) { }; const auto file = group.first.first; bool success = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names) - : model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_); + : source_loader(file).load_file_tensors(file, group.first.second, callback, names, enable_mmap_); if (!success || loaded != names) return false; for (auto& entry : batch) { @@ -1316,7 +1331,7 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { void ModelManager::release_all() { clear_all_prefetched_params(); runtime_residencies_.clear(); - workspace_reclaimers_.clear(); + runtime_owners_.clear(); for (auto& state : tensor_states_) { state->pin_count = 0; state->applied_lora_epoch = UINT64_MAX; @@ -1643,11 +1658,17 @@ void ModelManager::evict_compute_backend_params(const std::vector& if (tensors.empty()) { return; } - std::vector required_states; - if (!resolve_required_tensor_states(tensors, required_states)) { + std::unordered_set target_states; + // Session cleanup may follow unregistration, after tensor metadata has been freed. + for (auto* tensor : tensors) { + const auto found = tensor_states_by_tensor_.find(tensor); + if (found != tensor_states_by_tensor_.end()) { + target_states.insert(found->second); + } + } + if (target_states.empty()) { return; } - std::unordered_set target_states(required_states.begin(), required_states.end()); for (const auto& block : compute_staging_blocks_) { const bool intersects = std::any_of( @@ -1743,12 +1764,12 @@ WeightResidencyInfo ModelManager::inspect_compute_backend_params( return info; } -void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) { - workspace_reclaimers_[owner_id] = std::move(reclaim); +void ModelManager::register_runtime_owner(uintptr_t owner_id) { + runtime_owners_.insert(owner_id); } void ModelManager::remove_runtime_owner(uintptr_t owner_id) { - workspace_reclaimers_.erase(owner_id); + runtime_owners_.erase(owner_id); for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { if (it->first.first == owner_id) { it = runtime_residencies_.erase(it); @@ -1806,8 +1827,8 @@ ModelManager::CapacityCheck ModelManager::check_capacity( // A wrapped report may reflect real pressure; reclaim before accepting the tracked estimate. LOG_VERBOSE("ignoring impossible free memory report on %s (%.2f MB free / %.2f MB total), using tracked residency", ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0)); - free_bytes = total_bytes; - result.requires_reclamation = true; + free_bytes = total_bytes; + result.wrapped_device_report = true; } if (total_bytes > 0) { free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0); @@ -1818,6 +1839,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( if (request.max_backend_bytes > 0) { const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), other_runtime_resident_bytes(request.owner_id, request.compute_backend)); + result.requires_reclamation = result.requires_reclamation || resident > request.max_backend_bytes; result.available_budget_bytes = resident < request.max_backend_bytes ? request.max_backend_bytes - resident : 0; @@ -1842,6 +1864,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( if (entry.second > 0) { const size_t resident = add(compute_backend_resident_bytes(entry.first), other_runtime_resident_bytes(request.owner_id, entry.first)); + result.requires_reclamation = result.requires_reclamation || resident > entry.second; const size_t available = resident < entry.second ? entry.second - resident : 0; result.available_budget_bytes = std::min(result.available_budget_bytes, available); } @@ -1849,6 +1872,139 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return result; } +DeviceMemoryCapacity ModelManager::inspect_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& params, + bool log_details) const { + std::vector states; + if (request.compute_backend == nullptr || !resolve_required_tensor_states(params, states, request.compute_backend)) { + DeviceMemoryCapacity result; + result.valid = false; + return result; + } + return check_capacity(request, states, log_details); +} + +std::vector ModelManager::evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const { + std::vector result; + if (backend == nullptr) { + return result; + } + std::vector protected_list; + if (!resolve_required_tensor_states(protected_params, protected_list)) { + return result; + } + std::unordered_set protected_states(protected_list.begin(), protected_list.end()); + for (const auto& entry : prefetch_blocks_) { + protected_states.insert(entry.second->states.begin(), entry.second->states.end()); + } + const auto device = ggml_backend_get_device(backend); + const bool cpu = sd_backend_is_cpu(backend); + auto on_device = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr) { + return false; + } + if (cpu) { + return ggml_backend_buffer_is_host(buffer); + } + if (ggml_backend_buffer_is_host(buffer)) { + return false; + } + const auto type = ggml_backend_buffer_get_type(buffer); + const auto split = split_buffer_devices_.find(type); + if (split != split_buffer_devices_.end()) { + return std::any_of(split->second.begin(), split->second.end(), [&](const auto& entry) { + return ggml_backend_get_device(entry.first) == device; + }); + } + return ggml_backend_buft_get_device(type) == device; + }; + auto append = [&](ggml_backend_buffer_t buffer, const std::vector& states, bool storage) { + if (!on_device(buffer) || states.empty()) { + return; + } + WeightResidencyBlock candidate; + candidate.buffer = buffer; + candidate.backend = backend; + candidate.bytes = ggml_backend_buffer_get_size(buffer); + for (auto* state : states) { + if (state == nullptr) { + continue; + } + if (state->pin_count != 0 || protected_states.count(state) != 0) { + return; + } + if (storage) { + if (state->staged_to_compute_backend) { + return; + } + if (state->residency_mode != ResidencyMode::Disk) { + if (cpu || state->params_backend == nullptr || sd_backend_is_cpu(state->params_backend) || + state->compute_backend != state->params_backend || + state->sources.empty() || state->applied_lora_epoch != UINT64_MAX || + state->split_buffer_type != nullptr || state->tensor->view_src != nullptr) { + return; + } + candidate.requires_demotion = true; + } + } + candidate.tensors.push_back(state->tensor); + candidate.last_use_epoch = std::max(candidate.last_use_epoch, state->last_use_epoch); + } + result.push_back(std::move(candidate)); + }; + for (const auto& block : compute_staging_blocks_) { + std::vector states; + states.reserve(block->staged_tensors.size()); + for (const auto& entry : block->staged_tensors) { + states.push_back(entry.first); + } + append(block->buffer, states, false); + } + for (const auto& block : params_storage_blocks_) { + // Loader-retained mappings cannot be reclaimed by freeing a storage block. + if (block->mmap_tensor_stores.empty()) { + append(block->buffer, block->states, true); + } + } + return result; +} + +bool ModelManager::evict_compute_backend_block(const WeightResidencyBlock& block) { + const auto candidates = evictable_compute_backend_blocks(block.backend, {}); + const auto found = std::find_if(candidates.begin(), candidates.end(), [&](const auto& candidate) { + return candidate.buffer == block.buffer && candidate.tensors == block.tensors && + candidate.requires_demotion == block.requires_demotion; + }); + if (found == candidates.end()) { + return false; + } + for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end(); ++it) { + if ((*it)->buffer == block.buffer) { + free_compute_staging_block(**it); + compute_staging_blocks_.erase(it); + return true; + } + } + for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end(); ++it) { + if ((*it)->buffer == block.buffer) { + if (block.requires_demotion) { + for (auto* state : (*it)->states) { + if (state != nullptr) { + state->residency_mode = ResidencyMode::Disk; + } + } + } + free_params_storage_block(**it); + params_storage_blocks_.erase(it); + return true; + } + } + return false; +} + bool ModelManager::fits_compute_backend_capacity( const DeviceMemoryRequest& request, const std::vector& required_params) const { @@ -1865,184 +2021,5 @@ bool ModelManager::ensure_compute_backend_capacity( const std::vector& required_params, const std::vector>& preferred_eviction_order, const std::vector& protected_params) { - std::vector required_states; - if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) { - return false; - } - - ggml_backend_t compute_backend = request.compute_backend; - if (compute_backend == nullptr) { - LOG_ERROR("model manager cannot reclaim memory for a null compute backend"); - return false; - } - if (sd_backend_is_cpu(compute_backend)) { - return true; - } - - auto fits = [&]() { - const auto capacity = check_capacity(request, required_states); - return !capacity.requires_reclamation && capacity.fits(); - }; - if (fits()) { - return true; - } - for (const auto& entry : workspace_reclaimers_) { - if (entry.first != request.owner_id) { - entry.second(); - if (fits()) { - return true; - } - } - } - - std::unordered_set protected_states; - std::vector resolved_protected; - if (!resolve_required_tensor_states(protected_params, resolved_protected)) { - return false; - } - protected_states.insert(resolved_protected.begin(), resolved_protected.end()); - for (const auto& entry : prefetch_blocks_) { - if (entry.second != nullptr) { - protected_states.insert(entry.second->states.begin(), entry.second->states.end()); - } - } - - std::unordered_set eviction_states; - auto add_evictable_state = [&](TensorState* state) { - if (state == nullptr || state->compute_backend != compute_backend || - state->pin_count > 0 || protected_states.find(state) != protected_states.end()) { - return; - } - const bool reloadable = state->residency_mode == ResidencyMode::Disk || - state->compute_backend != state->params_backend; - const bool resident = state->compute_backend == state->params_backend - ? state->loaded_to_params_backend - : state->staged_to_compute_backend; - if (reloadable && resident) { - eviction_states.insert(state); - } - }; - size_t last_eviction_count = 0; - auto release_eviction_states = [&]() { - if (eviction_states.size() != last_eviction_count) { - last_eviction_count = eviction_states.size(); - release_compute_staging_blocks(false, &eviction_states); - release_params_storage_blocks(false, &eviction_states); - } - // Device memory reports can lag behind releases even when no new candidates are added. - return fits(); - }; - - for (const auto& candidate_params : preferred_eviction_order) { - std::vector candidate_states; - if (!resolve_required_tensor_states(candidate_params, candidate_states)) { - return false; - } - for (TensorState* state : candidate_states) { - add_evictable_state(state); - } - if (release_eviction_states()) { - return true; - } - } - - std::vector global_candidates; - global_candidates.reserve(tensor_states_.size()); - for (const auto& state : tensor_states_) { - if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) { - global_candidates.push_back(state.get()); - } - } - std::stable_sort(global_candidates.begin(), - global_candidates.end(), - [](const TensorState* lhs, const TensorState* rhs) { - return lhs->last_use_epoch < rhs->last_use_epoch; - }); - for (TensorState* state : global_candidates) { - add_evictable_state(state); - if (release_eviction_states()) { - return true; - } - } - - // The loader retains mmap buffers after storage blocks are released, so - // demoting mapped states would only force duplicate allocations on reload. - std::unordered_set mapped_states; - for (const auto& block : params_storage_blocks_) { - if (!block->mmap_tensor_stores.empty()) { - mapped_states.insert(block->states.begin(), block->states.end()); - } - } - std::unordered_set required_set(required_states.begin(), required_states.end()); - std::vector demoted_states; - auto revert_demotions = [&]() { - for (TensorState* state : demoted_states) { - if (state->loaded_to_params_backend) { - state->residency_mode = ResidencyMode::ParamBackend; - } - } - }; - auto can_demote_to_disk = [&](TensorState* state) { - return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr && - state->residency_mode == ResidencyMode::ParamBackend && - state->params_backend != nullptr && state->params_backend == state->compute_backend && - state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) && - state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend && - protected_states.find(state) == protected_states.end() && - required_set.find(state) == required_set.end() && - mapped_states.find(state) == mapped_states.end() && - !state->sources.empty() && state->applied_lora_epoch == UINT64_MAX && - state->split_buffer_type == nullptr; - }; - for (TensorState* state : global_candidates) { - if (!can_demote_to_disk(state)) { - continue; - } - state->residency_mode = ResidencyMode::Disk; - demoted_states.push_back(state); - add_evictable_state(state); - if (release_eviction_states()) { - revert_demotions(); - return true; - } - } - revert_demotions(); - - auto capacity = check_capacity(request, required_states); - if (!capacity.requires_reclamation && capacity.fits()) { - return true; - } - if (sd_backend_is(compute_backend, "Vulkan") && - capacity.required_budget_bytes <= capacity.available_budget_bytes) { - ggml_backend_synchronize(compute_backend); - sd_backend_clear_buffer_cache(compute_backend); - // Vulkan heap usage can lag behind buffer releases. Fast memory queries - // must not turn that reporting delay into an allocation failure. - const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); - for (;;) { - capacity = check_capacity(request, required_states); - if (!capacity.requires_reclamation && capacity.fits()) { - return true; - } - if (capacity.required_budget_bytes > capacity.available_budget_bytes || - std::chrono::steady_clock::now() >= deadline) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(10)); - } - } - capacity = check_capacity(request, required_states, true); - const std::string available_device = capacity.available_device_bytes == SIZE_MAX - ? "unknown" - : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0)); - const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX - ? "unlimited" - : sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0)); - LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget", - ggml_backend_name(compute_backend), - capacity.required_device_bytes / (1024.0 * 1024.0), - capacity.required_budget_bytes / (1024.0 * 1024.0), - available_device.c_str(), - available_budget.c_str()); - return false; + return execution_engine()->ensure_capacity(*this, request, required_params, preferred_eviction_order, protected_params); } diff --git a/src/model_manager.h b/src/model_manager.h index 81d2546a..f592a63f 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -91,6 +91,8 @@ private: }; ModelLoader model_loader_; + // Auxiliary models retain their own catalogs and conversion settings. + std::map model_sources_; std::vector> tensor_states_; std::map tensor_states_by_tensor_; mutable std::list resolved_tensor_states_; @@ -101,7 +103,7 @@ private: std::map> prefetch_blocks_; std::map prefetch_backends_; std::map, RuntimeResidency> runtime_residencies_; - std::map> workspace_reclaimers_; + std::unordered_set runtime_owners_; bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; @@ -131,6 +133,7 @@ private: std::vector& required_states, ggml_backend_t compute_backend = nullptr) const; bool should_ignore(const TensorState& state) const; + const ModelLoader& source_loader(ModelLoader::FileId file) const; bool is_optional_missing_tensor(const std::string& name) const; bool validate_tensor(const TensorState& state) const; @@ -148,17 +151,7 @@ private: size_t compute_backend_alloc_size(const std::vector& states, bool missing_only) const; size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const; - struct CapacityCheck { - size_t required_device_bytes = 0; - size_t required_budget_bytes = 0; - size_t available_device_bytes = SIZE_MAX; - size_t available_budget_bytes = SIZE_MAX; - bool requires_reclamation = false; - bool fits() const { - return required_device_bytes <= available_device_bytes && - required_budget_bytes <= available_budget_bytes; - } - }; + using CapacityCheck = DeviceMemoryCapacity; CapacityCheck check_capacity(const DeviceMemoryRequest& request, const std::vector& states, bool log_details = false) const; @@ -186,6 +179,8 @@ public: bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false); bool del_file(ModelLoader::FileId id); bool refresh_files(); + ModelLoader::FileId register_model_source(ModelLoader loader); + bool unregister_model_source(ModelLoader::FileId file); ModelLoader::FileVersions source_versions(const std::set& components, const ModelLoader& loader) const; size_t registered_params_size(const std::set& components) const; @@ -267,12 +262,13 @@ public: } bool validate_registered_tensors(); + bool validate_param_tensors(const std::vector& tensors); bool load_all_params_eagerly(); bool assign_compute_backend(const std::vector& tensors, ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; - void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) override; + void register_runtime_owner(uintptr_t owner_id) override; void remove_runtime_owner(uintptr_t owner_id) override; bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, const std::vector& required_params) const override; @@ -290,6 +286,14 @@ public: const std::vector& required_params, const std::vector>& preferred_eviction_order, const std::vector& protected_params) override; + DeviceMemoryCapacity inspect_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& params, + bool log_details = false) const override; + std::vector evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const override; + bool evict_compute_backend_block(const WeightResidencyBlock& block) override; WeightPrefetchResult prefetch_params( uintptr_t owner_id, const std::vector& tensors) override; diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index adeea95d..d6453a21 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -74,7 +74,7 @@ void ModelManager::invalidate_sources(const std::unordered_set& st } bool ModelManager::set_loader(ModelLoader loader) { - if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { + if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { return state->pin_count != 0; })) { LOG_ERROR("cannot update model sources during execution"); @@ -87,7 +87,10 @@ bool ModelManager::set_loader(ModelLoader loader) { auto key = std::make_pair(state.source_file, state.source_version); auto found = scoped.find(key); if (found == scoped.end()) - found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first; + found = scoped.emplace(key, model_sources_.count(key.first) != 0 + ? source_loader(key.first).file_tensors(key.first, key.second) + : loader.file_tensors(key.first, key.second)) + .first; return found->second; }; bool lora_changed = false; @@ -147,7 +150,9 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::setsource_file != 0) { - versions[state->source_file] = loader.file_revision(state->source_file); + versions[state->source_file] = model_sources_.count(state->source_file) != 0 + ? source_loader(state->source_file).file_revision(state->source_file) + : loader.file_revision(state->source_file); continue; } for (const auto& source : find_tensor_sources(*state, sources)) { @@ -167,3 +172,34 @@ size_t ModelManager::registered_params_size(const std::set& comp } return bytes; } + +const ModelLoader& ModelManager::source_loader(ModelLoader::FileId file) const { + auto found = model_sources_.find(file); + return found == model_sources_.end() ? model_loader_ : found->second; +} + +ModelLoader::FileId ModelManager::register_model_source(ModelLoader loader) { + const auto versions = loader.file_versions(); + if (versions.size() != 1) { + LOG_ERROR("a model source must contain exactly one root model file"); + return 0; + } + const auto file = versions.begin()->first; + if (model_loader_.file_revision(file) != 0 || model_sources_.count(file) != 0) { + LOG_ERROR("model source is already registered"); + return 0; + } + loader.set_n_threads(n_threads_); + model_sources_.emplace(file, std::move(loader)); + return file; +} + +bool ModelManager::unregister_model_source(ModelLoader::FileId file) { + if (std::any_of(tensor_states_.begin(), tensor_states_.end(), [file](const auto& state) { + return state->source_file == file; + })) { + LOG_ERROR("cannot unregister a model source while its weights are registered"); + return false; + } + return model_sources_.erase(file) != 0; +} diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 1f7d5e45..50a4fb7a 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -9,10 +9,10 @@ #include "core/rng.hpp" #include "diffusion_engine.h" #include "model/vae/vae.hpp" +#include "pipeline/image_upscale.h" #include "request.h" #include "runtime/denoiser.hpp" #include "runtime/image_preprocess.h" -#include "upscaler.h" namespace sd::pipeline { @@ -751,7 +751,7 @@ namespace sd::pipeline { static sd::Tensor upscale_hires_latent(StableDiffusionGGML* sd, const sd::Tensor& latent, const GenerationRequest& request, - UpscalerGGML* upscaler) { + ImageUpscaleModel* upscaler) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling hires latent upscale"); return {}; @@ -837,7 +837,7 @@ namespace sd::pipeline { sd->generation_progress.report("upscaling", "hires"); sd::Tensor upscaled_tensor; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { - upscaled_tensor = upscaler->upscale_tensor(decoded); + upscaled_tensor = upscale_image_tensor(*upscaler, decoded, sd->n_threads, request.hires.upscale_tile_size); if (upscaled_tensor.empty()) { LOG_ERROR("hires model upscale failed"); return {}; @@ -1034,22 +1034,18 @@ namespace sd::pipeline { LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); sd->generation_progress.report("upscaling", "hires"); - std::unique_ptr hires_upscaler; + std::unique_ptr hires_upscaler; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling generation before hires model load"); return false; } LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); - hires_upscaler = std::make_unique(sd->n_threads, - false, - request.hires.upscale_tile_size, - sd->backend_spec, - sd->params_backend_spec); - const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); - hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (!hires_upscaler->load_from_file(request.hires.model_path, - sd->n_threads)) { + hires_upscaler = load_image_upscale_model(sd->model_manager, + sd->backend_manager, + request.hires.model_path, + sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER)); + if (hires_upscaler == nullptr) { LOG_ERROR("load hires model upscaler failed"); return false; } diff --git a/src/pipeline/image_upscale.cpp b/src/pipeline/image_upscale.cpp new file mode 100644 index 00000000..1ccb7b55 --- /dev/null +++ b/src/pipeline/image_upscale.cpp @@ -0,0 +1,138 @@ +#include "pipeline/image_upscale.h" + +#include "core/ggml_extend_backend.h" +#include "core/util.h" +#include "model/upscaler/esrgan.hpp" +#include "model/upscaler/mosr.hpp" +#include "model_manager.h" +#include "runtime/tiling.h" + +namespace sd::pipeline { + ImageUpscaleModel::ImageUpscaleModel(const std::shared_ptr& manager) + : manager_(manager) {} + + ImageUpscaleModel::~ImageUpscaleModel() { + if (auto* value = runner()) { + value->runner_end(); + } + if (auto manager = manager_.lock()) { + GGML_ASSERT(manager->unregister_param_tensors(weights_)); + if (source_file_ != 0) { + GGML_ASSERT(manager->unregister_model_source(source_file_)); + } + } + } + + GGMLRunner* ImageUpscaleModel::runner() const { + return mosr_ != nullptr ? static_cast(mosr_.get()) : esrgan_.get(); + } + + int ImageUpscaleModel::scale() const { + if (mosr_ != nullptr) { + return mosr_->config.scale; + } + return esrgan_ != nullptr ? esrgan_->config.scale : 1; + } + + std::unique_ptr load_image_upscale_model( + const std::shared_ptr& manager, + SDBackendManager& backends, + const std::string& path, + size_t max_vram_bytes, + bool conv_direct, + ggml_type weight_type) { + if (manager == nullptr) { + return nullptr; + } + const auto backend = backends.runtime_backend(SDBackendModule::UPSCALER); + const auto params_backend = backends.params_backend(SDBackendModule::UPSCALER); + if (backend == nullptr || params_backend == nullptr) { + LOG_ERROR("failed to initialize upscaler backends"); + return nullptr; + } + ModelLoader loader; + if (!loader.init_from_file_and_convert_name(path, "", VERSION_ESRGAN)) { + LOG_ERROR("init upscaler model loader failed: '%s'", path.c_str()); + return nullptr; + } + loader.set_wtype_override(weight_type); + const auto tensors = loader.get_tensor_storage_map(); + auto model = std::unique_ptr(new ImageUpscaleModel(manager)); + model->source_file_ = manager->register_model_source(std::move(loader)); + if (model->source_file_ == 0) { + return nullptr; + } + std::map weights; + if (tensors.count("gblocks.0.weight") != 0) { + model->mosr_ = std::make_unique(backend, tensors, manager); + if (model->mosr_->mosr_net == nullptr) { + LOG_ERROR("unsupported MoSR weights: '%s'", path.c_str()); + return nullptr; + } + model->mosr_->get_param_tensors(weights); + } else { + model->esrgan_ = std::make_unique(backend, tensors, manager); + if (model->esrgan_->rrdb_net == nullptr) { + LOG_ERROR("unsupported ESRGAN weights: '%s'", path.c_str()); + return nullptr; + } + model->esrgan_->get_param_tensors(weights); + } + for (const auto& entry : weights) { + model->weights_.push_back(entry.second); + } + const auto mode = backends.params_backend_is_disk(SDBackendModule::UPSCALER) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend; + if (!manager->register_param_tensors(ModelComponent::Upscaler, std::move(weights), mode, + backend, params_backend, nullptr, false, false, nullptr, + model->source_file_, VERSION_ESRGAN) || + !manager->validate_param_tensors(model->weights_)) { + LOG_ERROR("register upscaler weights failed: '%s'", path.c_str()); + return nullptr; + } + model->runner()->set_max_graph_vram_bytes(max_vram_bytes); + model->runner()->set_conv2d_direct_enabled(conv_direct); + LOG_INFO("Upscaler architecture: %s, weight type: %s", model->runner()->get_desc().c_str(), ggml_type_name(weight_type)); + return model; + } + + Tensor upscale_image_tensor(ImageUpscaleModel& model, const Tensor& input, int n_threads, int tile_size) { + if (input.empty() || input.dim() != 4 || (input.shape()[2] != 3 && input.shape()[2] != 4)) { + LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor"); + return {}; + } + struct RunnerEnd { + GGMLRunner* runner; + ~RunnerEnd() { runner->runner_end(); } + } cleanup{model.runner()}; + const bool has_alpha = input.shape()[2] == 4; + auto rgb = has_alpha ? sd::ops::slice(input, 2, 0, 3) : input; + const auto compute = [&](const Tensor& tensor) { + return model.mosr_ != nullptr ? model.mosr_->compute(n_threads, tensor) : model.esrgan_->compute(n_threads, tensor); + }; + const int scale = model.scale(); + Tensor output; + if (tile_size <= 0 || (input.shape()[0] <= tile_size && input.shape()[1] <= tile_size)) { + output = compute(rgb); + } else { + output = process_tiles_2d(rgb, + static_cast(input.shape()[0] * scale), + static_cast(input.shape()[1] * scale), + scale, tile_size, tile_size, 0.25f, false, false, compute); + } + if (output.empty()) { + LOG_ERROR("upscaler compute failed"); + return {}; + } + if (has_alpha) { + auto alpha = sd::ops::slice(input, 2, 3, 4); + auto shape = alpha.shape(); + shape[0] = output.shape()[0]; + shape[1] = output.shape()[1]; + alpha = sd::ops::interpolate(alpha, shape, sd::ops::InterpolateMode::Bilinear); + output = sd::ops::concat(output, alpha, 2); + } + return output; + } +} diff --git a/src/pipeline/image_upscale.h b/src/pipeline/image_upscale.h new file mode 100644 index 00000000..1de9ca00 --- /dev/null +++ b/src/pipeline/image_upscale.h @@ -0,0 +1,57 @@ +#ifndef __SD_PIPELINE_IMAGE_UPSCALE_H__ +#define __SD_PIPELINE_IMAGE_UPSCALE_H__ + +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "ggml.h" + +class ModelManager; +class SDBackendManager; +struct GGMLRunner; +struct ESRGAN; +struct MoSR; + +namespace sd::pipeline { + class ImageUpscaleModel { + friend std::unique_ptr load_image_upscale_model( + const std::shared_ptr&, + SDBackendManager&, + const std::string&, + size_t, + bool, + ggml_type); + friend Tensor upscale_image_tensor(ImageUpscaleModel&, const Tensor&, int, int); + + std::weak_ptr manager_; + std::unique_ptr esrgan_; + std::unique_ptr mosr_; + std::vector weights_; + uint64_t source_file_ = 0; + + explicit ImageUpscaleModel(const std::shared_ptr& manager); + + public: + ~ImageUpscaleModel(); + GGMLRunner* runner() const; + int scale() const; + }; + + std::unique_ptr load_image_upscale_model( + const std::shared_ptr& manager, + SDBackendManager& backends, + const std::string& path, + size_t max_vram_bytes = 0, + bool conv_direct = false, + ggml_type weight_type = GGML_TYPE_F16); + + Tensor upscale_image_tensor(ImageUpscaleModel& model, + const Tensor& input, + int n_threads, + int tile_size); +} + +#endif // __SD_PIPELINE_IMAGE_UPSCALE_H__ diff --git a/src/upscaler.cpp b/src/upscaler.cpp index c9bbbcb6..82aa1f46 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -2,8 +2,9 @@ #include "core/ggml_extend_backend.h" #include "core/util.h" #include "generation_info.h" +#include "model/upscaler/esrgan.hpp" +#include "model/upscaler/mosr.hpp" #include "model_loader.h" -#include "runtime/tiling.h" #include "stable-diffusion.h" #include @@ -19,188 +20,44 @@ UpscalerGGML::UpscalerGGML(int n_threads, direct(direct), tile_size(tile_size), backend_spec(std::move(backend_spec)), - params_backend_spec(std::move(params_backend_spec)) { -} + params_backend_spec(std::move(params_backend_spec)) {} -UpscalerGGML::~UpscalerGGML() { - // ModelManager holds raw ggml tensor pointers owned by the runner context. - model_manager.reset(); - esrgan_upscaler.reset(); - mosr_upscaler.reset(); -} +UpscalerGGML::~UpscalerGGML() = default; void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; - if (runner() != nullptr) { - runner()->set_max_graph_vram_bytes(max_vram_bytes); + if (model != nullptr) { + model->runner()->set_max_graph_vram_bytes(max_vram_bytes); } } -GGMLRunner* UpscalerGGML::runner() { - if (mosr_upscaler) { - return mosr_upscaler.get(); - } - return esrgan_upscaler.get(); -} - int UpscalerGGML::get_scale() { - if (mosr_upscaler) { - return mosr_upscaler->config.scale; - } - if (esrgan_upscaler) { - return esrgan_upscaler->config.scale; - } - return 1; + return model != nullptr ? model->scale() : 1; } -sd::Tensor UpscalerGGML::compute(const sd::Tensor& input_tensor) { - if (mosr_upscaler) { - return mosr_upscaler->compute(n_threads, input_tensor); - } - return esrgan_upscaler->compute(n_threads, input_tensor); -} - -bool UpscalerGGML::load_from_file(const std::string& esrgan_path, - int n_threads) { +bool UpscalerGGML::load_from_file(const std::string& path, int n_threads) { sd::ParallelScope tensor_scope(&tensor_executor); ggml_log_set(sd_ggml_log_callback, nullptr); - + model.reset(); std::string error; - if (!backend_manager.init(backend_spec.c_str(), - params_backend_spec.c_str(), - /*split_mode_spec=*/nullptr, - &error)) { + if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), nullptr, &error)) { LOG_ERROR("upscaler backend config failed: %s", error.c_str()); return false; } - auto backend_for = [&](SDBackendModule module) { - ggml_backend_t module_backend = backend_manager.runtime_backend(module); - if (module_backend == nullptr) { - LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module)); - } - return module_backend; - }; - auto params_backend_for = [&](SDBackendModule module) { - ggml_backend_t module_backend = backend_manager.params_backend(module); - if (module_backend == nullptr) { - LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module)); - } - return module_backend; - }; - auto ensure_backend_pair = [&](SDBackendModule module) { - if (backend_for(module) == nullptr) { - return false; - } - return params_backend_for(module) != nullptr; - }; - if (!ensure_backend_pair(SDBackendModule::UPSCALER)) { - return false; - } - model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(false); - - ModelLoader model_loader; - if (!model_loader.init_from_file_and_convert_name(esrgan_path, "", VERSION_ESRGAN)) { - LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str()); - return false; - } - model_loader.set_wtype_override(model_data_type); - LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); - const auto& tensor_storage_map = model_loader.get_tensor_storage_map(); - std::map tensors; - if (tensor_storage_map.count("gblocks.0.weight") > 0) { - mosr_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - tensor_storage_map, - model_manager); - if (mosr_upscaler->mosr_net == nullptr) { - LOG_ERROR("init mosr model from metadata failed: '%s'", esrgan_path.c_str()); - return false; - } - mosr_upscaler->get_param_tensors(tensors); - } else { - esrgan_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - tensor_storage_map, - model_manager); - if (esrgan_upscaler->rrdb_net == nullptr) { - LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); - return false; - } - esrgan_upscaler->get_param_tensors(tensors); - } - LOG_INFO("Upscaler architecture: %s", runner()->get_desc().c_str()); - runner()->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (direct) { - runner()->set_conv2d_direct_enabled(true); - } - - if (!model_manager->set_loader(model_loader) || - !model_manager->register_param_tensors(ModelComponent::Upscaler, - std::move(tensors), - backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, - backend_for(SDBackendModule::UPSCALER), - params_backend_for(SDBackendModule::UPSCALER)) || - !model_manager->validate_registered_tensors()) { - LOG_ERROR("register upscaler tensors with model manager failed"); - return false; - } - return true; + model = sd::pipeline::load_image_upscale_model(model_manager, backend_manager, path, + max_graph_vram_bytes, direct, model_data_type); + return model != nullptr; } sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { sd::ParallelScope tensor_scope(&tensor_executor); - if (input_tensor.empty() || input_tensor.dim() != 4 || - (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) { - LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor"); + if (model == nullptr) { return {}; } - - const bool has_alpha = input_tensor.shape()[2] == 4; - sd::Tensor rgb; - if (has_alpha) { - rgb = sd::ops::slice(input_tensor, 2, 0, 3); - } - const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; - sd::Tensor upscaled; - const int scale = get_scale(); - if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { - upscaled = compute(model_input); - } else { - auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { - auto output_tile = compute(input_tile); - if (output_tile.empty()) { - LOG_ERROR("upscaler compute failed while processing a tile"); - return {}; - } - return output_tile; - }; - - upscaled = process_tiles_2d(model_input, - static_cast(input_tensor.shape()[0] * scale), - static_cast(input_tensor.shape()[1] * scale), - scale, - tile_size, - tile_size, - 0.25f, - false, - false, - on_processing); - } - runner()->runner_end(); - if (upscaled.empty()) { - LOG_ERROR("upscaler compute failed"); - return {}; - } - if (has_alpha) { - auto alpha = sd::ops::slice(input_tensor, 2, 3, 4); - auto alpha_shape = alpha.shape(); - alpha_shape[0] = upscaled.shape()[0]; - alpha_shape[1] = upscaled.shape()[1]; - alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear); - upscaled = sd::ops::concat(upscaled, alpha, 2); - } - return upscaled; + return sd::pipeline::upscale_image_tensor(*model, input_tensor, n_threads, tile_size); } sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) { diff --git a/src/upscaler.h b/src/upscaler.h index b93065fc..f3d09384 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -3,9 +3,8 @@ #include "core/ggml_extend_backend.h" #include "core/tensor.hpp" -#include "model/upscaler/esrgan.hpp" -#include "model/upscaler/mosr.hpp" #include "model_manager.h" +#include "pipeline/image_upscale.h" #include "stable-diffusion.h" #include @@ -15,9 +14,7 @@ struct UpscalerGGML { SDBackendManager backend_manager; std::shared_ptr model_manager; ggml_type model_data_type = GGML_TYPE_F16; - std::shared_ptr esrgan_upscaler; - std::shared_ptr mosr_upscaler; - std::string esrgan_path; + std::unique_ptr model; int n_threads; sd::ParallelExecutor tensor_executor; bool direct = false; @@ -36,9 +33,7 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); - GGMLRunner* runner(); int get_scale(); - sd::Tensor compute(const sd::Tensor& input_tensor); sd::Tensor upscale_tensor(const sd::Tensor& input_tensor); sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor); };