Commit ddc367a for stable-diffusion.cpp
commit ddc367ab5d502113e0df3b5d9415112ce89fac24
Author: leejet <leejet714@gmail.com>
Date: Tue Nov 10 23:45:01 2026 +0800
refactor: unify runner execution and weight residency (#2321)
diff --git a/docs/backend.md b/docs/backend.md
index da90e48..cc6545d 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 0000000..5ad241a
--- /dev/null
+++ b/src/core/execution_engine.cpp
@@ -0,0 +1,992 @@
+#include "core/execution_engine.h"
+
+#include <algorithm>
+#include <chrono>
+#include <exception>
+#include <thread>
+#include <unordered_set>
+
+#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<sd::ExecutionEngine> DeviceResidencyManager::execution_engine() {
+ if (execution_engine_ == nullptr) {
+ execution_engine_ = std::make_shared<sd::ExecutionEngine>(this);
+ }
+ return execution_engine_;
+}
+
+namespace sd {
+ class GraphExecution {
+ ExecutionEngine& engine_;
+ ExecutionSession& session_;
+ const ComputeTask& task_;
+ std::shared_ptr<DeviceResidencyManager> manager_;
+ std::unordered_set<const ggml_tensor*> params_tensor_set_;
+ ggml_status last_compute_status_ = GGML_STATUS_FAILED;
+ Tensor<float> output_;
+
+ using GraphCutPlan = ggml_graph_cut::Plan;
+ ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
+ std::vector<ggml_tensor*> 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<float> 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<const ggml_tensor*>* external_outputs = nullptr);
+ std::vector<DeviceMemoryRequest> memory_requests(const std::vector<BackendBufferSize>& sizes, size_t pending_cache_bytes) const;
+ bool fits(const std::vector<DeviceMemoryRequest>& requests, const std::vector<ggml_tensor*>& params) const;
+ bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
+ bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& 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<ggml_tensor*>(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<ggml_tensor*> GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) {
+ std::vector<ggml_tensor*> used_params;
+ if (gf == nullptr || params_tensor_set_.empty()) {
+ return used_params;
+ }
+
+ std::unordered_set<const ggml_tensor*> seen_params;
+ const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
+ seen_params.reserve(static_cast<size_t>(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<float> 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 : "<unnamed>",
+ 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<float>(tensor);
+ }
+
+ void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) {
+ GGML_ASSERT(gf != nullptr);
+ std::unordered_set<const ggml_tensor*> 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<size_t>(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 : "<unnamed>",
+ tensor ? ggml_op_name(tensor->op) : "<null>",
+ 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<ggml_tensor*> 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<ggml_backend_t> 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<const ggml_tensor*>* 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<DeviceMemoryRequest> GraphExecution::memory_requests(
+ const std::vector<BackendBufferSize>& sizes,
+ size_t pending_cache_bytes) const {
+ std::vector<DeviceMemoryRequest> 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<size_t>(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<DeviceMemoryRequest>& requests,
+ const std::vector<ggml_tensor*>& 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<void()>& 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<void()>& 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<BeforeComputeCallback*>(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<const ggml_tensor*> 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<size_t>(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<float>(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<bool()>& 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<ggml_backend_t, size_t> 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<const ggml_tensor*> 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<DeviceResidencyManager>& 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<size_t>& limits) {
+ backend_limits_ = limits;
+ weight_assignments_.clear();
+ node_assignments_.clear();
+ primary_notice_logged_ = false;
+ }
+
+ void ExecutionSession::set_runtime_backends(const std::vector<ggml_backend_t>& 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<DeviceResidencyManager>& manager) {
+ GGML_ASSERT(manager.get() == manager_);
+ auto session = std::make_unique<ExecutionSession>(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<ggml_tensor*> 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<ggml_tensor*>& required_params,
+ const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+ const std::vector<ggml_tensor*>& 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<const ggml_tensor*, size_t> 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 0000000..79ced1d
--- /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 <cstdint>
+#include <functional>
+#include <map>
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#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<ggml_tensor*> weights;
+ std::map<ggml_tensor*, const void*> inputs;
+ std::vector<std::pair<ggml_tensor*, std::string>> debug_tensors;
+ int n_threads = 1;
+ bool no_return = false;
+ std::function<bool()> read_outputs;
+ };
+
+ struct ComputeResult {
+ ggml_status status = GGML_STATUS_FAILED;
+ Tensor<float> output;
+ };
+
+ class ExecutionSession {
+ friend class ExecutionEngine;
+ friend class GraphExecution;
+
+ ggml_backend_t backend_;
+ std::weak_ptr<DeviceResidencyManager> residency_manager_;
+ RunnerCache cache_;
+ GraphCutTensorCache cut_cache_;
+ ComputeWorkspace workspace_;
+ std::unordered_set<ggml_tensor*> params_;
+ std::vector<ggml_backend_t> extra_backends_;
+ size_t max_backend_bytes_ = 0;
+ bool layer_split_enabled_ = false;
+ std::vector<size_t> backend_limits_;
+ ggml_graph_cut::PlanCache plan_cache_;
+ std::unordered_map<const ggml_tensor*, ggml_backend_t> weight_assignments_;
+ std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments_;
+ bool primary_notice_logged_ = false;
+ std::map<ggml_backend_t, size_t> 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<uintptr_t>(this); }
+
+ public:
+ ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& 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<size_t>& limits);
+ void set_runtime_backends(const std::vector<ggml_backend_t>& backends);
+ };
+
+ class ExecutionEngine {
+ friend class GraphExecution;
+
+ std::map<uintptr_t, std::unique_ptr<ExecutionSession>> 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<DeviceResidencyManager>& 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<ggml_tensor*>& required_params,
+ const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+ const std::vector<ggml_tensor*>& 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 39d5429..7a27811 100644
--- a/src/core/ggml_runner.cpp
+++ b/src/core/ggml_runner.cpp
@@ -1,15 +1,11 @@
#include <algorithm>
#include <exception>
-#include <map>
#include <utility>
#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<ggml_tensor*> GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf
return used_params;
}
-void GGMLRunner::evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& 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<uintptr_t>(this),
- runtime_backend, retained_runtime_buffer_bytes());
- for (auto backend : extra_runtime_backends) {
- manager->update_runtime_residency(reinterpret_cast<uintptr_t>(this),
- backend, retained_runtime_buffer_bytes(backend));
- }
- }
-}
-
-std::optional<sd::Tensor<float>> 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 : "<unnamed>",
- 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<float>(tensor);
-}
-
-void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) {
- GGML_ASSERT(gf != nullptr);
- std::unordered_set<const ggml_tensor*> 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<size_t>(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 : "<unnamed>",
- tensor ? ggml_op_name(tensor->op) : "<null>",
- 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<ggml_tensor*> 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<ggml_backend_t> 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<uintptr_t>(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<uintptr_t>(this));
- std::vector<ggml_tensor*> tensors;
- for (auto tensor : params_tensor_set_) {
- auto* parameter = manager->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
- if (parameter != nullptr)
- tensors.push_back(parameter);
- }
- manager->evict_compute_backend_params(tensors);
- manager->remove_runtime_owner(reinterpret_cast<uintptr_t>(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<DeviceResidencyManager> manager)
- : runtime_backend(backend),
- cache_(backend),
- cut_cache_(backend),
- workspace_(backend),
+ : execution_engine_(manager != nullptr ? manager->execution_engine() : std::make_shared<sd::ExecutionEngine>()),
+ 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<sd::Tensor<float>> 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<sd::Tensor<float>> 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;
- }
- success = output.has_value();
- if (success) {
- cache_.graph_end(true);
- last_compute_status_ = GGML_STATUS_SUCCESS;
+ LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what());
}
- 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<size_t>& 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<ggml_backend_t>& 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<const ggml_tensor*>* 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<DeviceMemoryRequest> GGMLRunner::memory_requests(
- const std::vector<BackendBufferSize>& sizes,
- size_t pending_cache_bytes) const {
- std::vector<DeviceMemoryRequest> 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<size_t>(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<uintptr_t>(this), pending,
- retained, limit});
- }
- return requests;
-}
-
-bool GGMLRunner::fits(const std::vector<DeviceMemoryRequest>& requests,
- const std::vector<ggml_tensor*>& 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<void()>& 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<void()>& 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<BeforeComputeCallback*>(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<const ggml_tensor*> 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<size_t>(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<float>(tensor);
- print_sd_tensor(debug_tensor, false, entry.second.c_str());
- }
- }
-
- return true;
-}
-
-std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& 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<uintptr_t>(this),
- graph, plan, params_tensor_set_,
- segmented && manager != nullptr && manager->prefetch_enabled());
-
- std::map<ggml_backend_t, size_t> 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<Tensor<float>> output = Tensor<float>();
- 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<const ggml_tensor*> 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 d300be7..9524824 100644
--- a/src/core/ggml_runner.h
+++ b/src/core/ggml_runner.h
@@ -12,9 +12,7 @@
#include <utility>
#include <vector>
-#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<ggml_backend_t, size_t> logged_compute_bytes_;
- size_t logged_segment_count_ = 0;
+ std::shared_ptr<sd::ExecutionEngine> 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<const ggml_tensor*>* external_outputs = nullptr);
-
- std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
- size_t pending_cache_bytes) const;
-
- bool fits(const std::vector<DeviceMemoryRequest>& requests,
- const std::vector<ggml_tensor*>& params) const;
-
- bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
-
- std::optional<sd::Tensor<float>> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
-
protected:
typedef std::function<ggml_cgraph*()> 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<std::max_align_t[]> 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<size_t> graph_cut_layer_split_backend_vram_limits_;
-
- std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
+ bool graph_active_ = false;
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
std::weak_ptr<DeviceResidencyManager> 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<const ggml_tensor*> params_tensor_set_;
- std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
- std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
- bool graph_cut_layer_split_primary_notice_logged_ = false;
template <typename T>
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -239,8 +210,6 @@ protected:
std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf);
- void evict_compute_backend_param_tensors(const std::vector<ggml_tensor*>& 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<sd::Tensor<float>> 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<sd::Tensor<float>> 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 faf7be2..7317410 100644
--- a/src/device_residency_manager.h
+++ b/src/device_residency_manager.h
@@ -3,12 +3,17 @@
#include <cstdint>
#include <functional>
+#include <memory>
#include <vector>
#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<ggml_tensor*> tensors;
+ size_t bytes = 0;
+ uint64_t last_use_epoch = 0;
+ bool requires_demotion = false;
+};
+
struct DeviceResidencyManager {
+private:
+ std::shared_ptr<sd::ExecutionEngine> 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<bool()> reclaim) = 0;
- virtual void remove_runtime_owner(uintptr_t owner_id) = 0;
+ std::shared_ptr<sd::ExecutionEngine> 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<ggml_tensor*>& required_params) const = 0;
@@ -62,6 +99,18 @@ struct DeviceResidencyManager {
const std::vector<ggml_tensor*>& required_params,
const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
const std::vector<ggml_tensor*>& protected_params) = 0;
+ virtual DeviceMemoryCapacity inspect_compute_backend_capacity(
+ const DeviceMemoryRequest& request,
+ const std::vector<ggml_tensor*>& params,
+ bool log_details = false) const {
+ DeviceMemoryCapacity result;
+ result.valid = false;
+ return result;
+ }
+ virtual std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
+ ggml_backend_t backend,
+ const std::vector<ggml_tensor*>& 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<ggml_tensor*>& tensors) = 0;
diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp
index 7cd86fd..00b88f4 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 8b61827..bfdf262 100644
--- a/src/model_manager.cpp
+++ b/src/model_manager.cpp
@@ -1,11 +1,9 @@
#include "model_manager.h"
#include <algorithm>
-#include <chrono>
#include <cstdint>
#include <iterator>
#include <mutex>
-#include <thread>
#include <tuple>
#include <unordered_set>
@@ -150,7 +148,7 @@ bool ModelManager::set_loras(std::vector<LoraSpec> 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<LoraSpec> loras, SDVersion version) {
std::set<std::string> ModelManager::tensor_names() const {
std::set<std::string> 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<ggml_tensor*> new_tensors;
std::vector<std::unique_ptr<TensorState>> new_states;
@@ -424,6 +422,20 @@ bool ModelManager::validate_registered_tensors() {
return ok;
}
+bool ModelManager::validate_param_tensors(const std::vector<ggml_tensor*>& tensors) {
+ std::vector<TensorState*> 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<TensorState*>& states) {
std::vector<TensorState*> need_load;
need_load.reserve(states.size());
@@ -500,7 +512,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
}
bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
- std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> states_by_staging_target;
+ std::map<std::tuple<ggml_backend_t, ggml_backend_buffer_type_t, ModelLoader::FileId>, std::vector<TensorState*>> 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::vector<TensorStat
if (staging_buft == nullptr) {
staging_buft = ggml_backend_get_default_buffer_type(state->compute_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<TensorState*>& target_states = pair.second;
if (target_states.empty()) {
continue;
@@ -646,7 +658,7 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
std::map<ApplyTarget, LoraApplyGroup> 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<TensorState*>& states,
std::vector<ParamsStorageBlock*>& created_storage_blocks) {
- std::map<std::pair<ggml_backend_buffer_type_t, int>, std::vector<TensorState*>> states_by_buffer_type;
+ std::map<std::tuple<ggml_backend_buffer_type_t, int, ModelLoader::FileId>, std::vector<TensorState*>> 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<TensorState*>& states,
if (params_buft == nullptr) {
return false;
}
- states_by_buffer_type[{params_buft, static_cast<int>(state->residency_mode)}].push_back(state);
+ states_by_buffer_type[{params_buft, static_cast<int>(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<TensorState*>& 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<TensorState*>& 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<ggml_tensor*>&
if (tensors.empty()) {
return;
}
- std::vector<TensorState*> required_states;
- if (!resolve_required_tensor_states(tensors, required_states)) {
+ std::unordered_set<TensorState*> 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<TensorState*> 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<bool()> 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,200 +1872,154 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
return result;
}
-bool ModelManager::fits_compute_backend_capacity(
+DeviceMemoryCapacity ModelManager::inspect_compute_backend_capacity(
const DeviceMemoryRequest& request,
- const std::vector<ggml_tensor*>& required_params) const {
+ const std::vector<ggml_tensor*>& params,
+ bool log_details) const {
std::vector<TensorState*> states;
- if (!resolve_required_tensor_states(required_params, states, request.compute_backend)) {
- return false;
+ if (request.compute_backend == nullptr || !resolve_required_tensor_states(params, states, request.compute_backend)) {
+ DeviceMemoryCapacity result;
+ result.valid = false;
+ return result;
}
- const auto capacity = check_capacity(request, states);
- return !capacity.requires_reclamation && capacity.fits();
+ return check_capacity(request, states, log_details);
}
-bool ModelManager::ensure_compute_backend_capacity(
- const DeviceMemoryRequest& request,
- const std::vector<ggml_tensor*>& required_params,
- const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
- const std::vector<ggml_tensor*>& protected_params) {
- std::vector<TensorState*> 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;
+std::vector<WeightResidencyBlock> ModelManager::evictable_compute_backend_blocks(
+ ggml_backend_t backend,
+ const std::vector<ggml_tensor*>& protected_params) const {
+ std::vector<WeightResidencyBlock> result;
+ if (backend == nullptr) {
+ return result;
}
- if (sd_backend_is_cpu(compute_backend)) {
- return true;
+ std::vector<TensorState*> protected_list;
+ if (!resolve_required_tensor_states(protected_params, protected_list)) {
+ return result;
}
-
- auto fits = [&]() {
- const auto capacity = check_capacity(request, required_states);
- return !capacity.requires_reclamation && capacity.fits();
- };
- if (fits()) {
- return true;
+ std::unordered_set<TensorState*> 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());
}
- for (const auto& entry : workspace_reclaimers_) {
- if (entry.first != request.owner_id) {
- entry.second();
- if (fits()) {
- return true;
- }
+ 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;
}
- }
-
- std::unordered_set<TensorState*> protected_states;
- std::vector<TensorState*> 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());
+ if (cpu) {
+ return ggml_backend_buffer_is_host(buffer);
}
- }
-
- std::unordered_set<TensorState*> 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;
+ if (ggml_backend_buffer_is_host(buffer)) {
+ return false;
}
- 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);
+ 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;
};
- 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<TensorState*> candidate_states;
- if (!resolve_required_tensor_states(candidate_params, candidate_states)) {
- return false;
+ auto append = [&](ggml_backend_buffer_t buffer, const std::vector<TensorState*>& states, bool storage) {
+ if (!on_device(buffer) || states.empty()) {
+ return;
}
- for (TensorState* state : candidate_states) {
- add_evictable_state(state);
+ 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);
}
- if (release_eviction_states()) {
- return true;
+ result.push_back(std::move(candidate));
+ };
+ for (const auto& block : compute_staging_blocks_) {
+ std::vector<TensorState*> states;
+ states.reserve(block->staged_tensors.size());
+ for (const auto& entry : block->staged_tensors) {
+ states.push_back(entry.first);
}
+ append(block->buffer, states, false);
}
-
- std::vector<TensorState*> 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;
+ 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;
+}
- // The loader retains mmap buffers after storage blocks are released, so
- // demoting mapped states would only force duplicate allocations on reload.
- std::unordered_set<TensorState*> mapped_states;
- for (const auto& block : params_storage_blocks_) {
- if (!block->mmap_tensor_stores.empty()) {
- mapped_states.insert(block->states.begin(), block->states.end());
+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;
}
}
- std::unordered_set<TensorState*> required_set(required_states.begin(), required_states.end());
- std::vector<TensorState*> demoted_states;
- auto revert_demotions = [&]() {
- for (TensorState* state : demoted_states) {
- if (state->loaded_to_params_backend) {
- state->residency_mode = ResidencyMode::ParamBackend;
+ 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;
+ }
+ }
}
- }
- };
- 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();
+ free_params_storage_block(**it);
+ params_storage_blocks_.erase(it);
return true;
}
}
- revert_demotions();
+ return false;
+}
- auto capacity = check_capacity(request, required_states);
- if (!capacity.requires_reclamation && capacity.fits()) {
- return true;
+bool ModelManager::fits_compute_backend_capacity(
+ const DeviceMemoryRequest& request,
+ const std::vector<ggml_tensor*>& required_params) const {
+ std::vector<TensorState*> states;
+ if (!resolve_required_tensor_states(required_params, states, request.compute_backend)) {
+ return false;
}
- 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;
+ const auto capacity = check_capacity(request, states);
+ return !capacity.requires_reclamation && capacity.fits();
+}
+
+bool ModelManager::ensure_compute_backend_capacity(
+ const DeviceMemoryRequest& request,
+ const std::vector<ggml_tensor*>& required_params,
+ const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
+ const std::vector<ggml_tensor*>& protected_params) {
+ 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 81d2546..f592a63 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<ModelLoader::FileId, ModelLoader> model_sources_;
std::vector<std::unique_ptr<TensorState>> tensor_states_;
std::map<const ggml_tensor*, TensorState*> tensor_states_by_tensor_;
mutable std::list<ResolvedTensorStates> resolved_tensor_states_;
@@ -101,7 +103,7 @@ private:
std::map<uintptr_t, std::unique_ptr<PrefetchBlock>> prefetch_blocks_;
std::map<ggml_backend_t, ggml_backend_t> prefetch_backends_;
std::map<std::pair<uintptr_t, ggml_backend_t>, RuntimeResidency> runtime_residencies_;
- std::map<uintptr_t, std::function<bool()>> workspace_reclaimers_;
+ std::unordered_set<uintptr_t> runtime_owners_;
bool warned_split_lora_skip_ = false;
std::set<std::string> common_ignore_tensors_;
std::vector<LoraSpec> loras_;
@@ -131,6 +133,7 @@ private:
std::vector<TensorState*>& 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<TensorState*>& 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<TensorState*>& 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<ModelComponent>& components, const ModelLoader& loader) const;
size_t registered_params_size(const std::set<ModelComponent>& components) const;
@@ -267,12 +262,13 @@ public:
}
bool validate_registered_tensors();
+ bool validate_param_tensors(const std::vector<ggml_tensor*>& tensors);
bool load_all_params_eagerly();
bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
ggml_backend_t compute_backend) override;
bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
- void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> 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<ggml_tensor*>& required_params) const override;
@@ -290,6 +286,14 @@ public:
const std::vector<ggml_tensor*>& required_params,
const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
const std::vector<ggml_tensor*>& protected_params) override;
+ DeviceMemoryCapacity inspect_compute_backend_capacity(
+ const DeviceMemoryRequest& request,
+ const std::vector<ggml_tensor*>& params,
+ bool log_details = false) const override;
+ std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
+ ggml_backend_t backend,
+ const std::vector<ggml_tensor*>& protected_params) const override;
+ bool evict_compute_backend_block(const WeightResidencyBlock& block) override;
WeightPrefetchResult prefetch_params(
uintptr_t owner_id,
const std::vector<ggml_tensor*>& tensors) override;
diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp
index adeea95..d6453a2 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<TensorState*>& 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::set<ModelComp
continue;
}
if (state->source_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<ModelComponent>& 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 1f7d5e4..50a4fb7 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<float> upscale_hires_latent(StableDiffusionGGML* sd,
const sd::Tensor<float>& 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<float> 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<UpscalerGGML> hires_upscaler;
+ std::unique_ptr<ImageUpscaleModel> 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<UpscalerGGML>(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 0000000..1ccb7b5
--- /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<ModelManager>& 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<GGMLRunner*>(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<ImageUpscaleModel> load_image_upscale_model(
+ const std::shared_ptr<ModelManager>& 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<ImageUpscaleModel>(new ImageUpscaleModel(manager));
+ model->source_file_ = manager->register_model_source(std::move(loader));
+ if (model->source_file_ == 0) {
+ return nullptr;
+ }
+ std::map<std::string, ggml_tensor*> weights;
+ if (tensors.count("gblocks.0.weight") != 0) {
+ model->mosr_ = std::make_unique<MoSR>(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<ESRGAN>(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<float> upscale_image_tensor(ImageUpscaleModel& model, const Tensor<float>& 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<float>& tensor) {
+ return model.mosr_ != nullptr ? model.mosr_->compute(n_threads, tensor) : model.esrgan_->compute(n_threads, tensor);
+ };
+ const int scale = model.scale();
+ Tensor<float> 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<int>(input.shape()[0] * scale),
+ static_cast<int>(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 0000000..1de9ca0
--- /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 <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#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<ImageUpscaleModel> load_image_upscale_model(
+ const std::shared_ptr<ModelManager>&,
+ SDBackendManager&,
+ const std::string&,
+ size_t,
+ bool,
+ ggml_type);
+ friend Tensor<float> upscale_image_tensor(ImageUpscaleModel&, const Tensor<float>&, int, int);
+
+ std::weak_ptr<ModelManager> manager_;
+ std::unique_ptr<ESRGAN> esrgan_;
+ std::unique_ptr<MoSR> mosr_;
+ std::vector<ggml_tensor*> weights_;
+ uint64_t source_file_ = 0;
+
+ explicit ImageUpscaleModel(const std::shared_ptr<ModelManager>& manager);
+
+ public:
+ ~ImageUpscaleModel();
+ GGMLRunner* runner() const;
+ int scale() const;
+ };
+
+ std::unique_ptr<ImageUpscaleModel> load_image_upscale_model(
+ const std::shared_ptr<ModelManager>& 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<float> upscale_image_tensor(ImageUpscaleModel& model,
+ const Tensor<float>& input,
+ int n_threads,
+ int tile_size);
+}
+
+#endif // __SD_PIPELINE_IMAGE_UPSCALE_H__
diff --git a/src/upscaler.cpp b/src/upscaler.cpp
index c9bbbcb..82aa1f4 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 <cstdlib>
@@ -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);
- }
-}
-
-GGMLRunner* UpscalerGGML::runner() {
- if (mosr_upscaler) {
- return mosr_upscaler.get();
+ if (model != nullptr) {
+ model->runner()->set_max_graph_vram_bytes(max_vram_bytes);
}
- 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;
-}
-
-sd::Tensor<float> UpscalerGGML::compute(const sd::Tensor<float>& input_tensor) {
- if (mosr_upscaler) {
- return mosr_upscaler->compute(n_threads, input_tensor);
- }
- return esrgan_upscaler->compute(n_threads, input_tensor);
+ return model != nullptr ? model->scale() : 1;
}
-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<ModelManager>();
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<std::string, ggml_tensor*> tensors;
- if (tensor_storage_map.count("gblocks.0.weight") > 0) {
- mosr_upscaler = std::make_shared<MoSR>(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<ESRGAN>(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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& 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<float> rgb;
- if (has_alpha) {
- rgb = sd::ops::slice(input_tensor, 2, 0, 3);
- }
- const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
- sd::Tensor<float> 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<float>& input_tile) -> sd::Tensor<float> {
- 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<int>(input_tensor.shape()[0] * scale),
- static_cast<int>(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 b93065f..f3d0938 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 <memory>
@@ -15,9 +14,7 @@ struct UpscalerGGML {
SDBackendManager backend_manager;
std::shared_ptr<ModelManager> model_manager;
ggml_type model_data_type = GGML_TYPE_F16;
- std::shared_ptr<ESRGAN> esrgan_upscaler;
- std::shared_ptr<MoSR> mosr_upscaler;
- std::string esrgan_path;
+ std::unique_ptr<sd::pipeline::ImageUpscaleModel> 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<float> compute(const sd::Tensor<float>& input_tensor);
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
};