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

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

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-09-06 22:30:45 +08:00
1 parent 6c57cc3b38
commit 462d675018
78 files changed
+2963 -2051
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/compute_workspace.cpp
+257
View File
@@ -0,0 +1,257 @@
#include "core/compute_workspace.h"
#include <algorithm>
#include <cstring>
#include <map>
#include <unordered_map>
#include <unordered_set>
#include "core/ggml_extend_backend.h"
#include "core/ggml_graph_cut.h"
#include "ggml-cpu.h"
#include "ggml/src/ggml-impl.h"
namespace sd {
ComputeWorkspace::~ComputeWorkspace() {
segment_end();
release();
ggml_backend_free(cpu_backend_);
}
void ComputeWorkspace::set_extra_backends(const std::vector<ggml_backend_t>& backends) {
if (extra_backends_ != backends) {
GGML_ASSERT(!active_);
release();
extra_backends_ = backends;
}
}
bool ComputeWorkspace::needs_scheduler(ggml_cgraph* graph) const {
if (!extra_backends_.empty()) {
return true;
}
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
if (!ggml_backend_supports_op(backend_, ggml_graph_node(graph, i))) {
return true;
}
}
return false;
}
ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) {
std::vector<ggml_backend_t> backends{backend_};
backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end());
if (!sd_backend_is_cpu(backend_)) {
if (cpu_backend_ == nullptr) {
cpu_backend_ = sd_backend_cpu_init();
}
if (cpu_backend_ == nullptr) {
return nullptr;
}
backends.push_back(cpu_backend_);
}
std::vector<ggml_backend_buffer_type_t> bufts;
for (auto backend : backends) {
auto buft = backend == cpu_backend_
? ggml_backend_dev_host_buffer_type(ggml_backend_get_device(backend_))
: nullptr;
bufts.push_back(buft != nullptr ? buft : ggml_backend_get_default_buffer_type(backend));
}
return ggml_backend_sched_new(backends.data(), bufts.data(), static_cast<int>(backends.size()),
graph_size, false, false);
}
bool ComputeWorkspace::measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const {
return measurement.scheduler == needs_scheduler(graph);
}
bool ComputeWorkspace::prepare(const Measurement& measurement) {
GGML_ASSERT(!active_);
if (measurement.buffers.empty()) {
return false;
}
const bool grows = std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
[&](const BackendBufferSize& size) { return size.bytes > bytes(size.backend); });
if (measurement.scheduler != (scheduler_ != nullptr) || grows) {
release();
}
return true;
}
bool ComputeWorkspace::release_excess(const Measurement& measurement) {
return std::any_of(measurement.buffers.begin(), measurement.buffers.end(),
[&](const BackendBufferSize& size) { return bytes(size.backend) > size.bytes; }) &&
release();
}
bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
GGML_ASSERT(!active_);
const bool use_scheduler = needs_scheduler(graph);
if (use_scheduler) {
if (allocator_ != nullptr) {
release();
}
const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
release();
scheduler_ = make_scheduler(capacity);
scheduler_capacity_ = capacity;
}
if (scheduler_ == nullptr) {
return false;
}
ggml_backend_sched_reset(scheduler_);
assign_nodes(scheduler_, graph);
// Scheduler allocation rewrites sources. Split the execution graph only once.
if (!ggml_backend_sched_alloc_graph(scheduler_, graph)) {
release();
return false;
}
} else {
if (scheduler_ != nullptr) {
release();
}
if (allocator_ == nullptr) {
allocator_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_));
}
auto signature = ggml_graph_cut::graph_layout(graph, true);
if (signature != reservation_) {
if (!ggml_gallocr_reserve(allocator_, graph)) {
release();
return false;
}
reservation_ = std::move(signature);
++reservations_;
}
if (!ggml_gallocr_alloc_graph(allocator_, graph)) {
release();
return false;
}
}
active_ = true;
return true;
}
ComputeWorkspace::Measurement ComputeWorkspace::measure(
ggml_cgraph* graph,
size_t direct_bytes,
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
const AssignNodes& assign_nodes) {
if (!needs_scheduler(graph)) {
return {{{backend_, direct_bytes}}, false};
}
std::vector<const ggml_tensor*> tensors;
std::unordered_set<const ggml_tensor*> seen;
auto visit = [&](const ggml_tensor* tensor) {
if (tensor != nullptr && seen.insert(tensor).second) {
tensors.push_back(tensor);
}
};
for (int i = 0; i < graph->n_nodes; ++i) {
visit(graph->nodes[i]);
}
for (int i = 0; i < graph->n_leafs; ++i) {
visit(graph->leafs[i]);
}
for (size_t i = 0; i < tensors.size(); ++i) {
visit(tensors[i]->view_src);
for (auto source : tensors[i]->src) {
visit(source);
}
}
const size_t graph_size = tensors.size() + 8;
auto context = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
if (context == nullptr) {
return {};
}
std::unordered_map<const ggml_tensor*, ggml_tensor*> copies;
std::map<ggml_backend_t, ggml_backend_buffer_t> external_buffers;
for (auto tensor : tensors) {
auto copy = ggml_dup_tensor(context, tensor);
*copy = *tensor;
copies[tensor] = copy;
}
for (const auto& entry : copies) {
auto source = entry.first;
auto copy = entry.second;
copy->view_src = source->view_src == nullptr ? nullptr : copies.at(source->view_src);
for (int i = 0; i < GGML_MAX_SRC; ++i) {
copy->src[i] = source->src[i] == nullptr ? nullptr : copies.at(source->src[i]);
}
auto external = external_backend(source);
if (external != nullptr && source->view_src == nullptr) {
auto& buffer = external_buffers[external];
if (buffer == nullptr) {
buffer = ggml_backend_alloc_buffer(external, 0);
GGML_ASSERT(buffer != nullptr);
ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
}
copy->buffer = buffer;
copy->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
copy->extra = nullptr;
}
}
auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
copy_graph->n_nodes = graph->n_nodes;
copy_graph->n_leafs = graph->n_leafs;
for (int i = 0; i < graph->n_nodes; ++i) {
copy_graph->nodes[i] = copies.at(graph->nodes[i]);
}
for (int i = 0; i < graph->n_leafs; ++i) {
copy_graph->leafs[i] = copies.at(graph->leafs[i]);
}
Measurement result;
result.scheduler = true;
auto scheduler = make_scheduler(graph_size);
if (scheduler != nullptr) {
assign_nodes(scheduler, copy_graph);
std::vector<size_t> sizes(extra_backends_.size() + 2);
ggml_backend_sched_reserve_size(scheduler, copy_graph, sizes.data());
result.buffers.push_back({backend_, sizes[0]});
for (size_t i = 0; i < extra_backends_.size(); ++i) {
result.buffers.push_back({extra_backends_[i], sizes[i + 1]});
}
ggml_backend_sched_free(scheduler);
}
for (const auto& entry : external_buffers) {
ggml_backend_buffer_free(entry.second);
}
ggml_free(context);
return result;
}
void ComputeWorkspace::synchronize() const {
if (scheduler_ != nullptr) {
ggml_backend_sched_synchronize(scheduler_);
} else {
ggml_backend_synchronize(backend_);
}
}
void ComputeWorkspace::segment_end() {
if (active_) {
synchronize();
active_ = false;
}
}
bool ComputeWorkspace::release() {
if (active_) {
return false;
}
ggml_gallocr_free(allocator_);
allocator_ = nullptr;
ggml_backend_sched_free(scheduler_);
scheduler_ = nullptr;
scheduler_capacity_ = 0;
reservation_.clear();
return true;
}
size_t ComputeWorkspace::bytes(ggml_backend_t backend) const {
if (scheduler_ != nullptr) {
return ggml_backend_sched_get_buffer_size(scheduler_, backend);
}
return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0;
}
}
src/core/compute_workspace.h
+64
View File
@@ -0,0 +1,64 @@
#ifndef __SD_CORE_COMPUTE_WORKSPACE_H__
#define __SD_CORE_COMPUTE_WORKSPACE_H__
#include <functional>
#include <vector>
#include "ggml-alloc.h"
#include "ggml-backend.h"
namespace sd {
struct BackendBufferSize {
ggml_backend_t backend = nullptr;
size_t bytes = 0;
};
class ComputeWorkspace {
ggml_backend_t backend_;
std::vector<ggml_backend_t> extra_backends_;
ggml_backend_t cpu_backend_ = nullptr;
ggml_gallocr_t allocator_ = nullptr;
ggml_backend_sched_t scheduler_ = nullptr;
size_t scheduler_capacity_ = 0;
std::vector<uint64_t> reservation_;
bool active_ = false;
size_t reservations_ = 0;
ggml_backend_sched_t make_scheduler(size_t graph_size);
bool needs_scheduler(ggml_cgraph* graph) const;
public:
struct Measurement {
std::vector<BackendBufferSize> buffers;
bool scheduler = false;
};
using AssignNodes = std::function<void(ggml_backend_sched_t, ggml_cgraph*)>;
explicit ComputeWorkspace(ggml_backend_t backend)
: backend_(backend) {}
~ComputeWorkspace();
ComputeWorkspace(const ComputeWorkspace&) = delete;
ComputeWorkspace& operator=(const ComputeWorkspace&) = delete;
void set_extra_backends(const std::vector<ggml_backend_t>& backends);
bool measurement_matches(ggml_cgraph* graph, const Measurement& measurement) const;
bool prepare(const Measurement& measurement);
bool release_excess(const Measurement& measurement);
bool allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes);
Measurement measure(
ggml_cgraph* graph,
size_t direct_bytes,
const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
const AssignNodes& assign_nodes);
void synchronize() const;
void segment_end();
bool release();
bool active() const { return active_; }
ggml_backend_sched_t scheduler() const { return scheduler_; }
ggml_backend_t cpu_backend() const { return cpu_backend_; }
size_t bytes(ggml_backend_t backend) const;
size_t reservation_count() const { return reservations_; }
};
}
#endif // __SD_CORE_COMPUTE_WORKSPACE_H__
src/core/ggml_extend.hpp
+150 -883
View File
File diff suppressed because it is too large. Load diff
src/core/ggml_graph_cut.cpp
+235 -339
View File
@@ -2,6 +2,7 @@
#include <algorithm>
#include <cctype>
#include <climits>
#include <cmath>
#include <cstring>
#include <map>
@@ -67,25 +68,6 @@ namespace sd::ggml_graph_cut {
return -1;
}
static Plan::InputShape input_shape(const ggml_tensor* tensor) {
Plan::InputShape shape;
if (tensor == nullptr) {
return shape;
}
shape.type = tensor->type;
for (int i = 0; i < GGML_MAX_DIMS; ++i) {
shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
}
return shape;
}
static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
return segment.compute_buffer_size +
segment.input_param_bytes +
segment.input_previous_cut_bytes +
segment.output_bytes;
}
static std::string lower_ascii_copy(std::string value) {
std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
return static_cast<char>(std::tolower(c));
@@ -291,55 +273,6 @@ namespace sd::ggml_graph_cut {
return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
}
static bool is_segment_output_needed_after(const Plan& plan,
size_t end_segment_index,
int output_node_index) {
if (end_segment_index + 1 >= plan.segments.size()) {
return false;
}
for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
const auto& segment = plan.segments[seg_idx];
for (const auto& input_ref : segment.input_refs) {
if (input_ref.type == Segment::INPUT_PREVIOUS_CUT &&
input_ref.node_index == output_node_index) {
return true;
}
}
}
return false;
}
static Segment make_segment_seed(const Plan& plan,
size_t start_segment_index,
size_t end_segment_index) {
GGML_ASSERT(start_segment_index < plan.segments.size());
GGML_ASSERT(end_segment_index < plan.segments.size());
GGML_ASSERT(start_segment_index <= end_segment_index);
Segment seed;
const auto& start_segment = plan.segments[start_segment_index];
const auto& target_segment = plan.segments[end_segment_index];
std::unordered_set<int> seen_output_node_indices;
for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
const bool is_boundary_segment = seg_idx == end_segment_index;
for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
if ((is_boundary_segment ||
is_segment_output_needed_after(plan, end_segment_index, output_node_index)) &&
seen_output_node_indices.insert(output_node_index).second) {
seed.output_node_indices.push_back(output_node_index);
}
}
}
if (start_segment_index == end_segment_index) {
seed.group_name = target_segment.group_name;
} else {
seed.group_name = sd_format("%s..%s",
start_segment.group_name.c_str(),
target_segment.group_name.c_str());
}
return seed;
}
static void build_segment(ggml_cgraph* gf,
Plan& plan,
Segment& segment,
@@ -416,31 +349,7 @@ namespace sd::ggml_graph_cut {
}
return a.display_name < b.display_name;
});
segment.input_refs = input_refs;
for (const auto& input : input_refs) {
ggml_tensor* current_input = input_tensor(gf, input);
size_t tensor_bytes = current_input == nullptr
? 0
: (input.type == Segment::INPUT_PREVIOUS_CUT
? cache_tensor_bytes(current_input)
: ggml_nbytes(current_input));
switch (input.type) {
case Segment::INPUT_PREVIOUS_CUT:
segment.input_previous_cut_bytes += tensor_bytes;
break;
case Segment::INPUT_PARAM:
segment.input_param_bytes += tensor_bytes;
break;
case Segment::INPUT_EXTERNAL:
default:
segment.input_external_bytes += tensor_bytes;
break;
}
}
for (int output_node_index : segment.output_node_indices) {
ggml_tensor* output = ggml_graph_node(gf, output_node_index);
segment.output_bytes += cache_tensor_bytes(output);
}
segment.input_refs = input_refs;
segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
for (int output_node_index : segment.output_node_indices) {
@@ -449,6 +358,70 @@ namespace sd::ggml_graph_cut {
plan.segments.push_back(std::move(segment));
}
static bool validate_plan(ggml_cgraph* gf,
const Plan& plan,
std::string* validation_error) {
auto fail = [&](const std::string& reason) {
if (validation_error != nullptr) {
*validation_error = reason;
}
return false;
};
if (!plan.has_cuts) {
return true;
}
if (plan.segments.size() <= 1) {
return fail("fewer than two segments");
}
const int n_nodes = ggml_graph_n_nodes(gf);
std::unordered_set<int> completed_outputs;
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
const Segment& segment = plan.segments[segment_index];
const std::string segment_label = "segment " + std::to_string(segment_index) +
" ('" + segment.group_name + "')";
if (segment.internal_node_indices.empty() || segment.output_node_indices.empty()) {
return fail(segment_label + " has no internal nodes or outputs");
}
for (const Segment::InputRef& input : segment.input_refs) {
if (input.type == Segment::INPUT_PREVIOUS_CUT) {
if (input.node_index < 0 || input.node_index >= n_nodes ||
completed_outputs.find(input.node_index) == completed_outputs.end()) {
return fail(segment_label + " references an unavailable cut node " +
std::to_string(input.node_index));
}
} else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) {
return fail(segment_label + " references an invalid leaf " +
std::to_string(input.leaf_index));
}
}
std::unordered_set<int> segment_nodes;
segment_nodes.reserve(segment.internal_node_indices.size());
for (int node_index : segment.internal_node_indices) {
if (node_index < 0 || node_index >= n_nodes) {
return fail(segment_label + " contains an invalid node " +
std::to_string(node_index));
}
if (!segment_nodes.insert(node_index).second) {
return fail(segment_label + " contains duplicate node " +
std::to_string(node_index));
}
}
for (int output_index : segment.output_node_indices) {
if (output_index < 0 || output_index >= n_nodes ||
segment_nodes.find(output_index) == segment_nodes.end()) {
return fail(segment_label + " has an output outside its node set: " +
std::to_string(output_index));
}
if (completed_outputs.find(output_index) != completed_outputs.end()) {
return fail(segment_label + " repeats output node " +
std::to_string(output_index));
}
completed_outputs.insert(output_index);
}
}
return true;
}
bool is_graph_cut_tensor(const ggml_tensor* tensor) {
if (tensor == nullptr || tensor->name[0] == '\0') {
return false;
@@ -509,26 +482,80 @@ namespace sd::ggml_graph_cut {
return ggml_nbytes(cache_src);
}
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
GGML_ASSERT(gf != nullptr);
if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
return false;
}
for (const auto& input_shape_ref : plan.input_shapes) {
if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
return false;
std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
std::vector<const ggml_tensor*> tensors;
std::unordered_map<const ggml_tensor*, size_t> indices;
auto add = [&](const ggml_tensor* tensor) {
if (tensor != nullptr && indices.emplace(tensor, tensors.size() + 1).second) {
tensors.push_back(tensor);
}
ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
if (leaf == nullptr || input_shape_ref.type != leaf->type) {
return false;
};
for (int i = 0; i < graph->n_leafs; ++i) {
add(graph->leafs[i]);
}
for (int i = 0; i < graph->n_nodes; ++i) {
add(graph->nodes[i]);
}
for (size_t i = 0; i < tensors.size(); ++i) {
add(tensors[i]->view_src);
for (auto source : tensors[i]->src) {
add(source);
}
}
std::vector<uint64_t> signature;
signature.reserve(tensors.size() * 24);
signature.push_back(graph->n_nodes);
signature.push_back(graph->n_leafs);
for (int i = 0; i < graph->n_leafs; ++i) {
signature.push_back(indices.at(graph->leafs[i]));
}
for (int i = 0; i < graph->n_nodes; ++i) {
signature.push_back(indices.at(graph->nodes[i]));
}
for (auto tensor : tensors) {
signature.push_back(tensor->op);
signature.push_back(tensor->type);
signature.push_back(tensor->flags);
signature.push_back(tensor->view_offs);
if (include_bindings) {
signature.push_back(tensor->data != nullptr);
auto buffer = tensor_buffer(tensor);
signature.push_back(reinterpret_cast<uintptr_t>(buffer == nullptr ? nullptr : ggml_backend_buffer_get_type(buffer)));
}
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
return false;
}
signature.push_back(tensor->ne[d]);
signature.push_back(tensor->nb[d]);
}
signature.push_back(tensor->view_src == nullptr ? 0 : indices.at(tensor->view_src));
for (auto source : tensor->src) {
signature.push_back(source == nullptr ? 0 : indices.at(source));
}
for (int value : tensor->op_params) {
signature.push_back(static_cast<uint32_t>(value));
}
}
return true;
return signature;
}
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
GGML_ASSERT(gf != nullptr);
if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
plan.layout != graph_layout(gf, false)) {
return false;
}
for (int i = 0; i < gf->n_leafs; ++i) {
if (plan.leaf_names[i] != gf->leafs[i]->name) {
return false;
}
}
std::vector<std::pair<int, std::string>> cut_markers;
for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
auto node = ggml_graph_node(gf, i);
if (is_graph_cut_tensor(node)) {
cut_markers.emplace_back(i, node->name);
}
}
return cut_markers == plan.cut_markers;
}
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
@@ -578,26 +605,6 @@ namespace sd::ggml_graph_cut {
return tensors;
}
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
const Plan& plan,
size_t current_segment_index) {
GGML_ASSERT(gf != nullptr);
std::unordered_set<std::string> future_input_names;
for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
const auto& segment = plan.segments[seg_idx];
for (const auto& input_ref : segment.input_refs) {
if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
continue;
}
ggml_tensor* current_input = input_tensor(gf, input_ref);
if (current_input != nullptr && current_input->name[0] != '\0') {
future_input_names.insert(current_input->name);
}
}
}
return future_input_names;
}
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
const Segment& segment,
ggml_context** graph_ctx_out) {
@@ -662,6 +669,10 @@ namespace sd::ggml_graph_cut {
continue;
}
ggml_set_output(output);
if (output->view_src != nullptr) {
// A consumed output view does not keep its storage alive in gallocr.
ggml_set_output(output->view_src);
}
}
for (int node_idx : segment.internal_node_indices) {
ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
@@ -716,6 +727,10 @@ namespace sd::ggml_graph_cut {
if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) {
saved_output_flags[output] = output->flags;
}
if (output != nullptr && output->view_src != nullptr &&
saved_output_flags.find(output->view_src) == saved_output_flags.end()) {
saved_output_flags[output->view_src] = output->view_src->flags;
}
}
ggml_context* graph_ctx = nullptr;
@@ -744,6 +759,46 @@ namespace sd::ggml_graph_cut {
return buffer_size;
}
static size_t measure_graph_compute_buffer(
ggml_backend_t backend,
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set) {
struct TensorRuntimeBinding {
ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr;
void* extra = nullptr;
};
std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
auto mark_external = [&](ggml_tensor* tensor) {
if (tensor == nullptr || saved_bindings.find(tensor) != saved_bindings.end()) {
return;
}
saved_bindings[tensor] = {tensor->buffer, tensor->data, tensor->extra};
tensor->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
};
for (int i = 0; i < leaf_count(gf); ++i) {
ggml_tensor* leaf = leaf_tensor(gf, i);
if (!is_params_tensor(params_tensor_set, leaf)) {
continue;
}
mark_external(leaf);
mark_external(leaf->view_src);
}
ggml_gallocr_t allocr = ggml_gallocr_new(
ggml_backend_get_default_buffer_type(backend));
size_t sizes[1] = {0};
ggml_gallocr_reserve_n_size(allocr, gf, nullptr, nullptr, sizes);
ggml_gallocr_free(allocr);
for (const auto& kv : saved_bindings) {
kv.first->buffer = kv.second.buffer;
kv.first->data = kv.second.data;
kv.first->extra = kv.second.extra;
}
return sizes[0];
}
Plan build_plan(ggml_backend_t backend,
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
@@ -756,24 +811,22 @@ namespace sd::ggml_graph_cut {
if (n_nodes <= 0) {
return plan;
}
plan.n_nodes = n_nodes;
plan.n_leafs = gf->n_leafs;
plan.layout = graph_layout(gf, false);
for (int i = 0; i < gf->n_leafs; ++i) {
ggml_tensor* leaf = gf->leafs[i];
if (is_params_tensor(params_tensor_set, leaf)) {
continue;
}
auto shape = input_shape(leaf);
shape.leaf_index = i;
plan.input_shapes.push_back(shape);
plan.leaf_names.emplace_back(gf->leafs[i]->name);
}
plan.compute_buffer_size =
measure_graph_compute_buffer(backend, gf, params_tensor_set);
std::unordered_map<const ggml_tensor*, int> producer_index;
producer_index.reserve(static_cast<size_t>(n_nodes));
for (int i = 0; i < n_nodes; ++i) {
producer_index[ggml_graph_node(gf, i)] = i;
ggml_tensor* node = ggml_graph_node(gf, i);
producer_index[node] = i;
if (is_graph_cut_tensor(node)) {
plan.cut_markers.push_back({i, node->name});
}
}
std::vector<Segment> grouped_segments;
std::unordered_map<std::string, size_t> group_to_segment;
for (int i = 0; i < n_nodes; ++i) {
@@ -824,11 +877,24 @@ namespace sd::ggml_graph_cut {
if (final_output_index < 0) {
final_output_index = n_nodes - 1;
}
ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
Segment final_segment;
final_segment.group_name = "ggml_runner.final";
Segment final_segment;
final_segment.group_name = "ggml_runner.final";
if (final_output_index >= 0 &&
available_cut_output_node_indices.find(final_output_index) ==
available_cut_output_node_indices.end()) {
final_segment.output_node_indices.push_back(final_output_index);
}
for (int i = 0; i < n_nodes; ++i) {
ggml_tensor* node = ggml_graph_node(gf, i);
if (i == final_output_index || node == nullptr ||
(node->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 ||
available_cut_output_node_indices.find(i) !=
available_cut_output_node_indices.end()) {
continue;
}
final_segment.output_node_indices.push_back(i);
}
if (!final_segment.output_node_indices.empty()) {
build_segment(gf,
plan,
final_segment,
@@ -839,223 +905,53 @@ namespace sd::ggml_graph_cut {
log_desc);
}
return plan;
}
Plan apply_max_vram_budget(ggml_cgraph* gf,
const Plan& base_plan,
size_t max_graph_vram_bytes,
ggml_backend_t backend,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc) {
GGML_ASSERT(backend != nullptr);
GGML_ASSERT(gf != nullptr);
int64_t t_budget_begin = ggml_time_ms();
if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
return base_plan;
}
const int n_nodes = ggml_graph_n_nodes(gf);
std::unordered_map<const ggml_tensor*, int> producer_index;
producer_index.reserve(static_cast<size_t>(n_nodes));
for (int i = 0; i < n_nodes; ++i) {
producer_index[ggml_graph_node(gf, i)] = i;
}
Plan merged_plan;
merged_plan.available = true;
merged_plan.has_cuts = base_plan.has_cuts;
merged_plan.valid = base_plan.valid;
merged_plan.n_nodes = base_plan.n_nodes;
merged_plan.n_leafs = base_plan.n_leafs;
std::unordered_set<int> available_cut_output_node_indices;
available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
size_t start_segment_index = 0;
while (start_segment_index < base_plan.segments.size()) {
Plan single_plan;
auto single_available_cut_output_node_indices = available_cut_output_node_indices;
auto single_seed = make_segment_seed(base_plan,
start_segment_index,
start_segment_index);
build_segment(gf,
single_plan,
single_seed,
producer_index,
single_available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
GGML_ASSERT(!single_plan.segments.empty());
size_t best_end_segment_index = start_segment_index;
bool can_merge_next_segment = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
const size_t next_end_segment_index = best_end_segment_index + 1;
Plan candidate_plan;
auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
auto candidate_seed = make_segment_seed(base_plan,
start_segment_index,
next_end_segment_index);
build_segment(gf,
candidate_plan,
candidate_seed,
producer_index,
candidate_available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
GGML_ASSERT(!candidate_plan.segments.empty());
const auto& candidate_segment = candidate_plan.segments.back();
const size_t candidate_bytes = graph_cut_segment_vram_bytes(candidate_segment);
if (candidate_bytes > max_graph_vram_bytes) {
break;
std::unordered_set<std::string> future_cut_names;
for (auto segment = plan.segments.rbegin(); segment != plan.segments.rend(); ++segment) {
segment->future_cut_names = future_cut_names;
segment->live_cut_names = future_cut_names;
for (const auto& input : segment->input_refs) {
if (input.type != Segment::INPUT_PREVIOUS_CUT) {
continue;
}
best_end_segment_index = next_end_segment_index;
segment->live_cut_names.insert(input.display_name);
future_cut_names.insert(input.display_name);
}
auto best_seed = make_segment_seed(base_plan,
start_segment_index,
best_end_segment_index);
build_segment(gf,
merged_plan,
best_seed,
producer_index,
available_cut_output_node_indices,
backend,
params_tensor_set,
log_desc);
start_segment_index = best_end_segment_index + 1;
}
if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
std::string plan_validation_error;
plan.valid = validate_plan(gf, plan, &plan_validation_error);
if (!plan.valid && log_desc != nullptr) {
LOG_WARN("%s graph cut plan validation failed (%s); using monolithic execution",
log_desc,
max_graph_vram_bytes / 1024.0 / 1024.0,
base_plan.segments.size(),
merged_plan.segments.size());
plan_validation_error.c_str());
}
if (log_desc != nullptr) {
LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
log_desc,
ggml_time_ms() - t_budget_begin);
}
return merged_plan;
return plan;
}
Plan resolve_plan(ggml_backend_t backend,
ggml_cgraph* gf,
PlanCache* cache,
size_t max_graph_vram_bytes,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc) {
GGML_ASSERT(backend != nullptr);
GGML_ASSERT(gf != nullptr);
GGML_ASSERT(cache != nullptr);
int64_t t_prepare_begin = ggml_time_ms();
Plan base_plan;
int64_t t_plan_begin = ggml_time_ms();
if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
base_plan = cache->graph_cut_plan;
} else {
base_plan = build_plan(backend, gf, params_tensor_set, log_desc);
cache->graph_cut_plan = base_plan;
cache->graph_cut_plan.available = true;
cache->budgeted_graph_cut_plan.available = false;
if (log_desc != nullptr) {
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
}
if (cache->graph_cut_plan.available &&
plan_matches_graph(gf, cache->graph_cut_plan)) {
return cache->graph_cut_plan;
}
Plan resolved_plan = base_plan;
if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
if (cache->budgeted_graph_cut_plan.available &&
cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
resolved_plan = cache->budgeted_graph_cut_plan;
} else {
resolved_plan = apply_max_vram_budget(gf,
base_plan,
max_graph_vram_bytes,
backend,
params_tensor_set,
log_desc);
cache->budgeted_graph_cut_plan = resolved_plan;
cache->budgeted_graph_cut_plan.available = true;
cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
}
}
return resolved_plan;
}
void annotate_residency(Plan& plan,
size_t max_graph_vram_bytes,
bool prefetch_enabled) {
// Cached plans may be reused with a smaller live budget.
for (auto& seg : plan.segments) {
seg.residency = SegmentResidency::STREAMED;
}
if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
return;
}
bool any_param_bearing = false;
for (const auto& seg : plan.segments) {
if (seg.input_param_bytes > 0) {
any_param_bearing = true;
break;
}
}
if (!any_param_bearing) {
return;
}
// Leave room for the largest active streamed segment.
size_t worst_streamed_footprint = 0;
size_t prefetch_headroom = 0;
for (const auto& seg : plan.segments) {
const size_t seg_footprint = seg.input_param_bytes +
seg.compute_buffer_size +
seg.output_bytes +
seg.input_previous_cut_bytes +
seg.input_external_bytes;
if (seg_footprint > worst_streamed_footprint) {
worst_streamed_footprint = seg_footprint;
}
prefetch_headroom = std::max(prefetch_headroom, seg.input_param_bytes);
}
constexpr size_t safety = 512ull * 1024 * 1024;
if (worst_streamed_footprint > SIZE_MAX - safety) {
return;
}
size_t reserved = safety + worst_streamed_footprint;
if (prefetch_enabled) {
if (prefetch_headroom > SIZE_MAX - reserved) {
return;
}
reserved += prefetch_headroom;
}
if (max_graph_vram_bytes <= reserved) {
return;
}
const size_t available = max_graph_vram_bytes - reserved;
size_t cumulative = 0;
for (auto& seg : plan.segments) {
if (cumulative + seg.input_param_bytes > available) {
break;
}
seg.residency = SegmentResidency::RESIDENT;
cumulative += seg.input_param_bytes;
int64_t t_plan_begin = ggml_time_ms();
Plan plan = build_plan(backend, gf, params_tensor_set, log_desc);
cache->graph_cut_plan = plan;
if (log_desc != nullptr) {
LOG_INFO("%s build cached graph cut plan done (taking %lld ms)",
log_desc,
ggml_time_ms() - t_plan_begin);
}
return plan;
}
} // namespace sd::ggml_graph_cut
src/core/ggml_graph_cut.h
+13 -41
View File
@@ -6,19 +6,13 @@
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
namespace sd::ggml_graph_cut {
// Streaming residency for a segment's params.
enum class SegmentResidency : uint8_t {
STREAMED = 0,
RESIDENT = 1,
};
struct Segment {
enum InputType {
INPUT_EXTERNAL = 0,
@@ -33,38 +27,28 @@ namespace sd::ggml_graph_cut {
int node_index = -1;
};
size_t compute_buffer_size = 0;
size_t output_bytes = 0;
size_t input_external_bytes = 0;
size_t input_previous_cut_bytes = 0;
size_t input_param_bytes = 0;
size_t compute_buffer_size = 0;
std::string group_name;
std::vector<int> internal_node_indices;
std::vector<int> output_node_indices;
std::vector<InputRef> input_refs;
SegmentResidency residency = SegmentResidency::STREAMED;
std::unordered_set<std::string> future_cut_names;
std::unordered_set<std::string> live_cut_names;
};
struct Plan {
struct InputShape {
int leaf_index = -1;
ggml_type type = GGML_TYPE_COUNT;
std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
};
bool available = false;
bool has_cuts = false;
bool valid = true;
int n_nodes = 0;
int n_leafs = 0;
std::vector<InputShape> input_shapes;
bool available = false;
bool has_cuts = false;
bool valid = true;
size_t compute_buffer_size = 0;
std::vector<uint64_t> layout;
std::vector<std::string> leaf_names;
std::vector<std::pair<int, std::string>> cut_markers;
std::vector<Segment> segments;
};
struct PlanCache {
Plan graph_cut_plan;
Plan budgeted_graph_cut_plan;
size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
};
static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
@@ -89,13 +73,12 @@ namespace sd::ggml_graph_cut {
ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
size_t cache_tensor_bytes(const ggml_tensor* tensor);
// Plans ignore runtime bindings; allocator reservations must include them.
std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings);
bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
const Plan& plan,
size_t current_segment_index);
ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
const Segment& segment,
ggml_context** graph_ctx_out);
@@ -109,23 +92,12 @@ namespace sd::ggml_graph_cut {
ggml_cgraph* gf,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
Plan apply_max_vram_budget(ggml_cgraph* gf,
const Plan& base_plan,
size_t max_graph_vram_bytes,
ggml_backend_t backend,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
Plan resolve_plan(ggml_backend_t backend,
ggml_cgraph* gf,
PlanCache* cache,
size_t max_graph_vram_bytes,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
// Mark leading segments resident after reserving streamed execution headroom.
void annotate_residency(Plan& plan,
size_t max_graph_vram_bytes,
bool prefetch_enabled);
} // namespace sd::ggml_graph_cut
#endif // __SD_CORE_GGML_GRAPH_CUT_H__
src/core/ggml_runner.cpp
+296
View File
@@ -0,0 +1,296 @@
#include <algorithm>
#include <map>
#include <utility>
#include "core/ggml_extend.hpp"
#include "core/segment_graph_bindings.h"
#include "core/segment_weight_pipeline.h"
using namespace sd;
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, size_t direct_bytes) {
auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
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, direct_bytes, external_backend, assign_nodes);
}
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) {
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_status status;
if (scheduler != nullptr) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str());
multi_device_eval_callback_warned = true;
}
status = ggml_backend_sched_graph_compute(scheduler, graph);
} else {
status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph,
sd_get_backend_eval_callback(),
sd_get_backend_eval_callback_data());
}
workspace_.synchronize();
if (status != GGML_STATUS_SUCCESS) {
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);
ggml_graph_cut::Plan plan;
if (!resolve_graph_cut_plan(graph, &plan)) {
return std::nullopt;
}
const auto full_measurement = measure(graph, plan.compute_buffer_size);
if (full_measurement.buffers.empty()) {
return std::nullopt;
}
auto manager = residency_manager.lock();
const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
manager != nullptr && manager->segmented_compute_enabled() &&
plan.valid && plan.has_cuts && plan.segments.size() > 1 &&
!fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
if (!segmented) {
ggml_graph_cut::Segment segment;
segment.group_name = "graph";
segment.compute_buffer_size = plan.compute_buffer_size;
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);
}
plan.segments = {std::move(segment)};
}
const bool segments_changed = plan.segments.size() != logged_segment_count_;
if (segments_changed && (segmented || logged_segment_count_ > 1)) {
LOG_DEBUG("%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() {
runner.workspace_.segment_end();
bindings.restore();
weights.segment_end();
ggml_free(context);
runner.sync_runtime_residency();
}
} segment_cleanup{*this, weights, bindings, segment_context};
auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
if (!workspace_.prepare(measurement)) {
return fail_segment("workspace preparation");
}
const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
const 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)) && workspace_.release_excess(measurement)) {
sync_runtime_residency();
requests = memory_requests(measurement.buffers, new_cache_bytes);
}
return weights.ensure_segment_capacity(index, requests);
};
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_.measurement_matches(segment_graph, measurement)) {
measurement = measure(segment_graph, segment.compute_buffer_size);
}
if (!workspace_.prepare(measurement) || !ensure_capacity()) {
return fail_segment("workspace capacity check");
}
if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
pin_multi_device_nodes(scheduler, current);
})) {
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());
}
if (!execute_segment(segment_graph, n_threads) ||
!cache_.capture(segment_graph) ||
!cut_cache_.capture(graph, segment, get_desc().c_str())) {
return fail_segment("execution or 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<float>(result, "output");
if (!output.has_value()) {
return fail_segment("output readback");
}
}
}
// 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_DEBUG("%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;
}
src/core/layer_split_partition.cpp
+12 -4
View File
@@ -145,19 +145,24 @@ namespace sd {
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
backend_vram_limits,
primary_backend_vram_limit);
// Existing placements may already occupy the reported free VRAM. Reuse
// them; execution checks missing weights and reclaims memory as needed.
const bool reuse_assignments = std::all_of(seen_params.begin(), seen_params.end(), [&](ggml_tensor* param) {
return param_assignments.count(param) != 0;
});
std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
size_t current_backend = 0;
int64_t current_used = 0;
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
int64_t bytes = segment_param_bytes[seg_idx];
while (current_backend + 1 < split_backends.size() &&
while (!reuse_assignments && current_backend + 1 < split_backends.size() &&
bytes > 0 &&
current_used + bytes > backend_capacities[current_backend]) {
current_backend++;
current_used = 0;
}
if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
if (!reuse_assignments && bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
desc,
seg_idx,
@@ -167,7 +172,6 @@ namespace sd {
return false;
}
current_used += bytes;
backend_by_segment[seg_idx] = split_backends[current_backend];
for (ggml_tensor* param : segment_params[seg_idx]) {
ggml_backend_t target_backend = split_backends[current_backend];
@@ -186,12 +190,16 @@ namespace sd {
ggml_get_name(param));
return false;
}
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
if (reuse_assignments) {
current_backend = backend_idx;
}
assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
assignment.tensors_by_backend[backend_idx].push_back(param);
assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
}
backend_by_segment[seg_idx] = split_backends[current_backend];
}
const int n_nodes = ggml_graph_n_nodes(gf);
src/core/layer_stream_prefetch.cpp
-124
View File
@@ -1,124 +0,0 @@
#include "core/layer_stream_prefetch.h"
#include <utility>
#include "core/ggml_graph_cut.h"
#include "weight_manager.h"
namespace sd {
static ggml_tensor* canonical_param(
ggml_tensor* tensor,
const std::unordered_set<const ggml_tensor*>& params) {
for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
if (params.find(current) != params.end()) {
return current;
}
}
return nullptr;
}
LayerStreamPrefetch::LayerStreamPrefetch(
const std::shared_ptr<RunnerWeightManager>& manager,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled)
: manager_(manager),
owner_id_(owner_id),
enabled_(enabled && manager != nullptr) {
segment_params_.resize(plan.segments.size());
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
std::unordered_set<ggml_tensor*> seen;
for (ggml_tensor* tensor :
ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
ggml_tensor* param = canonical_param(tensor, params);
if (param != nullptr && seen.insert(param).second) {
segment_params_[segment_index].push_back(param);
}
}
}
}
LayerStreamPrefetch::~LayerStreamPrefetch() {
clear();
}
size_t LayerStreamPrefetch::next_parameter_segment(size_t segment_index) const {
for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
if (!segment_params_[next].empty()) {
return next;
}
}
return SIZE_MAX;
}
void LayerStreamPrefetch::disable() {
clear();
enabled_ = false;
}
bool LayerStreamPrefetch::activate(size_t segment_index) {
if (!enabled_ || queued_segment_ == SIZE_MAX) {
return true;
}
if (queued_segment_ != segment_index) {
return true;
}
auto manager = manager_.lock();
if (manager == nullptr ||
!manager->activate_prefetched_params(owner_id_, queued_params_)) {
disable();
return false;
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
return true;
}
bool LayerStreamPrefetch::enqueue_next(size_t segment_index) {
if (!enabled_) {
return true;
}
if (queued_segment_ != SIZE_MAX) {
return true;
}
const size_t next_segment = next_parameter_segment(segment_index);
if (next_segment == SIZE_MAX) {
return true;
}
std::unordered_set<ggml_tensor*> active_params(
segment_params_[segment_index].begin(),
segment_params_[segment_index].end());
std::vector<ggml_tensor*> params;
params.reserve(segment_params_[next_segment].size());
for (ggml_tensor* param : segment_params_[next_segment]) {
if (active_params.find(param) == active_params.end()) {
params.push_back(param);
}
}
if (params.empty()) {
return true;
}
auto manager = manager_.lock();
if (manager == nullptr || !manager->prefetch_params(owner_id_, params)) {
disable();
return false;
}
queued_params_ = std::move(params);
queued_segment_ = next_segment;
return true;
}
void LayerStreamPrefetch::clear() {
if (auto manager = manager_.lock()) {
manager->clear_prefetched_params(owner_id_);
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
}
src/core/layer_stream_prefetch.h
-48
View File
@@ -1,48 +0,0 @@
#ifndef __SD_CORE_LAYER_STREAM_PREFETCH_H__
#define __SD_CORE_LAYER_STREAM_PREFETCH_H__
#include <cstddef>
#include <cstdint>
#include <memory>
#include <unordered_set>
#include <vector>
struct ggml_cgraph;
struct ggml_tensor;
struct RunnerWeightManager;
namespace sd::ggml_graph_cut {
struct Plan;
}
namespace sd {
class LayerStreamPrefetch {
private:
std::weak_ptr<RunnerWeightManager> manager_;
uintptr_t owner_id_ = 0;
std::vector<std::vector<ggml_tensor*>> segment_params_;
std::vector<ggml_tensor*> queued_params_;
size_t queued_segment_ = SIZE_MAX;
bool enabled_ = true;
size_t next_parameter_segment(size_t segment_index) const;
void disable();
public:
LayerStreamPrefetch(
const std::shared_ptr<RunnerWeightManager>& manager,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled = true);
~LayerStreamPrefetch();
bool enabled() const { return enabled_; }
bool activate(size_t segment_index);
bool enqueue_next(size_t segment_index);
void clear();
};
}
#endif // __SD_CORE_LAYER_STREAM_PREFETCH_H__
src/core/runner_cache.cpp
+211
View File
@@ -0,0 +1,211 @@
#include "core/runner_cache.h"
#include <algorithm>
#include <iterator>
#include <unordered_set>
#include "core/ggml_graph_cut.h"
#include "core/util.h"
namespace sd {
static std::unordered_set<const ggml_tensor*> cache_graph_tensors(ggml_cgraph* graph) {
std::unordered_set<const ggml_tensor*> tensors;
for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
tensors.insert(ggml_graph_node(graph, i));
}
for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
tensors.insert(ggml_graph_cut::leaf_tensor(graph, i));
}
return tensors;
}
CachedTensor::~CachedTensor() {
ggml_backend_buffer_free(buffer);
ggml_free(context);
}
std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source) {
if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
return nullptr;
}
auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) {
return nullptr;
}
entry->tensor = ggml_dup_tensor(entry->context, source);
// Cut views are rebound with their original strides and offsets.
std::copy(std::begin(source->nb), std::end(source->nb), std::begin(entry->tensor->nb));
ggml_set_name(entry->tensor, name.c_str());
entry->buffer = ggml_backend_alloc_ctx_tensors(entry->context, backend);
if (entry->buffer == nullptr) {
return nullptr;
}
if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) {
std::vector<uint8_t> data(ggml_nbytes(source));
ggml_backend_tensor_get(source, data.data(), 0, data.size());
ggml_backend_tensor_set(entry->tensor, data.data(), 0, data.size());
} else {
ggml_backend_tensor_copy(source, entry->tensor);
}
return entry;
}
static ggml_tensor* cached_tensor(const CachedTensors& tensors, const std::string& name) {
auto entry = tensors.find(name);
return entry == tensors.end() ? nullptr : entry->second->tensor;
}
static size_t resident_bytes(const CachedTensors& tensors, ggml_backend_dev_t device) {
size_t bytes = 0;
for (const auto& entry : tensors) {
auto buffer = entry.second->buffer;
if (!ggml_backend_buffer_is_host(buffer) &&
ggml_backend_buft_get_device(ggml_backend_buffer_get_type(buffer)) == device) {
const size_t size = ggml_backend_buffer_get_size(buffer);
bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
}
}
return bytes;
}
ggml_tensor* RunnerCache::get(const std::string& name) const {
return cached_tensor(committed_, name);
}
void RunnerCache::stage(const std::string& name, ggml_tensor* tensor) {
if (tensor != nullptr) {
ggml_set_output(tensor);
outputs_[name] = tensor;
}
}
size_t RunnerCache::pending_bytes(ggml_cgraph* graph) const {
if (outputs_.empty()) {
return 0;
}
auto tensors = cache_graph_tensors(graph);
auto buft = ggml_backend_get_default_buffer_type(backend_);
size_t bytes = 0;
for (const auto& output : outputs_) {
if (pending_.count(output.first) || !tensors.count(output.second)) {
continue;
}
const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, output.second),
ggml_backend_buft_get_alignment(buft));
bytes = size > SIZE_MAX - bytes ? SIZE_MAX : bytes + size;
}
return bytes;
}
size_t RunnerCache::resident_bytes(ggml_backend_dev_t device) const {
const size_t committed = sd::resident_bytes(committed_, device);
const size_t pending = sd::resident_bytes(pending_, device);
return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
}
bool RunnerCache::capture(ggml_cgraph* graph) {
if (outputs_.empty()) {
return true;
}
const auto tensors = cache_graph_tensors(graph);
for (const auto& output : outputs_) {
if (pending_.count(output.first) || !tensors.count(output.second)) {
continue;
}
GGML_ASSERT(ggml_is_contiguous(output.second));
auto entry = CachedTensor::copy(backend_, output.first, output.second);
if (entry == nullptr) {
return false;
}
pending_[output.first] = std::move(entry);
}
ggml_backend_synchronize(backend_);
return true;
}
void RunnerCache::graph_end(bool success) {
// Graph inputs can still reference the previous generation until graph end.
if (success) {
for (auto& entry : pending_) {
committed_[entry.first] = std::move(entry.second);
}
}
pending_.clear();
outputs_.clear();
}
void RunnerCache::clear() {
graph_end(false);
committed_.clear();
}
ggml_tensor* GraphCutTensorCache::get(const std::string& name) const {
return cached_tensor(tensors_, name);
}
size_t GraphCutTensorCache::resident_bytes(ggml_backend_dev_t device) const {
return sd::resident_bytes(tensors_, device);
}
size_t GraphCutTensorCache::estimate_output_bytes(
ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const {
ggml_backend_buffer_type_t buffer_type =
ggml_backend_get_default_buffer_type(backend_);
if (buffer_type == nullptr) {
return SIZE_MAX;
}
const size_t alignment = ggml_backend_buft_get_alignment(buffer_type);
size_t total_size = 0;
for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
ggml_tensor* output = ggml_graph_cut::output_tensor(graph, segment, output_idx);
if (output == nullptr || !ggml_graph_cut::is_graph_cut_tensor(output) ||
!segment.future_cut_names.count(output->name)) {
continue;
}
ggml_tensor* source = ggml_graph_cut::cache_source_tensor(output);
const size_t tensor_size = GGML_PAD(
ggml_backend_buft_get_alloc_size(buffer_type, source), alignment);
total_size = tensor_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + tensor_size;
}
return total_size;
}
void GraphCutTensorCache::prune(const std::unordered_set<std::string>& keep_names) {
for (auto it = tensors_.begin(); it != tensors_.end();) {
it = keep_names.count(it->first) ? std::next(it) : tensors_.erase(it);
}
}
bool GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
size_t copied_bytes = 0;
size_t copied_count = 0;
for (int index : segment.output_node_indices) {
auto output = ggml_graph_node(graph, index);
if (!ggml_graph_cut::is_graph_cut_tensor(output) ||
!segment.future_cut_names.count(output->name)) {
continue;
}
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output));
if (entry == nullptr) {
LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
return false;
}
const size_t size = ggml_backend_buffer_get_size(entry->buffer);
copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
++copied_count;
tensors_[output->name] = std::move(entry);
}
ggml_backend_synchronize(backend_);
if (copied_count > 0) {
LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
}
return true;
}
}
src/core/runner_cache.h
+66
View File
@@ -0,0 +1,66 @@
#ifndef __SD_CORE_RUNNER_CACHE_H__
#define __SD_CORE_RUNNER_CACHE_H__
#include <map>
#include <memory>
#include <string>
#include <unordered_set>
#include "ggml-backend.h"
namespace sd::ggml_graph_cut {
struct Segment;
}
namespace sd {
struct CachedTensor {
ggml_context* context = nullptr;
ggml_backend_buffer_t buffer = nullptr;
ggml_tensor* tensor = nullptr;
~CachedTensor();
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source);
};
using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
class RunnerCache {
ggml_backend_t backend_;
CachedTensors committed_;
CachedTensors pending_;
std::map<std::string, ggml_tensor*> outputs_;
public:
explicit RunnerCache(ggml_backend_t backend)
: backend_(backend) {}
RunnerCache(const RunnerCache&) = delete;
RunnerCache& operator=(const RunnerCache&) = delete;
ggml_tensor* get(const std::string& name) const;
void stage(const std::string& name, ggml_tensor* tensor);
const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
size_t pending_bytes(ggml_cgraph* graph) const;
size_t resident_bytes(ggml_backend_dev_t device) const;
bool capture(ggml_cgraph* graph);
void graph_end(bool success);
void clear();
};
class GraphCutTensorCache {
ggml_backend_t backend_;
CachedTensors tensors_;
public:
explicit GraphCutTensorCache(ggml_backend_t backend)
: backend_(backend) {}
ggml_tensor* get(const std::string& name) const;
size_t resident_bytes(ggml_backend_dev_t device) const;
size_t estimate_output_bytes(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const;
bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
void prune(const std::unordered_set<std::string>& keep_names);
void clear() { tensors_.clear(); }
};
}
#endif // __SD_CORE_RUNNER_CACHE_H__
src/core/segment_graph_bindings.cpp
+137
View File
@@ -0,0 +1,137 @@
#include "core/segment_graph_bindings.h"
#include <algorithm>
#include <iterator>
#include "core/ggml_graph_cut.h"
#include "core/runner_cache.h"
#include "core/util.h"
#include "ggml.h"
namespace sd {
SegmentGraphBindings::SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
const ggml_graph_cut::Plan& plan,
ggml_cgraph* graph)
: tensor_cache_(tensor_cache),
graph_(graph) {
GGML_ASSERT(graph_ != nullptr);
for (int i = 0; i < ggml_graph_n_nodes(graph_); ++i) {
ggml_tensor* tensor = ggml_graph_node(graph_, i);
Topology topology{tensor->op, {}, tensor->view_src, tensor->flags};
std::copy(std::begin(tensor->src), std::end(tensor->src), topology.sources.begin());
topology_[tensor] = topology;
}
for (const auto& segment : plan.segments) {
for (const auto& input : segment.input_refs) {
if (input.type != ggml_graph_cut::Segment::INPUT_EXTERNAL) {
continue;
}
ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
if (tensor == nullptr || tensor->buffer == nullptr) {
continue;
}
external_bindings_[tensor] = {tensor->buffer, tensor->data, tensor->extra};
}
}
}
void SegmentGraphBindings::reset(const ggml_graph_cut::Segment& segment) {
restore();
for (const auto& input : segment.input_refs) {
ggml_tensor* tensor = ggml_graph_cut::input_tensor(graph_, input);
if (tensor == nullptr) {
continue;
}
switch (input.type) {
case ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT:
tensor->buffer = nullptr;
tensor->data = nullptr;
tensor->extra = nullptr;
break;
case ggml_graph_cut::Segment::INPUT_EXTERNAL: {
auto binding = external_bindings_.find(tensor);
if (binding != external_bindings_.end()) {
tensor->buffer = binding->second.buffer;
tensor->data = binding->second.data;
tensor->extra = binding->second.extra;
} else {
tensor->buffer = nullptr;
tensor->data = nullptr;
tensor->extra = nullptr;
}
break;
}
case ggml_graph_cut::Segment::INPUT_PARAM:
break;
}
}
for (int node_index : segment.internal_node_indices) {
ggml_tensor* node = ggml_graph_node(graph_, node_index);
if (node == nullptr) {
continue;
}
node->buffer = nullptr;
node->data = nullptr;
node->extra = nullptr;
}
}
void SegmentGraphBindings::restore() {
for (const auto& entry : topology_) {
entry.first->op = entry.second.op;
entry.first->view_src = entry.second.view_source;
entry.first->flags = entry.second.flags;
std::copy(entry.second.sources.begin(), entry.second.sources.end(), std::begin(entry.first->src));
}
}
bool SegmentGraphBindings::bind_cached_inputs(
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
std::unordered_map<ggml_tensor*, ggml_tensor*> cached_view_sources;
for (const auto& input : segment.input_refs) {
if (input.type != ggml_graph_cut::Segment::INPUT_PREVIOUS_CUT) {
continue;
}
ggml_tensor* input_tensor = ggml_graph_cut::input_tensor(graph_, input);
if (input_tensor == nullptr) {
continue;
}
ggml_tensor* cached_tensor = tensor_cache_.get(input.display_name);
if (cached_tensor == nullptr) {
LOG_ERROR("%s missing graph cut cache tensor: %s",
log_desc,
input.display_name.c_str());
return false;
}
if (input_tensor->view_src != nullptr) {
cached_view_sources[topology_.at(input_tensor).view_source] = cached_tensor;
input_tensor->view_src = cached_tensor;
input_tensor->buffer = nullptr;
input_tensor->data = cached_tensor->data == nullptr
? nullptr
: static_cast<void*>(static_cast<char*>(cached_tensor->data) +
input_tensor->view_offs);
input_tensor->extra = cached_tensor->extra;
} else {
input_tensor->buffer = cached_tensor->buffer;
input_tensor->data = cached_tensor->data;
input_tensor->extra = cached_tensor->extra;
}
for (int source_index = 0; source_index < GGML_MAX_SRC; ++source_index) {
input_tensor->src[source_index] = nullptr;
}
input_tensor->op = GGML_OP_NONE;
}
// ggml flattens view chains, so descendants also need the cached root.
for (int node_index : segment.internal_node_indices) {
ggml_tensor* node = ggml_graph_node(graph_, node_index);
auto cached_source = cached_view_sources.find(topology_.at(node).view_source);
if (cached_source != cached_view_sources.end()) {
node->view_src = cached_source->second;
}
}
return true;
}
}
src/core/segment_graph_bindings.h
+52
View File
@@ -0,0 +1,52 @@
#ifndef __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
#define __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
#include <array>
#include <unordered_map>
#include "ggml.h"
struct ggml_backend_buffer;
struct ggml_cgraph;
struct ggml_tensor;
namespace sd {
class GraphCutTensorCache;
namespace ggml_graph_cut {
struct Plan;
struct Segment;
}
class SegmentGraphBindings {
public:
SegmentGraphBindings(GraphCutTensorCache& tensor_cache,
const ggml_graph_cut::Plan& plan,
ggml_cgraph* graph);
void reset(const ggml_graph_cut::Segment& segment);
void restore();
~SegmentGraphBindings() { restore(); }
bool bind_cached_inputs(const ggml_graph_cut::Segment& segment,
const char* log_desc);
private:
struct ExternalBinding {
ggml_backend_buffer* buffer = nullptr;
void* data = nullptr;
void* extra = nullptr;
};
GraphCutTensorCache& tensor_cache_;
ggml_cgraph* graph_ = nullptr;
std::unordered_map<ggml_tensor*, ExternalBinding> external_bindings_;
struct Topology {
ggml_op op;
std::array<ggml_tensor*, GGML_MAX_SRC> sources;
ggml_tensor* view_source;
int flags;
};
std::unordered_map<ggml_tensor*, Topology> topology_;
};
}
#endif // __SD_CORE_SEGMENT_GRAPH_BINDINGS_H__
src/core/segment_weight_pipeline.cpp
+205
View File
@@ -0,0 +1,205 @@
#include "core/segment_weight_pipeline.h"
#include <utility>
#include "core/ggml_graph_cut.h"
#include "device_residency_manager.h"
namespace sd {
static ggml_tensor* canonical_param(
ggml_tensor* tensor,
const std::unordered_set<const ggml_tensor*>& params) {
for (ggml_tensor* current = tensor; current != nullptr; current = current->view_src) {
if (params.find(current) != params.end()) {
return current;
}
}
return nullptr;
}
SegmentWeightPipeline::SegmentWeightPipeline(
const std::shared_ptr<DeviceResidencyManager>& residency_manager,
ggml_backend_t compute_backend,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled)
: residency_manager_(residency_manager),
compute_backend_(compute_backend),
owner_id_(owner_id),
enabled_(enabled && residency_manager != nullptr) {
segment_params_.resize(plan.segments.size());
for (size_t segment_index = 0; segment_index < plan.segments.size(); ++segment_index) {
std::unordered_set<ggml_tensor*> seen;
for (ggml_tensor* tensor :
ggml_graph_cut::param_tensors(graph, plan.segments[segment_index])) {
ggml_tensor* param = canonical_param(tensor, params);
if (param != nullptr && seen.insert(param).second) {
segment_params_[segment_index].push_back(param);
}
}
}
}
SegmentWeightPipeline::~SegmentWeightPipeline() {
segment_end();
clear();
}
size_t SegmentWeightPipeline::next_parameter_segment(size_t segment_index) const {
for (size_t next = segment_index + 1; next < segment_params_.size(); ++next) {
if (!segment_params_[next].empty()) {
return next;
}
}
return SIZE_MAX;
}
std::vector<std::vector<ggml_tensor*>> SegmentWeightPipeline::preferred_eviction_order() const {
return {segment_params_.rbegin(), segment_params_.rend()};
}
void SegmentWeightPipeline::disable() {
clear();
enabled_ = false;
}
void SegmentWeightPipeline::activate(size_t segment_index) {
if (!enabled_ || queued_segment_ == SIZE_MAX || queued_segment_ != segment_index) {
return;
}
auto manager = residency_manager_.lock();
if (manager == nullptr ||
!manager->activate_prefetched_params(owner_id_, queued_params_)) {
disable();
return;
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
bool SegmentWeightPipeline::ensure_segment_capacity(
size_t segment_index,
const std::vector<DeviceMemoryRequest>& requests) {
if (segment_index >= segment_params_.size()) {
return false;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
return segment_params_[segment_index].empty();
}
std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
protected_params.insert(protected_params.end(), queued_params_.begin(), queued_params_.end());
for (const auto& request : requests) {
if (!manager->ensure_compute_backend_capacity(request, segment_params_[segment_index],
preferred_eviction_order(), protected_params)) {
return false;
}
}
return true;
}
bool SegmentWeightPipeline::segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity) {
GGML_ASSERT(pinned_params_.empty());
activate(segment_index);
if (!ensure_capacity()) {
return false;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
return segment_params_[segment_index].empty();
}
if (!manager->prepare_params(segment_params_[segment_index])) {
return false;
}
pinned_params_ = segment_params_[segment_index];
return true;
}
void SegmentWeightPipeline::segment_end() {
if (auto manager = residency_manager_.lock()) {
manager->release_compute_backend_params(pinned_params_);
}
pinned_params_.clear();
}
void SegmentWeightPipeline::enqueue_next(
size_t segment_index,
const DeviceMemoryRequest& request) {
if (!enabled_ || queued_segment_ != SIZE_MAX) {
return;
}
const size_t next_segment = next_parameter_segment(segment_index);
if (next_segment == SIZE_MAX) {
return;
}
std::unordered_set<ggml_tensor*> active_params(
segment_params_[segment_index].begin(),
segment_params_[segment_index].end());
std::vector<ggml_tensor*> params;
params.reserve(segment_params_[next_segment].size());
for (ggml_tensor* param : segment_params_[next_segment]) {
if (active_params.find(param) == active_params.end()) {
params.push_back(param);
}
}
if (params.empty()) {
return;
}
auto manager = residency_manager_.lock();
if (manager == nullptr) {
disable();
return;
}
const WeightResidencyInfo residency =
manager->inspect_compute_backend_params(params);
if (residency.missing_bytes == 0) {
return;
}
if (!residency.async_prefetch_supported) {
disable();
return;
}
DeviceMemoryRequest backend_request = request;
backend_request.compute_backend = compute_backend_;
backend_request.owner_id = owner_id_;
std::vector<ggml_tensor*> protected_params = segment_params_[segment_index];
protected_params.insert(protected_params.end(), params.begin(), params.end());
if (!manager->ensure_compute_backend_capacity(backend_request,
params,
preferred_eviction_order(),
protected_params)) {
disable();
return;
}
switch (manager->prefetch_params(owner_id_, params)) {
case WeightPrefetchResult::Scheduled:
queued_params_ = std::move(params);
queued_segment_ = next_segment;
return;
case WeightPrefetchResult::AlreadyResident:
return;
case WeightPrefetchResult::Unsupported:
disable();
return;
case WeightPrefetchResult::Failed:
disable();
return;
}
disable();
}
void SegmentWeightPipeline::clear() {
if (auto manager = residency_manager_.lock()) {
manager->clear_prefetched_params(owner_id_);
}
queued_params_.clear();
queued_segment_ = SIZE_MAX;
}
}
src/core/segment_weight_pipeline.h
+62
View File
@@ -0,0 +1,62 @@
#ifndef __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
#define __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <unordered_set>
#include <vector>
#include "ggml-backend.h"
struct DeviceMemoryRequest;
struct DeviceResidencyManager;
struct ggml_cgraph;
struct ggml_tensor;
namespace sd::ggml_graph_cut {
struct Plan;
}
namespace sd {
class SegmentWeightPipeline {
private:
std::weak_ptr<DeviceResidencyManager> residency_manager_;
ggml_backend_t compute_backend_ = nullptr;
uintptr_t owner_id_ = 0;
std::vector<std::vector<ggml_tensor*>> segment_params_;
std::vector<ggml_tensor*> queued_params_;
std::vector<ggml_tensor*> pinned_params_;
size_t queued_segment_ = SIZE_MAX;
bool enabled_ = true;
size_t next_parameter_segment(size_t segment_index) const;
std::vector<std::vector<ggml_tensor*>> preferred_eviction_order() const;
void disable();
void activate(size_t segment_index);
void clear();
public:
SegmentWeightPipeline(
const std::shared_ptr<DeviceResidencyManager>& residency_manager,
ggml_backend_t compute_backend,
uintptr_t owner_id,
ggml_cgraph* graph,
const ggml_graph_cut::Plan& plan,
const std::unordered_set<const ggml_tensor*>& params,
bool enabled = true);
~SegmentWeightPipeline();
const std::vector<ggml_tensor*>& params(size_t index) const { return segment_params_[index]; }
bool ensure_segment_capacity(size_t segment_index,
const std::vector<DeviceMemoryRequest>& requests);
bool segment_start(size_t segment_index, const std::function<bool()>& ensure_capacity);
void segment_end();
// Prefetch is best effort; segment_start falls back to synchronous loading.
void enqueue_next(size_t segment_index,
const DeviceMemoryRequest& request);
};
}
#endif // __SD_CORE_SEGMENT_WEIGHT_PIPELINE_H__
src/core/tensor_ggml.hpp
+22 -3
View File
@@ -9,6 +9,7 @@
#include <type_traits>
#include "core/tensor.hpp"
#include "ggml-backend.h"
#include "ggml.h"
namespace sd {
@@ -54,10 +55,28 @@ namespace sd {
GGML_ABORT("ggml tensor type does not match sd::Tensor type");
}
Tensor<T> result(shape_from_ggml(tensor));
if (tensor->buffer != nullptr) {
ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
std::vector<uint8_t> strided_data;
void* destination = result.data();
if (!ggml_is_contiguous(tensor)) {
strided_data.resize(ggml_nbytes(tensor));
destination = strided_data.data();
}
auto buffer = tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
if (buffer != nullptr) {
ggml_backend_tensor_get(tensor, destination, 0, ggml_nbytes(tensor));
} else {
std::memcpy(result.data(), tensor->data, ggml_nbytes(tensor));
std::memcpy(destination, tensor->data, ggml_nbytes(tensor));
}
if (!strided_data.empty()) {
for (int64_t i = 0; i < result.numel(); ++i) {
int64_t index = i;
size_t offset = 0;
for (int d = 0; d < GGML_MAX_DIMS; ++d) {
offset += static_cast<size_t>(index % tensor->ne[d]) * tensor->nb[d];
index /= tensor->ne[d];
}
std::memcpy(result.data() + i, strided_data.data() + offset, sizeof(T));
}
}
return result;
}