DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 03:24:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: unify runner execution and weight residency (#2321)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-10 23:45:01 +08:00
1 parent 56070ea034
commit ddc367ab5d
15 files changed
+1679 -1253
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/backend.md
+18 -6
View File
@@ -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.
src/core/execution_engine.cpp
+992
View File
@@ -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;
}
}
src/core/execution_engine.h
+107
View File
@@ -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__
src/core/ggml_runner.cpp
+38 -780
View File
@@ -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;
LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what());
}
success = output.has_value();
if (success) {
cache_.graph_end(true);
last_compute_status_ = GGML_STATUS_SUCCESS;
}
return output;
return std::nullopt;
}
void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) {
graph_cut_layer_split_enabled = enabled;
if (!enabled) {
graph_cut_layer_split_assignments_.clear();
graph_cut_layer_split_node_assignments_.clear();
graph_cut_layer_split_primary_notice_logged_ = false;
}
execution_session_.set_layer_split_enabled(enabled);
}
void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector<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);
}
src/core/ggml_runner.h
+12 -64
View File
@@ -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);
src/device_residency_manager.h
+53 -4
View File
@@ -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;
src/model/diffusion/qwen_image_2_1.hpp
+1 -1
View File
@@ -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;
src/model_manager.cpp
+181 -204
View File
@@ -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,6 +1872,139 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
return result;
}
DeviceMemoryCapacity ModelManager::inspect_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& params,
bool log_details) const {
std::vector<TensorState*> states;
if (request.compute_backend == nullptr || !resolve_required_tensor_states(params, states, request.compute_backend)) {
DeviceMemoryCapacity result;
result.valid = false;
return result;
}
return check_capacity(request, states, log_details);
}
std::vector<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;
}
std::vector<TensorState*> protected_list;
if (!resolve_required_tensor_states(protected_params, protected_list)) {
return result;
}
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());
}
const auto device = ggml_backend_get_device(backend);
const bool cpu = sd_backend_is_cpu(backend);
auto on_device = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr) {
return false;
}
if (cpu) {
return ggml_backend_buffer_is_host(buffer);
}
if (ggml_backend_buffer_is_host(buffer)) {
return false;
}
const auto type = ggml_backend_buffer_get_type(buffer);
const auto split = split_buffer_devices_.find(type);
if (split != split_buffer_devices_.end()) {
return std::any_of(split->second.begin(), split->second.end(), [&](const auto& entry) {
return ggml_backend_get_device(entry.first) == device;
});
}
return ggml_backend_buft_get_device(type) == device;
};
auto append = [&](ggml_backend_buffer_t buffer, const std::vector<TensorState*>& states, bool storage) {
if (!on_device(buffer) || states.empty()) {
return;
}
WeightResidencyBlock candidate;
candidate.buffer = buffer;
candidate.backend = backend;
candidate.bytes = ggml_backend_buffer_get_size(buffer);
for (auto* state : states) {
if (state == nullptr) {
continue;
}
if (state->pin_count != 0 || protected_states.count(state) != 0) {
return;
}
if (storage) {
if (state->staged_to_compute_backend) {
return;
}
if (state->residency_mode != ResidencyMode::Disk) {
if (cpu || state->params_backend == nullptr || sd_backend_is_cpu(state->params_backend) ||
state->compute_backend != state->params_backend ||
state->sources.empty() || state->applied_lora_epoch != UINT64_MAX ||
state->split_buffer_type != nullptr || state->tensor->view_src != nullptr) {
return;
}
candidate.requires_demotion = true;
}
}
candidate.tensors.push_back(state->tensor);
candidate.last_use_epoch = std::max(candidate.last_use_epoch, state->last_use_epoch);
}
result.push_back(std::move(candidate));
};
for (const auto& block : compute_staging_blocks_) {
std::vector<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);
}
for (const auto& block : params_storage_blocks_) {
// Loader-retained mappings cannot be reclaimed by freeing a storage block.
if (block->mmap_tensor_stores.empty()) {
append(block->buffer, block->states, true);
}
}
return result;
}
bool ModelManager::evict_compute_backend_block(const WeightResidencyBlock& block) {
const auto candidates = evictable_compute_backend_blocks(block.backend, {});
const auto found = std::find_if(candidates.begin(), candidates.end(), [&](const auto& candidate) {
return candidate.buffer == block.buffer && candidate.tensors == block.tensors &&
candidate.requires_demotion == block.requires_demotion;
});
if (found == candidates.end()) {
return false;
}
for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end(); ++it) {
if ((*it)->buffer == block.buffer) {
free_compute_staging_block(**it);
compute_staging_blocks_.erase(it);
return true;
}
}
for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end(); ++it) {
if ((*it)->buffer == block.buffer) {
if (block.requires_demotion) {
for (auto* state : (*it)->states) {
if (state != nullptr) {
state->residency_mode = ResidencyMode::Disk;
}
}
}
free_params_storage_block(**it);
params_storage_blocks_.erase(it);
return true;
}
}
return false;
}
bool ModelManager::fits_compute_backend_capacity(
const DeviceMemoryRequest& request,
const std::vector<ggml_tensor*>& required_params) const {
@@ -1865,184 +2021,5 @@ bool ModelManager::ensure_compute_backend_capacity(
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;
}
if (sd_backend_is_cpu(compute_backend)) {
return true;
}
auto fits = [&]() {
const auto capacity = check_capacity(request, required_states);
return !capacity.requires_reclamation && capacity.fits();
};
if (fits()) {
return true;
}
for (const auto& entry : workspace_reclaimers_) {
if (entry.first != request.owner_id) {
entry.second();
if (fits()) {
return true;
}
}
}
std::unordered_set<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());
}
}
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;
}
const bool reloadable = state->residency_mode == ResidencyMode::Disk ||
state->compute_backend != state->params_backend;
const bool resident = state->compute_backend == state->params_backend
? state->loaded_to_params_backend
: state->staged_to_compute_backend;
if (reloadable && resident) {
eviction_states.insert(state);
}
};
size_t last_eviction_count = 0;
auto release_eviction_states = [&]() {
if (eviction_states.size() != last_eviction_count) {
last_eviction_count = eviction_states.size();
release_compute_staging_blocks(false, &eviction_states);
release_params_storage_blocks(false, &eviction_states);
}
// Device memory reports can lag behind releases even when no new candidates are added.
return fits();
};
for (const auto& candidate_params : preferred_eviction_order) {
std::vector<TensorState*> candidate_states;
if (!resolve_required_tensor_states(candidate_params, candidate_states)) {
return false;
}
for (TensorState* state : candidate_states) {
add_evictable_state(state);
}
if (release_eviction_states()) {
return true;
}
}
std::vector<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;
}
}
// 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());
}
}
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;
}
}
};
auto can_demote_to_disk = [&](TensorState* state) {
return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr &&
state->residency_mode == ResidencyMode::ParamBackend &&
state->params_backend != nullptr && state->params_backend == state->compute_backend &&
state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) &&
state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend &&
protected_states.find(state) == protected_states.end() &&
required_set.find(state) == required_set.end() &&
mapped_states.find(state) == mapped_states.end() &&
!state->sources.empty() && state->applied_lora_epoch == UINT64_MAX &&
state->split_buffer_type == nullptr;
};
for (TensorState* state : global_candidates) {
if (!can_demote_to_disk(state)) {
continue;
}
state->residency_mode = ResidencyMode::Disk;
demoted_states.push_back(state);
add_evictable_state(state);
if (release_eviction_states()) {
revert_demotions();
return true;
}
}
revert_demotions();
auto capacity = check_capacity(request, required_states);
if (!capacity.requires_reclamation && capacity.fits()) {
return true;
}
if (sd_backend_is(compute_backend, "Vulkan") &&
capacity.required_budget_bytes <= capacity.available_budget_bytes) {
ggml_backend_synchronize(compute_backend);
sd_backend_clear_buffer_cache(compute_backend);
// Vulkan heap usage can lag behind buffer releases. Fast memory queries
// must not turn that reporting delay into an allocation failure.
const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
for (;;) {
capacity = check_capacity(request, required_states);
if (!capacity.requires_reclamation && capacity.fits()) {
return true;
}
if (capacity.required_budget_bytes > capacity.available_budget_bytes ||
std::chrono::steady_clock::now() >= deadline) {
break;
}
std::this_thread::sleep_for(std::chrono::milliseconds(10));
}
}
capacity = check_capacity(request, required_states, true);
const std::string available_device = capacity.available_device_bytes == SIZE_MAX
? "unknown"
: sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));
const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX
? "unlimited"
: sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0));
LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
ggml_backend_name(compute_backend),
capacity.required_device_bytes / (1024.0 * 1024.0),
capacity.required_budget_bytes / (1024.0 * 1024.0),
available_device.c_str(),
available_budget.c_str());
return false;
return execution_engine()->ensure_capacity(*this, request, required_params, preferred_eviction_order, protected_params);
}
src/model_manager.h
+17 -13
View File
@@ -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;
src/model_manager_files.cpp
+39 -3
View File
@@ -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;
}
src/pipeline/image.cpp
+9 -13
View File
@@ -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;
}
src/pipeline/image_upscale.cpp
+138
View File
@@ -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;
}
}
src/pipeline/image_upscale.h
+57
View File
@@ -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__
src/upscaler.cpp
+15 -158
View File
@@ -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);
if (model != nullptr) {
model->runner()->set_max_graph_vram_bytes(max_vram_bytes);
}
}
GGMLRunner* UpscalerGGML::runner() {
if (mosr_upscaler) {
return mosr_upscaler.get();
}
return esrgan_upscaler.get();
}
int UpscalerGGML::get_scale() {
if (mosr_upscaler) {
return mosr_upscaler->config.scale;
}
if (esrgan_upscaler) {
return esrgan_upscaler->config.scale;
}
return 1;
return model != nullptr ? model->scale() : 1;
}
sd::Tensor<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);
}
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) {
src/upscaler.h
+2 -7
View File
@@ -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);
};