mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-14 12:56:33 +00:00
refactor: unify runner execution and weight residency (#2321)
This commit is contained in:
leejet authored and GitHub committed
1 parent
56070ea034
commit
ddc367ab5d
15 files changed
+1679
-1253
No files matched your search
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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__
|
||||
+38
-780
@@ -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);
|
||||
}
|
||||
+12
-64
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
+181
-204
@@ -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);
|
||||
}
|
||||
+17
-13
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
+9
-13
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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__
|
||||
+15
-158
@@ -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) {
|
||||
|
||||
+2
-7
@@ -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);
|
||||
};
|
||||
|
||||
Reference in new issue
Block a user