mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-09-25 12:40:41 +00:00
78 lines
3.5 KiB
C++
78 lines
3.5 KiB
C++
#ifndef __DEVICE_RESIDENCY_MANAGER_H__
|
|
#define __DEVICE_RESIDENCY_MANAGER_H__
|
|
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <vector>
|
|
|
|
#include "ggml-backend.h"
|
|
|
|
struct ggml_tensor;
|
|
|
|
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 DeviceResidencyManager {
|
|
virtual ~DeviceResidencyManager() = default;
|
|
|
|
virtual bool segmented_compute_enabled() const = 0;
|
|
virtual bool prefetch_enabled() const = 0;
|
|
virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function<bool()> reclaim) = 0;
|
|
virtual void remove_runtime_owner(uintptr_t owner_id) = 0;
|
|
// 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 WeightPrefetchResult prefetch_params(
|
|
uintptr_t owner_id,
|
|
const std::vector<ggml_tensor*>& tensors) = 0;
|
|
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__
|