#include "core/execution_engine.h"

#include <algorithm>
#include <chrono>
#include <exception>
#include <thread>
#include <unordered_set>

#include "core/ggml_extend_backend.h"
#include "core/ggml_tensor_utils.h"
#include "core/layer_split_partition.h"
#include "core/segment_graph_bindings.h"
#include "core/segment_weight_pipeline.h"
#include "core/tensor_ggml.hpp"
#include "core/util.h"

std::shared_ptr<sd::ExecutionEngine> DeviceResidencyManager::execution_engine() {
    if (execution_engine_ == nullptr) {
        execution_engine_ = std::make_shared<sd::ExecutionEngine>(this);
    }
    return execution_engine_;
}

namespace sd {
    class GraphExecution {
        ExecutionEngine& engine_;
        ExecutionSession& session_;
        const ComputeTask& task_;
        std::shared_ptr<DeviceResidencyManager> manager_;
        std::unordered_set<const ggml_tensor*> params_tensor_set_;
        ggml_status last_compute_status_ = GGML_STATUS_FAILED;
        Tensor<float> output_;

        using GraphCutPlan = ggml_graph_cut::Plan;
        ggml_tensor* canonical_param_tensor(ggml_tensor* tensor);
        std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* graph);
        ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const;
        // Weightless ops follow the most recent weight device. Views must stay
        // unpinned so the scheduler preserves cross-device copies for consumers.
        void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* graph, ggml_cgraph* original_graph = nullptr);
        Tensor<float> read_graph_tensor(ggml_tensor* tensor, const char* label);
        void copy_data_to_backend_tensor(ggml_cgraph* graph);
        const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* graph);
        bool assign_graph_cut_layer_split_backends(ggml_cgraph* graph);
        ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs = nullptr);
        std::vector<DeviceMemoryRequest> memory_requests(const std::vector<BackendBufferSize>& sizes, size_t pending_cache_bytes) const;
        bool fits(const std::vector<DeviceMemoryRequest>& requests, const std::vector<ggml_tensor*>& params) const;
        bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute);
        bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs);
        bool is_multi_device() const { return !session_.extra_backends_.empty(); }
        size_t retained_runtime_buffer_bytes(ggml_backend_t backend) const { return session_.resident_bytes(backend); }
        void sync_runtime_residency() { engine_.sync_residency(session_); }

    public:
        GraphExecution(ExecutionEngine& engine, ExecutionSession& session, const ComputeTask& task)
            : engine_(engine), session_(session), task_(task), manager_(session.residency_manager_.lock()) {
            params_tensor_set_.insert(task.weights.begin(), task.weights.end());
            if (manager_ != nullptr) {
                for (int i = 0; i < ggml_graph_cut::leaf_count(task.graph); ++i) {
                    auto* parameter = manager_->resolve_param_tensor(ggml_graph_cut::leaf_tensor(task.graph, i));
                    if (parameter != nullptr) {
                        params_tensor_set_.insert(parameter);
                    }
                }
                for (auto* tensor : params_tensor_set_) {
                    auto* parameter = manager_->resolve_param_tensor(const_cast<ggml_tensor*>(tensor));
                    if (parameter != nullptr) {
                        session_.params_.insert(parameter);
                    }
                }
            }
        }

        ComputeResult run() {
            if (!execute_graph(task_.graph, task_.n_threads, task_.no_return, task_.read_outputs)) {
                return {last_compute_status_, {}};
            }
            return {GGML_STATUS_SUCCESS, std::move(output_)};
        }
    };

    static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) {
        // Vulkan reports a preferred allocation chunk size, not a hard tensor limit.
        if (sd_backend_is(backend, "Vulkan")) {
            return true;
        }
        auto buft        = ggml_backend_get_default_buffer_type(backend);
        size_t alignment = ggml_backend_buft_get_alignment(buft);
        size_t max_size  = ggml_backend_buft_get_max_size(buft);
        max_size -= max_size % alignment;
        if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) {
            return false;
        }
        // Cross-backend copies must also fit the destination's allocation limit.
        for (auto src : node->src) {
            if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) {
                return false;
            }
        }
        return true;
    }

    static size_t add_bytes(size_t a, size_t b) {
        return b > SIZE_MAX - a ? SIZE_MAX : a + b;
    }

    ggml_tensor* GraphExecution::canonical_param_tensor(ggml_tensor* tensor) {
        for (auto* current = tensor; current != nullptr; current = current->view_src) {
            if (params_tensor_set_.count(current) != 0)
                return current;
        }
        return nullptr;
    }

    std::vector<ggml_tensor*> GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) {
        std::vector<ggml_tensor*> used_params;
        if (gf == nullptr || params_tensor_set_.empty()) {
            return used_params;
        }

        std::unordered_set<const ggml_tensor*> seen_params;
        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
        seen_params.reserve(static_cast<size_t>(n_leafs));
        for (int i = 0; i < n_leafs; ++i) {
            ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
            ggml_tensor* param_leaf = canonical_param_tensor(leaf);
            if (param_leaf != nullptr &&
                seen_params.insert(param_leaf).second) {
                used_params.push_back(param_leaf);
            }
        }
        return used_params;
    }

    ggml_backend_t GraphExecution::backend_for_weight(const ggml_tensor* tensor) const {
        if (tensor == nullptr || tensor->buffer == nullptr) {
            return nullptr;
        }
        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
            ggml_backend_buffer_is_host(tensor->buffer)) {
            return nullptr;
        }
        ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
        if (dev == nullptr) {
            return nullptr;
        }
        if (ggml_backend_get_device(session_.backend_) == dev) {
            return session_.backend_;
        }
        for (ggml_backend_t backend : session_.extra_backends_) {
            if (ggml_backend_get_device(backend) == dev) {
                return backend;
            }
        }
        return nullptr;
    }

    void GraphExecution::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) {
        if (sched == nullptr || gf == nullptr) {
            return;
        }
        ggml_backend_t current = session_.backend_;
        const int n_nodes      = ggml_graph_n_nodes(gf);
        for (int i = 0; i < n_nodes; i++) {
            ggml_tensor* node    = ggml_graph_node(gf, i);
            auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i));
            if (node_assignment != session_.node_assignments_.end()) {
                current = node_assignment->second;
            }
            for (int s = 0; s < GGML_MAX_SRC; s++) {
                ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
                if (weight_backend != nullptr) {
                    if (node_assignment == session_.node_assignments_.end()) {
                        current = weight_backend;
                    }
                }
            }
            if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
                node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
                continue;
            }
            if (!backend_can_allocate_op(current, node)) {
                auto cpu = session_.workspace_.cpu_backend();
                if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) {
                    ggml_backend_sched_set_tensor_backend(sched, node, cpu);
                }
            } else if (ggml_backend_supports_op(current, node)) {
                ggml_backend_sched_set_tensor_backend(sched, node, current);
            }
        }
    }

    Tensor<float> GraphExecution::read_graph_tensor(ggml_tensor* tensor, const char* label) {
        if (tensor == nullptr) {
            LOG_ERROR("%s %s tensor is null", task_.description.c_str(), label);
            return {};
        }
        if (tensor->type != GGML_TYPE_F32) {
            LOG_ERROR("%s %s tensor type mismatch: got %s",
                      task_.description.c_str(),
                      label,
                      ggml_type_name(tensor->type));
            return {};
        }
        ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
        if (buf == nullptr) {
            LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
                      task_.description.c_str(),
                      label,
                      tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
                      ggml_op_name(tensor->op),
                      tensor->buffer,
                      tensor->view_src,
                      tensor->view_src ? tensor->view_src->buffer : nullptr,
                      tensor->data);
            return {};
        }

        return sd::make_sd_tensor_from_ggml<float>(tensor);
    }

    void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) {
        GGML_ASSERT(gf != nullptr);
        std::unordered_set<const ggml_tensor*> graph_tensor_set;
        const int n_leafs = sd::ggml_graph_cut::leaf_count(gf);
        const int n_nodes = ggml_graph_n_nodes(gf);
        graph_tensor_set.reserve(static_cast<size_t>(n_leafs + n_nodes));
        for (int i = 0; i < n_leafs; ++i) {
            graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
        }
        for (int i = 0; i < n_nodes; ++i) {
            graph_tensor_set.insert(ggml_graph_node(gf, i));
        }

        for (auto& kv : task_.inputs) {
            auto tensor = kv.first;
            auto data   = kv.second;
            if (tensor == nullptr || data == nullptr) {
                continue;
            }
            const char* name = ggml_get_name(tensor);
            if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
                continue;
            }
            if (tensor->buffer == nullptr) {
                LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
                         task_.description.c_str(),
                         name != nullptr ? name : "",
                         (long long)tensor->ne[0],
                         (long long)tensor->ne[1],
                         (long long)tensor->ne[2],
                         (long long)tensor->ne[3],
                         ggml_type_name(tensor->type));
                continue;
            }

            ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
            if (buf == nullptr) {
                LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
                         task_.description.c_str(),
                         tensor && tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
                         tensor ? ggml_op_name(tensor->op) : "<null>",
                         data,
                         tensor ? tensor->view_src : nullptr,
                         (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr);
                continue;
            }

            ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor));
        }
    }

    const GraphExecution::GraphCutPlan& GraphExecution::resolve_graph_cut_plan(ggml_cgraph* gf) {
        GGML_ASSERT(gf != nullptr);
        return sd::ggml_graph_cut::resolve_plan(session_.backend_,
                                                gf,
                                                &session_.plan_cache_,
                                                params_tensor_set_,
                                                task_.description.c_str());
    }

    bool GraphExecution::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
        session_.node_assignments_.clear();
        if (!session_.layer_split_enabled_) {
            return true;
        }
        if (!is_multi_device()) {
            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", task_.description.c_str());
            return false;
        }

        const auto& plan = resolve_graph_cut_plan(gf);
        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
            auto manager = manager_;
            if (manager == nullptr) {
                LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str());
                return false;
            }
            std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
            if (!graph_params.empty() &&
                !manager->assign_compute_backend(graph_params, session_.backend_)) {
                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
                          task_.description.c_str(),
                          sd::layer_split_backend_device_display_name(session_.backend_).c_str());
                return false;
            }
            for (ggml_tensor* param : graph_params) {
                if (param != nullptr) {
                    session_.weight_assignments_[param] = session_.backend_;
                }
            }
            const int n_nodes = ggml_graph_n_nodes(gf);
            for (int i = 0; i < n_nodes; i++) {
                ggml_tensor* node = ggml_graph_node(gf, i);
                if (node != nullptr) {
                    session_.node_assignments_[node] = session_.backend_;
                }
            }
            if (!session_.primary_notice_logged_) {
                LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
                         task_.description.c_str(),
                         sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
                         graph_params.size());
                session_.primary_notice_logged_ = true;
            } else {
                LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
                            task_.description.c_str(),
                            sd::layer_split_backend_device_display_name(session_.backend_).c_str(),
                            graph_params.size());
            }
            return true;
        }

        std::vector<ggml_backend_t> split_backends;
        split_backends.reserve(session_.extra_backends_.size() + 1);
        split_backends.push_back(session_.backend_);
        for (ggml_backend_t backend : session_.extra_backends_) {
            if (backend != nullptr) {
                split_backends.push_back(backend);
            }
        }

        auto manager = manager_;
        if (manager == nullptr) {
            LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str());
            return false;
        }

        sd::GraphCutLayerSplitAssignment assignment;
        auto canonicalize_param = [this](ggml_tensor* tensor) {
            return canonical_param_tensor(tensor);
        };
        if (!sd::partition_graph_cut_layer_split(task_.description.c_str(),
                                                 gf,
                                                 plan,
                                                 split_backends,
                                                 session_.backend_limits_,
                                                 session_.max_backend_bytes_,
                                                 session_.weight_assignments_,
                                                 canonicalize_param,
                                                 &assignment)) {
            return false;
        }

        for (size_t i = 0; i < split_backends.size(); i++) {
            if (assignment.tensors_by_backend[i].empty()) {
                continue;
            }
            if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
                LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
                          task_.description.c_str(),
                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
                return false;
            }
        }

        session_.node_assignments_ = std::move(assignment.node_assignments);
        sd::log_graph_cut_layer_split_assignment(task_.description.c_str(), split_backends, assignment);

        return true;
    }

    ComputeWorkspace::Measurement GraphExecution::measure(ggml_cgraph* graph, const std::unordered_set<const ggml_tensor*>* external_outputs) {
        auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t {
            if (external_outputs != nullptr && external_outputs->count(tensor)) {
                return session_.backend_;
            }
            if (!params_tensor_set_.count(tensor)) {
                return nullptr;
            }
            auto placement = session_.weight_assignments_.find(tensor);
            return placement == session_.weight_assignments_.end() ? session_.backend_ : placement->second;
        };
        auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
            pin_multi_device_nodes(scheduler, copy, graph);
        };
        return session_.workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
    }

    std::vector<DeviceMemoryRequest> GraphExecution::memory_requests(
        const std::vector<BackendBufferSize>& sizes,
        size_t pending_cache_bytes) const {
        std::vector<DeviceMemoryRequest> requests;
        for (const auto& size : sizes) {
            const size_t retained    = retained_runtime_buffer_bytes(size.backend);
            const size_t reusable    = session_.workspace_.bytes(size.backend);
            const size_t cache_bytes = size.backend == session_.backend_ ? pending_cache_bytes : 0;
            const size_t pending     = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes);
            size_t limit             = session_.max_backend_bytes_;
            if (is_multi_device()) {
                size_t index = 0;
                if (size.backend != session_.backend_) {
                    auto position = std::find(session_.extra_backends_.begin(), session_.extra_backends_.end(), size.backend);
                    index         = static_cast<size_t>(position - session_.extra_backends_.begin()) + 1;
                }
                if (index < session_.backend_limits_.size()) {
                    limit = session_.backend_limits_[index];
                }
            }
            requests.push_back({size.backend, session_.owner_id(), pending,
                                retained, limit});
        }
        return requests;
    }

    bool GraphExecution::fits(const std::vector<DeviceMemoryRequest>& requests,
                              const std::vector<ggml_tensor*>& params) const {
        auto manager = manager_;
        if (manager == nullptr) {
            return params.empty();
        }
        for (const auto& request : requests) {
            if (!manager->fits_compute_backend_capacity(request, params)) {
                return false;
            }
        }
        return true;
    }

    bool GraphExecution::execute_segment(ggml_cgraph* graph, int n_threads, const std::function<void()>& on_before_compute) {
        if (sd_backend_is_cpu(session_.backend_)) {
            sd_backend_cpu_set_n_threads(session_.backend_, n_threads);
        }
        if (session_.workspace_.cpu_backend() != nullptr) {
            sd_backend_cpu_set_n_threads(session_.workspace_.cpu_backend(), n_threads);
        }
        auto scheduler = session_.workspace_.scheduler();
        GGML_ASSERT(scheduler != nullptr);
        ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data());
