stable-diffusion.cpp/src/device_residency_manager.h

#ifndef __DEVICE_RESIDENCY_MANAGER_H__
#define __DEVICE_RESIDENCY_MANAGER_H__

#include <cstdint>
#include <functional>
#include <memory>
#include <vector>

#include "ggml-backend.h"

struct ggml_tensor;

namespace sd {
    class ExecutionEngine;
}

enum class WeightPrefetchResult {
    Scheduled,
    AlreadyResident,
    Unsupported,
    Failed,
};

struct WeightResidencyInfo {
    bool async_prefetch_supported = false;
    size_t missing_bytes          = 0;
};

struct DeviceMemoryRequest {
    ggml_backend_t compute_backend  = nullptr;
    uintptr_t owner_id              = 0;
    size_t pending_allocation_bytes = 0;
    size_t runtime_resident_bytes   = 0;
    size_t max_backend_bytes        = 0;

    // Runtime buffers only; the manager accounts for weights separately.
    size_t runtime_peak_bytes() const {
        return pending_allocation_bytes > SIZE_MAX - runtime_resident_bytes
                   ? SIZE_MAX
                   : runtime_resident_bytes + pending_allocation_bytes;
    }
};

struct DeviceMemoryCapacity {
    size_t required_device_bytes  = 0;
    size_t required_budget_bytes  = 0;
    size_t available_device_bytes = SIZE_MAX;
    size_t available_budget_bytes = SIZE_MAX;
    bool requires_reclamation     = false;
    bool valid                    = true;
    bool wrapped_device_report    = false;

    bool fits(bool after_reclamation = false) const {
        return valid && !requires_reclamation &&
               (after_reclamation || !wrapped_device_report) &&
               required_device_bytes <= available_device_bytes &&
               required_budget_bytes <= available_budget_bytes;
    }
};

struct WeightResidencyBlock {
    ggml_backend_buffer_t buffer = nullptr;
    ggml_backend_t backend       = nullptr;
    std::vector<ggml_tensor*> tensors;
    size_t bytes            = 0;
    uint64_t last_use_epoch = 0;
    bool requires_demotion  = false;
};

struct DeviceResidencyManager {
private:
    std::shared_ptr<sd::ExecutionEngine> execution_engine_;

public:
    virtual ~DeviceResidencyManager() = default;

    std::shared_ptr<sd::ExecutionEngine> execution_engine();

    virtual bool segmented_compute_enabled() const          = 0;
    virtual bool prefetch_enabled() const                   = 0;
    virtual void register_runtime_owner(uintptr_t owner_id) = 0;
    virtual void remove_runtime_owner(uintptr_t owner_id)   = 0;
    // Capacity requests select their backend's weights; protection spans all backends.
    virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request,
                                               const std::vector<ggml_tensor*>& required_params) const = 0;
    virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
                                        ggml_backend_t compute_backend)                                = 0;
    virtual ggml_tensor* resolve_param_tensor(ggml_tensor* tensor) const { return nullptr; }
    virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
    virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
    virtual void evict_compute_backend_params(const std::vector<ggml_tensor*>& tensors)   = 0;
    virtual WeightResidencyInfo inspect_compute_backend_params(
        const std::vector<ggml_tensor*>& tensors) const          = 0;
    virtual void update_runtime_residency(uintptr_t owner_id,
                                          ggml_backend_t compute_backend,
                                          size_t resident_bytes) = 0;
    virtual bool ensure_compute_backend_capacity(
        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) = 0;
    virtual DeviceMemoryCapacity inspect_compute_backend_capacity(
        const DeviceMemoryRequest& request,
        const std::vector<ggml_tensor*>& params,
        bool log_details = false) const {
        DeviceMemoryCapacity result;
        result.valid = false;
        return result;
    }
    virtual std::vector<WeightResidencyBlock> evictable_compute_backend_blocks(
        ggml_backend_t backend,
        const std::vector<ggml_tensor*>& protected_params) const { return {}; }
    virtual bool evict_compute_backend_block(const WeightResidencyBlock& block) { return false; }
    virtual WeightPrefetchResult prefetch_params(
        uintptr_t owner_id,
        const std::vector<ggml_tensor*>& tensors) = 0;
    virtual WeightPrefetchResult prepare_prefetch_params(
        uintptr_t owner_id,
        const std::vector<ggml_tensor*>& tensors) {
        return prefetch_params(owner_id, tensors);
    }
    virtual void submit_prefetched_params(uintptr_t owner_id) {}
    virtual bool activate_prefetched_params(uintptr_t owner_id,
                                            const std::vector<ggml_tensor*>& tensors) = 0;
    virtual void clear_prefetched_params(uintptr_t owner_id)                          = 0;
};

// Transitional alias for model constructors that have not yet adopted the
// residency-oriented name. It does not introduce a second implementation.
using RunnerWeightManager = DeviceResidencyManager;

#endif  // __DEVICE_RESIDENCY_MANAGER_H__