refactor: unify runner execution and weight residency (#2321)

src/device_residency_manager.h

@@ -3,12 +3,17 @@
 
 #include <cstdint>
 #include <functional>
+#include <memory>
 #include <vector>
 
 #include "ggml-backend.h"
 
 struct ggml_tensor;
 
+namespace sd {
+    class ExecutionEngine;
+}
+
 enum class WeightPrefetchResult {
     Scheduled,
     AlreadyResident,
@@ -36,13 +41,45 @@ struct DeviceMemoryRequest {
     }
 };
 
+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;
 
-    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;
+    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;
@@ -62,6 +99,18 @@ struct DeviceResidencyManager {
         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;