stable-diffusion.cpp/src/core/ggml_graph_cut.cpp

#include "core/ggml_graph_cut.h"

#include <algorithm>
#include <cctype>
#include <climits>
#include <cmath>
#include <cstring>
#include <map>
#include <set>
#include <sstream>
#include <stack>
#include <unordered_map>

#include "core/ggml_extend_backend.h"
#include "core/util.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"

#include "ggml/src/ggml-impl.h"

namespace sd::ggml_graph_cut {

    static constexpr double MAX_VRAM_BYTES_PER_GIB = 1024.0 * 1024.0 * 1024.0;

    static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) {
        if (tensor == nullptr) {
            return "<null>";
        }
        if (tensor->name[0] != '\0') {
            return tensor->name;
        }
        return sd_format("<tensor@%p>", (const void*)tensor);
    }

    static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
        GGML_ASSERT(gf != nullptr);
        GGML_ASSERT(tensor != nullptr);
        for (int i = 0; i < gf->n_leafs; ++i) {
            if (gf->leafs[i] == tensor) {
                return i;
            }
        }
        return -1;
    }

    static bool is_params_tensor(const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                                 const ggml_tensor* tensor) {
        if (tensor == nullptr) {
            return false;
        }
        return params_tensor_set.find(tensor) != params_tensor_set.end() ||
               (tensor->view_src != nullptr &&
                params_tensor_set.find(tensor->view_src) != params_tensor_set.end());
    }

    static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
        GGML_ASSERT(gf != nullptr);
        if (name == nullptr || name[0] == '\0') {
            return -1;
        }
        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 && std::strcmp(node->name, name) == 0) {
                return i;
            }
        }
        return -1;
    }

    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));
        });
        return value;
    }

    static std::string normalize_backend_budget_key(const std::string& value) {
        return lower_ascii_copy(trim(value));
    }

    static bool is_default_max_vram_key(const std::string& key) {
        std::string normalized = normalize_backend_budget_key(key);
        return normalized == "all" || normalized == "default" || normalized == "*";
    }

    static bool parse_max_vram_budget_value(const std::string& text, float* value, std::string* error) {
        float parsed = 0.f;
        if (!parse_strict_float(text, parsed) || !std::isfinite(parsed)) {
            if (error != nullptr) {
                *error = "invalid --max-vram value '" + text + "'";
            }
            return false;
        }
        *value = parsed;
        return true;
    }

    static std::vector<std::string> backend_budget_keys(ggml_backend_t backend) {
        std::vector<std::string> keys;
        if (backend == nullptr) {
            return keys;
        }

        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
        if (dev != nullptr) {
            keys.push_back(normalize_backend_budget_key(ggml_backend_dev_name(dev)));
        }
        const char* backend_name = ggml_backend_name(backend);
        if (backend_name != nullptr) {
            keys.push_back(normalize_backend_budget_key(backend_name));
        }
        return keys;
    }

    void MaxVramAssignment::reset(float fallback_gib) {
        default_gib = fallback_gib;
        backend_gib.clear();
        resolved_backend_bytes.clear();
    }

    bool MaxVramAssignment::parse(const std::string& raw_spec, std::string* error) {
        const std::string in = trim(raw_spec);
        if (in.empty()) {
            return true;
        }

        for (const std::string& raw_part : split_string(in, ',')) {
            const std::string part = trim(raw_part);
            if (part.empty()) {
                continue;
            }

            const size_t eq = part.find('=');
            if (eq == std::string::npos) {
                float value = 0.f;
                if (!parse_max_vram_budget_value(part, &value, error)) {
                    return false;
                }
                default_gib = value;
                continue;
            }

            const std::string key        = trim(part.substr(0, eq));
            const std::string value_text = trim(part.substr(eq + 1));
            if (key.empty() || value_text.empty()) {
                if (error != nullptr) {
                    *error = "invalid --max-vram assignment '" + part + "'";
                }
                return false;
            }

            float value = 0.f;
            if (!parse_max_vram_budget_value(value_text, &value, error)) {
                return false;
            }

            if (is_default_max_vram_key(key)) {
                default_gib = value;
                continue;
            }

            const std::string backend_key = trim(key);
            if (backend_key.empty()) {
                if (error != nullptr) {
                    *error = "invalid --max-vram backend key in '" + part + "'";
                }
                return false;
            }
            backend_gib[backend_key] = value;
        }
        resolved_backend_bytes.clear();
        return true;
    }

    bool MaxVramAssignment::canonicalize_backend_keys(std::string* error) {
        if (backend_gib.empty()) {
            return true;
        }

        std::unordered_map<std::string, float> normalized;
        for (const auto& kv : backend_gib) {
            std::string resolved = sd_backend_resolve_name(kv.first);
            if (resolved.empty()) {
                if (error != nullptr) {
                    *error = "unknown --max-vram backend '" + kv.first + "'";
                }
                return false;
            }
            normalized[normalize_backend_budget_key(resolved)] = kv.second;
        }
        backend_gib = std::move(normalized);
        resolved_backend_bytes.clear();
        return true;
    }

    size_t MaxVramAssignment::bytes_for_backend(ggml_backend_t backend) {
        std::vector<std::string> keys = backend_budget_keys(backend);
        const std::string cache_key   = keys.empty() ? std::string("<none>") : keys.front();
        auto cached                   = resolved_backend_bytes.find(cache_key);
        if (cached != resolved_backend_bytes.end()) {
            return cached->second;
        }

        float budget_gib = default_gib;
        if (!backend_gib.empty()) {
            for (const std::string& key : keys) {
                auto backend_it = backend_gib.find(key);
                if (backend_it != backend_gib.end()) {
                    budget_gib = backend_it->second;
                    break;
                }
            }
        }

        const float resolved_gib          = resolve_max_vram_gib(budget_gib, backend);
        const size_t bytes                = max_vram_gib_to_bytes(resolved_gib);
        resolved_backend_bytes[cache_key] = bytes;
        return bytes;
    }

    size_t max_vram_gib_to_bytes(float max_vram) {
        if (max_vram <= 0.f) {
            return 0;
        }
        return static_cast<size_t>(static_cast<double>(max_vram) * MAX_VRAM_BYTES_PER_GIB);
    }

    static float max_vram_bytes_to_gib(size_t max_vram_bytes) {
        return static_cast<float>(static_cast<double>(max_vram_bytes) / MAX_VRAM_BYTES_PER_GIB);
    }

    static size_t resolve_auto_max_vram_bytes(float spare_vram, ggml_backend_t backend) {
        if (backend == nullptr) {
            LOG_WARN("--max-vram < 0 requested, but no backend is available; disabling graph splitting");
            return 0;
        }

        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
        if (dev == nullptr) {
            LOG_WARN("--max-vram < 0 requested, but no backend device is available; disabling graph splitting");
            return 0;
        }
        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
            LOG_WARN("--max-vram < 0 requested, but the main backend is CPU; disabling graph splitting");
            return 0;
        }

        size_t free_vram  = 0;
        size_t total_vram = 0;
        ggml_backend_dev_memory(dev, &free_vram, &total_vram);
        size_t spare_bytes = static_cast<size_t>(MAX_VRAM_BYTES_PER_GIB * spare_vram);

        if (free_vram <= spare_bytes) {
            LOG_WARN("--max-vram < 0 requested, but free VRAM is %.2f GiB; reserving %.2f GiB leaves no graph budget",
                     free_vram / MAX_VRAM_BYTES_PER_GIB, spare_vram);
            return 0;
        }

        const size_t max_vram_bytes = free_vram - spare_bytes;
        LOG_INFO("--max-vram < 0 auto-detected %.2f GiB free VRAM (%.2f GiB total), reserving %.2f GiB; using %.2f GiB",
                 free_vram / MAX_VRAM_BYTES_PER_GIB,
                 total_vram / MAX_VRAM_BYTES_PER_GIB,
                 spare_vram,
                 max_vram_bytes / MAX_VRAM_BYTES_PER_GIB);
        return max_vram_bytes;
    }

    float resolve_max_vram_gib(float max_vram, ggml_backend_t backend) {
        if (max_vram >= 0.f) {
            return max_vram;
        }
        return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
    }

    static void build_segment(ggml_cgraph* gf,
                              Plan& plan,
                              Segment& segment,
                              const std::unordered_map<const ggml_tensor*, int>& producer_index,
                              std::unordered_set<int>& available_cut_output_node_indices,
                              ggml_backend_t backend,
                              const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                              const char* log_desc) {
        std::set<int> internal_nodes;
        std::unordered_set<const ggml_tensor*> input_seen;
        std::vector<Segment::InputRef> input_refs;

        std::stack<ggml_tensor*> work_stack;
        for (int output_node_index : segment.output_node_indices) {
            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
            if (output != nullptr) {
                work_stack.push(output);
            }
        }

        while (!work_stack.empty()) {
            ggml_tensor* tensor = work_stack.top();
            work_stack.pop();

            if (tensor == nullptr) {
                continue;
            }

            auto producer_it = producer_index.find(tensor);
            if (producer_it == producer_index.end()) {
                if (input_seen.insert(tensor).second) {
                    Segment::InputRef input_ref;
                    input_ref.type         = is_params_tensor(params_tensor_set, tensor) ? Segment::INPUT_PARAM : Segment::INPUT_EXTERNAL;
                    input_ref.display_name = graph_cut_tensor_display_name(tensor);
                    input_ref.leaf_index   = graph_leaf_index(gf, tensor);
                    input_refs.push_back(std::move(input_ref));
                }
                continue;
            }

            int node_idx = producer_it->second;
            if (available_cut_output_node_indices.find(node_idx) != available_cut_output_node_indices.end()) {
                if (input_seen.insert(tensor).second) {
                    Segment::InputRef input_ref;
                    input_ref.type         = Segment::INPUT_PREVIOUS_CUT;
                    input_ref.display_name = graph_cut_tensor_display_name(tensor);
                    input_ref.node_index   = node_idx;
                    input_refs.push_back(std::move(input_ref));
                }
                continue;
            }

            if (!internal_nodes.insert(node_idx).second) {
                continue;
            }

            ggml_tensor* node = ggml_graph_node(gf, node_idx);
            for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
                if (node->src[src_idx] != nullptr) {
                    work_stack.push(node->src[src_idx]);
                }
            }
        }

        if (!internal_nodes.empty()) {
            segment.internal_node_indices.assign(internal_nodes.begin(), internal_nodes.end());
        }

        std::sort(input_refs.begin(),
                  input_refs.end(),
                  [](const Segment::InputRef& a, const Segment::InputRef& b) {
                      if (a.type != b.type) {
                          return a.type < b.type;
                      }
                      return a.display_name < b.display_name;
                  });
        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) {
            available_cut_output_node_indices.insert(output_node_index);
        }
        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;
        }
        return starts_with(tensor->name, GGML_RUNNER_CUT_PREFIX) &&
               ends_with(tensor->name, GGML_RUNNER_CUT_SUFFIX);
    }

    std::string make_graph_cut_name(const std::string& group, const std::string& output) {
        return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output + GGML_RUNNER_CUT_SUFFIX;
    }

    void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
        if (tensor == nullptr) {
            return;
        }
        auto name = make_graph_cut_name(group, output);
        ggml_set_name(tensor, name.c_str());
    }

    int leaf_count(ggml_cgraph* gf) {
        GGML_ASSERT(gf != nullptr);
        return gf->n_leafs;
    }

    ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
        GGML_ASSERT(gf != nullptr);
        if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
            return nullptr;
        }
        return gf->leafs[leaf_index];
    }

    ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
        if (tensor == nullptr) {
            return nullptr;
        }
        return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
    }

    ggml_tensor* cache_source_tensor(ggml_tensor* tensor) {
        if (tensor == nullptr) {
            return nullptr;
        }
        if (tensor_buffer(tensor) == nullptr && tensor->src[0] != nullptr &&
            ggml_nelements(tensor->src[0]) == ggml_nelements(tensor) &&
            ggml_nbytes(tensor->src[0]) == ggml_nbytes(tensor)) {
            return cache_source_tensor(tensor->src[0]);
        }
        return tensor->view_src ? tensor->view_src : tensor;
    }

    size_t cache_tensor_bytes(const ggml_tensor* tensor) {
        if (tensor == nullptr) {
            return 0;
        }
        const ggml_tensor* cache_src = tensor->view_src ? tensor->view_src : tensor;
        return ggml_nbytes(cache_src);
    }

    static bool can_ignore_op_params(ggml_op op) {
        // Exempt only parameters that cannot affect graph layout or backend allocation size.
        switch (op) {
            case GGML_OP_SCALE:
                return true;
            default:
                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);
            }
        };
        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) {
                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));
            }
            if (!can_ignore_op_params(tensor->op)) {
                for (int value : tensor->op_params) {
                    signature.push_back(static_cast<uint32_t>(value));
                }
            }
        }
        return signature;
    }

    static bool plan_matches_graph(ggml_cgraph* gf,
                                   const Plan& plan,
                                   const std::vector<uint64_t>& layout) {
        GGML_ASSERT(gf != nullptr);
        if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
            plan.layout != layout) {
            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;
    }

    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
        GGML_ASSERT(gf != nullptr);
        return plan_matches_graph(gf, plan, graph_layout(gf, false));
    }

    ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
        GGML_ASSERT(gf != nullptr);
        if (output_index >= segment.output_node_indices.size()) {
            return nullptr;
        }
        int node_index = segment.output_node_indices[output_index];
        if (node_index < 0 || node_index >= ggml_graph_n_nodes(gf)) {
            return nullptr;
        }
        return ggml_graph_node(gf, node_index);
    }

    ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref) {
        GGML_ASSERT(gf != nullptr);
        if (input_ref.type == Segment::INPUT_PREVIOUS_CUT) {
            if (input_ref.node_index < 0 || input_ref.node_index >= ggml_graph_n_nodes(gf)) {
                return nullptr;
            }
            return ggml_graph_node(gf, input_ref.node_index);
        }
        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
            return nullptr;
        }
        return leaf_tensor(gf, input_ref.leaf_index);
    }

    std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment) {
        GGML_ASSERT(gf != nullptr);
        std::vector<ggml_tensor*> tensors;
        std::unordered_set<ggml_tensor*> seen_tensors;
        tensors.reserve(segment.input_refs.size());
        seen_tensors.reserve(segment.input_refs.size());
        for (const auto& input_ref : segment.input_refs) {
            if (input_ref.type != Segment::INPUT_PARAM) {
                continue;
            }
            ggml_tensor* tensor = input_tensor(gf, input_ref);
            if (tensor == nullptr) {
                continue;
            }
            if (seen_tensors.insert(tensor).second) {
                tensors.push_back(tensor);
            }
        }
        return tensors;
    }

    ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
                                     const Segment& segment,
                                     ggml_context** graph_ctx_out) {
        GGML_ASSERT(gf != nullptr);
        GGML_ASSERT(graph_ctx_out != nullptr);

        // Collect leaf inputs and internal nodes, then any tensor they
        // reference that is not already represented, notably the view_src of a
        // view-typed input leaf. ggml_gallocr sizes its hash set from
        // n_nodes + n_leafs (plus a 25% margin that rounds down to zero for a
        // one-node segment), so every distinct tensor it will hash must be
        // counted here or a tiny segment overflows the hash set and aborts.
        std::vector<ggml_tensor*> leaves;
        std::unordered_set<ggml_tensor*> represented;
        for (const auto& input : segment.input_refs) {
            ggml_tensor* current_input = input_tensor(gf, input);
            if (current_input == nullptr) {
                continue;
            }
            if (represented.insert(current_input).second) {
                leaves.push_back(current_input);
            }
        }
        for (int node_idx : segment.internal_node_indices) {
            represented.insert(ggml_graph_node(gf, node_idx));
        }
        auto add_reference = [&](ggml_tensor* tensor) {
            if (tensor != nullptr && represented.insert(tensor).second) {
                leaves.push_back(tensor);
            }
        };
        for (int node_idx : segment.internal_node_indices) {
            ggml_tensor* node = ggml_graph_node(gf, node_idx);
            for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
                add_reference(node->src[src_idx]);
            }
            add_reference(node->view_src);
        }
        for (size_t i = 0; i < leaves.size(); ++i) {
            add_reference(leaves[i]->view_src);
        }

        const size_t graph_size = segment.internal_node_indices.size() + leaves.size() + 8;
        ggml_init_params params = {
            /*.mem_size   =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
            /*.mem_buffer =*/nullptr,
            /*.no_alloc   =*/true,
        };
        ggml_context* graph_ctx = ggml_init(params);
        GGML_ASSERT(graph_ctx != nullptr);
        ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
        GGML_ASSERT(segment_graph != nullptr);

        for (ggml_tensor* leaf : leaves) {
            GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
            segment_graph->leafs[segment_graph->n_leafs++] = leaf;
        }

        for (int output_node_index : segment.output_node_indices) {
            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
            if (output == nullptr) {
                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));
        }
        *graph_ctx_out = graph_ctx;
        return segment_graph;
    }

    size_t measure_segment_compute_buffer(ggml_backend_t backend,
                                          ggml_cgraph* gf,
                                          const Segment& segment,
                                          const char* log_desc) {
        GGML_ASSERT(backend != nullptr);
        GGML_ASSERT(gf != nullptr);
        if (segment.internal_node_indices.empty()) {
            return 0;
        }

        struct TensorRuntimeBinding {
            ggml_backend_buffer_t buffer = nullptr;
            void* data                   = nullptr;
            void* extra                  = nullptr;
        };
        std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
        auto mark_measurement_external = [&](ggml_tensor* tensor) {
            if (tensor == nullptr) {
                return;
            }
            auto save_tensor = [&](ggml_tensor* t) {
                if (t == nullptr || saved_bindings.find(t) != saved_bindings.end()) {
                    return;
                }
                saved_bindings[t] = {t->buffer, t->data, t->extra};
                // During real execution params and previous-cut inputs already
                // have backend/cache buffers, so gallocr must not reserve them.
                t->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
            };
            save_tensor(tensor);
            save_tensor(tensor->view_src);
        };
        for (const auto& input : segment.input_refs) {
            if (input.type != Segment::INPUT_PARAM &&
                input.type != Segment::INPUT_PREVIOUS_CUT) {
                continue;
            }
            mark_measurement_external(input_tensor(gf, input));
        }

        std::unordered_map<ggml_tensor*, int32_t> saved_output_flags;
        for (int output_node_index : segment.output_node_indices) {
            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
            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;
        ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
        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,
            segment_graph,
            nullptr,
            nullptr,
            sizes);
        size_t buffer_size = sizes[0];

        ggml_gallocr_free(allocr);
        ggml_free(graph_ctx);
        for (const auto& kv : saved_output_flags) {
            kv.first->flags = kv.second;
        }
        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 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,
                    const char* log_desc) {
        GGML_ASSERT(backend != nullptr);
        GGML_ASSERT(gf != nullptr);
        Plan plan;
        plan.available    = true;
        const int n_nodes = ggml_graph_n_nodes(gf);
        if (n_nodes <= 0) {
            return plan;
        }
        plan.layout = graph_layout(gf, false);
        for (int i = 0; i < gf->n_leafs; ++i) {
            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) {
            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) {
            ggml_tensor* node = ggml_graph_node(gf, i);
            if (!is_graph_cut_tensor(node)) {
                continue;
            }

            plan.has_cuts = true;
            std::string full_name(node->name);
            size_t prefix_len   = std::strlen(GGML_RUNNER_CUT_PREFIX);
            size_t suffix_len   = std::strlen(GGML_RUNNER_CUT_SUFFIX);
            std::string payload = full_name.substr(prefix_len, full_name.size() - prefix_len - suffix_len);
            size_t sep          = payload.find('|');
            std::string group   = sep == std::string::npos ? payload : payload.substr(0, sep);

            auto it = group_to_segment.find(group);
            if (it == group_to_segment.end()) {
                Segment segment;
                segment.group_name = group;
                segment.output_node_indices.push_back(i);
                group_to_segment[group] = grouped_segments.size();
                grouped_segments.push_back(std::move(segment));
            } else {
                auto& segment = grouped_segments[it->second];
                segment.output_node_indices.push_back(i);
            }
        }

        if (!plan.has_cuts) {
            return plan;
        }

        std::unordered_set<int> available_cut_output_node_indices;
        available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
        for (auto& segment : grouped_segments) {
            build_segment(gf,
                          plan,
                          segment,
                          producer_index,
                          available_cut_output_node_indices,
                          backend,
                          params_tensor_set,
                          log_desc);
        }

        int final_output_index = graph_node_index_by_name(gf, "ggml_runner_final_result_tensor");
        if (final_output_index < 0) {
            final_output_index = n_nodes - 1;
        }
        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,
                          producer_index,
                          available_cut_output_node_indices,
                          backend,
                          params_tensor_set,
                          log_desc);
        }

        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;
                }
                segment->live_cut_names.insert(input.display_name);
                future_cut_names.insert(input.display_name);
            }
        }

        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,
                     plan_validation_error.c_str());
        }

        return plan;
    }

    Plan resolve_plan(ggml_backend_t backend,
                      ggml_cgraph* gf,
                      PlanCache* cache,
                      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);

        const auto layout = graph_layout(gf, false);
        auto& plans       = cache->graph_cut_plans;
        for (auto it = plans.begin(); it != plans.end(); ++it) {
            if (it->available && plan_matches_graph(gf, *it, layout)) {
                plans.splice(plans.begin(), plans, it);
                return plans.front();
            }
        }

        int64_t t_plan_begin = ggml_time_ms();
        plans.push_front(build_plan(backend, gf, params_tensor_set, log_desc));
        if (plans.size() > PlanCache::MAX_PLANS) {
            plans.pop_back();
        }
        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 plans.front();
    }

}  // namespace sd::ggml_graph_cut