#ifndef SD_USE_UPSTREAM_GGML
        struct BeforeComputeCallback {
            ggml_backend_t backend;
            const std::function<void()>& callback;
            bool invoked = false;
        } before_compute{session_.backend_, on_before_compute};
        ggml_backend_sched_set_compute_callback(
            scheduler, [](ggml_backend_t backend, void* data) {
                auto& hook = *static_cast<BeforeComputeCallback*>(data);
                if (backend == hook.backend && !hook.invoked) {
                    hook.invoked = true;
                    hook.callback();
                }
            },
            &before_compute);
        const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
        ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr);
#else
        on_before_compute();
        const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph);
#endif
        session_.workspace_.synchronize();
        if (status != GGML_STATUS_SUCCESS) {
            last_compute_status_ = status;
            LOG_ERROR("%s compute failed: %s", task_.description.c_str(), ggml_status_to_string(status));
            return false;
        }
        const std::string description = task_.description;
        if (!task_.debug_tensors.empty()) {
            std::unordered_set<const ggml_tensor*> graph_tensors;
            const int leaf_count = ggml_graph_cut::leaf_count(graph);
            const int node_count = ggml_graph_n_nodes(graph);
            graph_tensors.reserve(static_cast<size_t>(leaf_count + node_count));
            for (int index = 0; index < leaf_count; ++index) {
                graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index));
            }
            for (int index = 0; index < node_count; ++index) {
                graph_tensors.insert(ggml_graph_node(graph, index));
            }

            for (const auto& entry : task_.debug_tensors) {
                ggml_tensor* tensor = entry.first;
                if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) {
                    continue;
                }
                ggml_backend_buffer_t buffer =
                    tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer;
                if (buffer == nullptr) {
                    LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
                             description.c_str(),
                             entry.second.c_str());
                    continue;
                }
                if (tensor->type != GGML_TYPE_F32) {
                    LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
                             description.c_str(),
                             entry.second.c_str(),
                             ggml_type_name(tensor->type));
                    continue;
                }
                auto debug_tensor = make_sd_tensor_from_ggml<float>(tensor);
                print_sd_tensor(debug_tensor, false, entry.second.c_str());
            }
        }

        return true;
    }

    bool GraphExecution::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function<bool()>& read_outputs) {
        if (!assign_graph_cut_layer_split_backends(graph)) {
            return false;
        }
        const auto params       = collect_used_param_tensors(graph);
        const auto& cached_plan = resolve_graph_cut_plan(graph);
        auto manager            = manager_;
        const bool can_segment  = !is_multi_device() && !sd_backend_is_cpu(session_.backend_) &&
                                  manager != nullptr && manager->segmented_compute_enabled() &&
                                  cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1;
        // An empty workspace request is a lower bound on the full graph's memory demand.
        bool segmented = can_segment && !fits(memory_requests({{session_.backend_, 0}}, 0), params);
        ComputeWorkspace::Measurement full_measurement;
        if (!segmented) {
            full_measurement = measure(graph);
            if (full_measurement.buffers.empty()) {
                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                return false;
            }
            if (can_segment) {
                // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
                constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
                auto requests                      = memory_requests(full_measurement.buffers, session_.cache_.pending_bytes(graph));
                for (auto& request : requests) {
                    request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
                }
                segmented = !fits(requests, params);
            }
        }
        ggml_graph_cut::Plan monolithic_plan;
        if (!segmented) {
            monolithic_plan.segments.emplace_back();
            auto& segment               = monolithic_plan.segments.back();
            segment.group_name          = "graph";
            segment.compute_buffer_size = cached_plan.compute_buffer_size;
            segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph));
            segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph));
            for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
                segment.internal_node_indices.push_back(i);
            }
            for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) {
                auto tensor = ggml_graph_cut::leaf_tensor(graph, i);
                ggml_graph_cut::Segment::InputRef input;
                input.leaf_index = i;
                input.type       = canonical_param_tensor(tensor) != nullptr
                                       ? ggml_graph_cut::Segment::INPUT_PARAM
                                       : ggml_graph_cut::Segment::INPUT_EXTERNAL;
                segment.input_refs.push_back(input);
            }
        }
        const auto& plan            = segmented ? cached_plan : monolithic_plan;
        const bool segments_changed = plan.segments.size() != session_.logged_segment_count_;
        if (segments_changed && (segmented || session_.logged_segment_count_ > 1)) {
            LOG_VERBOSE("%s using %zu segment%s", task_.description.c_str(),
                        plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
        }
        SegmentGraphBindings bindings(session_.cut_cache_, plan, graph);
        SegmentWeightPipeline weights(manager, session_.backend_, session_.owner_id(),
                                      graph, plan, params_tensor_set_,
                                      segmented && manager != nullptr && manager->prefetch_enabled());

        std::map<ggml_backend_t, size_t> peak_compute_bytes;
        auto track_compute_buffer = [&](ggml_backend_t backend) {
            if (backend != nullptr) {
                auto& peak = peak_compute_bytes[backend];
                peak       = std::max(peak, session_.workspace_.bytes(backend));
            }
        };
        for (size_t index = 0; index < plan.segments.size(); ++index) {
            const auto& segment = plan.segments[index];
            const bool last     = index + 1 == plan.segments.size();
            auto fail_segment   = [&](const char* phase) {
                LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", task_.description.c_str(),
                          index + 1, plan.segments.size(), segment.group_name.c_str(), phase);
                return false;
            };
            session_.cut_cache_.prune(segment.live_cut_names);
            bindings.reset(segment);
            if (!bindings.bind_cached_inputs(segment, task_.description.c_str())) {
                return fail_segment("input binding");
            }
            ggml_context* segment_context = nullptr;
            auto segment_graph            = segmented
                                                ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context)
                                                : graph;
            struct SegmentCleanup {
                GraphExecution& runner;
                SegmentWeightPipeline& weights;
                SegmentGraphBindings& bindings;
                ggml_context* context;
                ~SegmentCleanup() {
                    if (!runner.session_.workspace_.segment_end()) {
                        runner.last_compute_status_ = GGML_STATUS_FAILED;
                    }
                    bindings.restore();
                    weights.segment_end();
                    ggml_free(context);
                    runner.sync_runtime_residency();
                }
            } segment_cleanup{*this, weights, bindings, segment_context};

            std::unordered_set<const ggml_tensor*> cache_outputs;
            if (segmented) {
                for (int output_index : segment.output_node_indices) {
                    auto tensor = ggml_graph_node(graph, output_index);
                    if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) &&
                        GraphCutTensorCache::can_bind_output(session_.backend_, tensor)) {
                        cache_outputs.insert(tensor);
                    }
                }
            }
            auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement;
            if (!session_.workspace_.prepare(measurement)) {
                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                return fail_segment("workspace preparation");
            }
            const size_t cut_bytes = last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment);
            size_t new_cache_bytes = add_bytes(cut_bytes, session_.cache_.pending_bytes(segment_graph));
            auto ensure_capacity   = [&]() {
                sync_runtime_residency();
                auto requests = memory_requests(measurement.buffers, new_cache_bytes);
                if (fits(requests, weights.params(index))) {
                    return true;
                }
                if (session_.workspace_.release_excess(measurement)) {
                    sync_runtime_residency();
                    requests = memory_requests(measurement.buffers, new_cache_bytes);
                }
                const bool ready = weights.ensure_segment_capacity(index, requests);
                if (!ready && manager != nullptr) {
                    last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                }
                return ready;
            };
            if (!weights.segment_start(index, ensure_capacity)) {
                return fail_segment("weight preparation");
            }
            // Preparing weights can execute LoRA graphs and reclaim an idle workspace.
            if (!session_.workspace_.prepare(measurement)) {
                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                return fail_segment("workspace preparation");
            }
            if (!ensure_capacity()) {
                return fail_segment("workspace capacity check");
            }
            if (!session_.cut_cache_.bind_outputs(cache_outputs)) {
                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                return fail_segment("output cache allocation");
            }
            new_cache_bytes = add_bytes(last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph));
            if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
                    pin_multi_device_nodes(scheduler, current);
                })) {
                last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
                return fail_segment("workspace allocation");
            }
            for (const auto& size : measurement.buffers) {
                track_compute_buffer(size.backend);
            }
            if (session_.workspace_.scheduler() != nullptr) {
                track_compute_buffer(session_.workspace_.cpu_backend());
            }
            if (!ensure_capacity()) {
                return fail_segment("allocated capacity check");
            }
            copy_data_to_backend_tensor(segment_graph);
            auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes);
            if (!prefetch_requests.empty()) {
                weights.enqueue_next(index, prefetch_requests.front());
            }
            LOG_DEBUG("%s executing segment %zu/%zu: %s", task_.description.c_str(),
                      index + 1, plan.segments.size(), segment.group_name.c_str());
            if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) {
                return fail_segment("execution");
            }
            auto cache_status = session_.cache_.capture(segment_graph);
            if (cache_status == GGML_STATUS_SUCCESS) {
                cache_status = session_.cut_cache_.capture(graph, segment, task_.description.c_str());
            }
            if (cache_status != GGML_STATUS_SUCCESS) {
                last_compute_status_ = cache_status;
                return fail_segment("output caching");
            }
            sync_runtime_residency();
            if (last) {
                if (read_outputs && !read_outputs()) {
                    return fail_segment("output finalization");
                }
                if (!no_return) {
                    output_ = read_graph_tensor(task_.output, "output");
                    if (output_.empty()) {
                        return fail_segment("output readback");
                    }
                }
            }
            if (!session_.workspace_.segment_end()) {
                last_compute_status_ = GGML_STATUS_FAILED;
                return fail_segment("workspace synchronization");
            }
            // Final outputs and their callbacks may still be views of consumed cuts.
            session_.cut_cache_.prune(segment.future_cut_names);
        }
        if (segments_changed || peak_compute_bytes != session_.logged_compute_bytes_) {
            for (const auto& entry : peak_compute_bytes) {
                LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)",
                            task_.description.c_str(), entry.second / (1024.0 * 1024.0),
                            sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first),
                            plan.segments.size(), plan.segments.size() == 1 ? "" : "s");
            }
            session_.logged_compute_bytes_ = std::move(peak_compute_bytes);
            session_.logged_segment_count_ = plan.segments.size();
        }
        return true;
    }

    ExecutionSession::ExecutionSession(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager)
        : backend_(backend), residency_manager_(manager), cache_(backend), cut_cache_(backend), workspace_(backend) {
        GGML_ASSERT(backend_ != nullptr);
    }

    size_t ExecutionSession::resident_bytes(ggml_backend_t backend) const {
        size_t bytes = workspace_.bytes(backend);
        if (backend == backend_) {
            const auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; };
            bytes          = add(bytes, cache_.resident_bytes(ggml_backend_get_device(backend)));
            bytes          = add(bytes, cut_cache_.resident_bytes(ggml_backend_get_device(backend)));
        }
        return bytes;
    }

    void ExecutionSession::set_layer_split_enabled(bool enabled) {
        layer_split_enabled_ = enabled;
        if (!enabled) {
            weight_assignments_.clear();
            node_assignments_.clear();
            primary_notice_logged_ = false;
        }
    }

    void ExecutionSession::set_backend_limits(const std::vector<size_t>& limits) {
        backend_limits_ = limits;
        weight_assignments_.clear();
        node_assignments_.clear();
        primary_notice_logged_ = false;
    }

    void ExecutionSession::set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
        extra_backends_.clear();
        for (auto backend : backends) {
            if (backend != nullptr && backend != backend_ &&
                std::find(extra_backends_.begin(), extra_backends_.end(), backend) == extra_backends_.end()) {
                extra_backends_.push_back(backend);
            }
        }
        workspace_.set_extra_backends(extra_backends_);
        weight_assignments_.clear();
        node_assignments_.clear();
        primary_notice_logged_ = false;
    }

    ExecutionEngine::~ExecutionEngine() {
        for (auto& entry : sessions_) {
            end(*entry.second);
        }
    }

    ExecutionSession& ExecutionEngine::create_session(ggml_backend_t backend, const std::shared_ptr<DeviceResidencyManager>& manager) {
        GGML_ASSERT(manager.get() == manager_);
        auto session = std::make_unique<ExecutionSession>(backend, manager);
        auto& result = *session;
        sessions_.emplace(session->owner_id(), std::move(session));
        return result;
    }

    void ExecutionEngine::release_session(ExecutionSession& session) {
        end(session);
        sessions_.erase(session.owner_id());
    }

    bool ExecutionEngine::start(ExecutionSession& session) {
        if (session.started_) {
            return true;
        }
        sd_backend_set_buffer_cache_enabled(session.backend_, true);
        for (auto backend : session.extra_backends_) {
            sd_backend_set_buffer_cache_enabled(backend, true);
        }
        session.workspace_.set_extra_backends(session.extra_backends_);
        if (auto manager = session.residency_manager_.lock()) {
            manager->register_runtime_owner(session.owner_id());
        }
        session.started_ = true;
        return true;
    }

    void ExecutionEngine::end(ExecutionSession& session) {
        GGML_ASSERT(!session.executing_);
        if (!session.started_) {
            return;
        }
        session.workspace_.release();
        session.cache_.clear();
        session.cut_cache_.clear();
        session.logged_compute_bytes_.clear();
        session.logged_segment_count_ = 0;
        if (auto manager = session.residency_manager_.lock()) {
            manager->clear_prefetched_params(session.owner_id());
            std::vector<ggml_tensor*> params(session.params_.begin(), session.params_.end());
            manager->evict_compute_backend_params(params);
            manager->remove_runtime_owner(session.owner_id());
        }
        session.params_.clear();
        sd_backend_clear_buffer_cache(session.backend_);
        for (auto backend : session.extra_backends_) {
            sd_backend_clear_buffer_cache(backend);
        }
        session.started_ = false;
    }

    void ExecutionEngine::sync_residency(ExecutionSession& session) {
        if (auto manager = session.residency_manager_.lock()) {
            manager->update_runtime_residency(session.owner_id(), session.backend_, session.resident_bytes(session.backend_));
            for (auto backend : session.extra_backends_) {
                manager->update_runtime_residency(session.owner_id(), backend, session.resident_bytes(backend));
            }
        }
    }

    bool ExecutionEngine::ensure_capacity(
        DeviceResidencyManager& manager,
        const DeviceMemoryRequest& request,
        const std::vector<ggml_tensor*>& required_params,
        const std::vector<std::vector<ggml_tensor*>>& preferred_eviction_order,
        const std::vector<ggml_tensor*>& protected_params) {
        GGML_ASSERT(&manager == manager_);
        auto capacity      = [&]() { return manager.inspect_compute_backend_capacity(request, required_params); };
        const auto initial = capacity();
        if (!initial.valid) {
            return false;
        }
        if (initial.fits()) {
            return true;
        }
        for (auto& entry : sessions_) {
            auto& session = *entry.second;
            if (session.owner_id() != request.owner_id && session.workspace_.release()) {
                sync_residency(session);
                if (capacity().fits(true)) {
                    return true;
                }
            }
        }
        auto protected_weights = protected_params;
        protected_weights.insert(protected_weights.end(), required_params.begin(), required_params.end());
        auto candidates = manager.evictable_compute_backend_blocks(request.compute_backend, protected_weights);
        std::unordered_map<const ggml_tensor*, size_t> preferred_positions;
        for (size_t index = 0; index < preferred_eviction_order.size(); ++index) {
            for (auto* tensor : preferred_eviction_order[index]) {
                auto* weight = manager.resolve_param_tensor(tensor);
                preferred_positions.emplace(weight != nullptr ? weight : tensor, index);
            }
        }
        auto preferred_position = [&](const WeightResidencyBlock& block) {
            size_t position = 0;
            // A shared allocation follows the last of its members in the eviction order.
            for (auto* tensor : block.tensors) {
                const auto found = preferred_positions.find(tensor);
                if (found == preferred_positions.end()) {
                    return SIZE_MAX;
                }
                position = std::max(position, found->second);
            }
            return position;
        };
        std::stable_sort(candidates.begin(), candidates.end(), [&](const auto& lhs, const auto& rhs) {
            if (lhs.requires_demotion != rhs.requires_demotion) {
                return !lhs.requires_demotion;
            }
            if (!lhs.requires_demotion) {
                const auto left = preferred_position(lhs), right = preferred_position(rhs);
                if (left != right) {
                    return left < right;
                }
                if (left != SIZE_MAX) {
                    return false;
                }
            }
            return lhs.last_use_epoch < rhs.last_use_epoch;
        });
        for (const auto& block : candidates) {
            if (manager.evict_compute_backend_block(block) && capacity().fits(true)) {
                return true;
            }
        }
        auto available = capacity();
        if (available.fits(true)) {
            return true;
        }
        if (sd_backend_is(request.compute_backend, "Vulkan") &&
            available.required_budget_bytes <= available.available_budget_bytes) {
            ggml_backend_synchronize(request.compute_backend);
            sd_backend_clear_buffer_cache(request.compute_backend);
            const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100);
            do {
                available = capacity();
                if (available.fits(true)) {
                    return true;
                }
                if (available.required_budget_bytes > available.available_budget_bytes) {
                    break;
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(10));
            } while (std::chrono::steady_clock::now() < deadline);
        }
        available               = manager.inspect_compute_backend_capacity(request, required_params, true);
        const auto device_bytes = available.available_device_bytes == SIZE_MAX
                                      ? std::string("unknown")
                                      : sd_format("%.2f MB", available.available_device_bytes / (1024.0 * 1024.0));
        const auto budget_bytes = available.available_budget_bytes == SIZE_MAX
                                      ? std::string("unlimited")
                                      : sd_format("%.2f MB", available.available_budget_bytes / (1024.0 * 1024.0));
        LOG_WARN("execution engine cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget",
                 ggml_backend_name(request.compute_backend),
                 available.required_device_bytes / (1024.0 * 1024.0),
                 available.required_budget_bytes / (1024.0 * 1024.0), device_bytes.c_str(), budget_bytes.c_str());
        return false;
    }

    void ExecutionEngine::clear_cache(ExecutionSession& session) {
        session.cache_.clear();
        sync_residency(session);
    }

    void ExecutionEngine::discard_graph(ExecutionSession& session) {
        GGML_ASSERT(!session.executing_);
        session.cache_.graph_end(false);
        session.cut_cache_.clear();
        session.workspace_.release();
        sync_residency(session);
    }

    ComputeResult ExecutionEngine::execute(ExecutionSession& session, const ComputeTask& task) {
        if (session.executing_ || task.graph == nullptr || !start(session)) {
            return {};
        }
        session.executing_ = true;
        ComputeResult result;
        try {
            GraphExecution execution(*this, session, task);
            result = execution.run();
        } catch (const std::bad_alloc&) {
            result.status = GGML_STATUS_ALLOC_FAILED;
            LOG_ERROR("%s graph allocation failed", task.description.c_str());
        } catch (const std::exception& error) {
            LOG_ERROR("%s graph execution failed on %s: %s", task.description.c_str(),
                      ggml_backend_name(session.backend_), error.what());
        } catch (...) {
            LOG_ERROR("%s graph execution failed on %s: unknown exception", task.description.c_str(),
                      ggml_backend_name(session.backend_));
        }
        if (!session.workspace_.segment_end()) {
            result.status = GGML_STATUS_FAILED;
        }
        const bool success = result.status == GGML_STATUS_SUCCESS;
        session.cache_.graph_end(success);
        session.cut_cache_.clear();
        if (!success) {
            result.output = {};
            session.workspace_.release();
        }
        session.executing_ = false;
        sync_residency(session);
        return result;
    }
}