1use crate::types::{GpuVector, GpuSearchResult, GpuDistanceMetric, GpuDeviceInfo, GpuCapabilities, GpuMemoryStats, HnswConfig};
4use crate::error::Result;
5
6pub trait GpuBackend {
8 fn device_info(&self) -> GpuDeviceInfo;
10
11 fn supports_operations(&self) -> GpuCapabilities;
13
14 fn memory_stats(&self) -> GpuMemoryStats;
16}
17
18pub trait GpuVectorStorage {
20 fn add_vectors(&mut self, vectors: &[GpuVector]) -> Result<Vec<usize>>;
22
23 fn search(&self, query: &[f32], limit: usize) -> Result<Vec<GpuSearchResult>>;
25
26 fn remove_vectors(&mut self, ids: &[String]) -> Result<()>;
28
29 fn vector_count(&self) -> usize;
31
32 fn dimension(&self) -> usize;
34
35 fn get_vector(&self, id: &str) -> Result<Option<GpuVector>>;
37
38 fn clear(&mut self) -> Result<()>;
40}
41
42pub trait GpuContext {
44 fn create_storage(&self, dimension: usize, metric: GpuDistanceMetric) -> Result<Box<dyn GpuVectorStorage>>;
46
47 fn create_storage_with_config(&self, dimension: usize, metric: GpuDistanceMetric, config: HnswConfig) -> Result<Box<dyn GpuVectorStorage>>;
49
50 fn memory_stats(&self) -> GpuMemoryStats;
52
53 fn device_info(&self) -> GpuDeviceInfo;
55}
56
57pub trait GpuBufferManager {
59 fn allocate_buffer(&mut self, size: usize) -> Result<GpuBuffer>;
61
62 fn deallocate_buffer(&mut self, buffer: GpuBuffer) -> Result<()>;
64
65 fn pool_stats(&self) -> BufferPoolStats;
67}
68
69#[derive(Debug, Clone)]
71pub struct GpuBuffer {
72 pub id: usize,
74 pub size: usize,
76 pub buffer_type: BufferType,
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub enum BufferType {
83 Vector,
85 Metadata,
87 Temporary,
89 HnswGraph,
91}
92
93#[derive(Debug, Clone)]
95pub struct BufferPoolStats {
96 pub total_buffers: usize,
98 pub available_buffers: usize,
100 pub utilization: f32,
102 pub total_memory: usize,
104}
105
106pub trait GpuMonitor {
108 fn get_vram_stats(&self) -> VramStats;
110
111 fn validate_all_vram(&self) -> Result<()>;
113
114 fn generate_vram_report(&self) -> String;
116}
117
118#[derive(Debug, Clone)]
120pub struct VramStats {
121 pub total_vram: usize,
123 pub allocated_vram: usize,
125 pub available_vram: usize,
127 pub utilization: f32,
129 pub buffer_count: usize,
131}
132
133#[derive(Debug, Clone)]
135pub struct VramBufferInfo {
136 pub buffer_id: usize,
138 pub size: usize,
140 pub buffer_type: BufferType,
142 pub allocated_at: u64,
144}