hive_gpu/traits.rs
1//! Core traits for Hive GPU
2
3use crate::error::Result;
4use crate::types::{
5 GpuCapabilities, GpuDeviceInfo, GpuDistanceMetric, GpuMemoryStats, GpuSearchResult, GpuVector,
6 HnswConfig,
7};
8
9/// Core GPU backend trait
10pub trait GpuBackend {
11 /// Get device information
12 fn device_info(&self) -> GpuDeviceInfo;
13
14 /// Get device capabilities
15 fn supports_operations(&self) -> GpuCapabilities;
16
17 /// Get memory statistics
18 fn memory_stats(&self) -> GpuMemoryStats;
19}
20
21/// GPU vector storage trait
22pub trait GpuVectorStorage {
23 /// Add vectors to storage
24 fn add_vectors(&mut self, vectors: &[GpuVector]) -> Result<Vec<usize>>;
25
26 /// Search for similar vectors
27 fn search(&self, query: &[f32], limit: usize) -> Result<Vec<GpuSearchResult>>;
28
29 /// Remove vectors by IDs
30 fn remove_vectors(&mut self, ids: &[String]) -> Result<()>;
31
32 /// Get total vector count
33 fn vector_count(&self) -> usize;
34
35 /// Get vector dimension
36 fn dimension(&self) -> usize;
37
38 /// Get vector by ID
39 fn get_vector(&self, id: &str) -> Result<Option<GpuVector>>;
40
41 /// Clear all vectors
42 fn clear(&mut self) -> Result<()>;
43}
44
45/// GPU context trait for creating storage
46pub trait GpuContext {
47 /// Create vector storage with default configuration
48 fn create_storage(
49 &self,
50 dimension: usize,
51 metric: GpuDistanceMetric,
52 ) -> Result<Box<dyn GpuVectorStorage>>;
53
54 /// Create vector storage with HNSW configuration
55 fn create_storage_with_config(
56 &self,
57 dimension: usize,
58 metric: GpuDistanceMetric,
59 config: HnswConfig,
60 ) -> Result<Box<dyn GpuVectorStorage>>;
61
62 /// Get memory statistics
63 fn memory_stats(&self) -> GpuMemoryStats;
64
65 /// Get device information
66 ///
67 /// Returns detailed information about the GPU device including name,
68 /// VRAM statistics, compute capabilities, and backend-specific details.
69 ///
70 /// # Errors
71 ///
72 /// Returns an error if device information cannot be queried.
73 ///
74 /// # Examples
75 ///
76 /// ```no_run
77 /// # #[cfg(all(target_os = "macos", feature = "metal-native"))]
78 /// # {
79 /// use hive_gpu::metal::MetalNativeContext;
80 /// use hive_gpu::traits::GpuContext;
81 ///
82 /// # fn example() -> Result<(), hive_gpu::error::HiveGpuError> {
83 /// let context = MetalNativeContext::new()?;
84 /// let info = context.device_info()?;
85 ///
86 /// println!("Device: {} ({})", info.name, info.backend);
87 /// println!("Total VRAM: {} MB", info.total_vram_mb());
88 /// println!("Available VRAM: {} MB", info.available_vram_mb());
89 /// println!("Usage: {:.1}%", info.vram_usage_percent());
90 /// # Ok(())
91 /// # }
92 /// # }
93 /// ```
94 fn device_info(&self) -> Result<GpuDeviceInfo>;
95}
96
97/// GPU buffer management trait
98pub trait GpuBufferManager {
99 /// Allocate buffer with specified size
100 fn allocate_buffer(&mut self, size: usize) -> Result<GpuBuffer>;
101
102 /// Deallocate buffer
103 fn deallocate_buffer(&mut self, buffer: GpuBuffer) -> Result<()>;
104
105 /// Get buffer pool statistics
106 fn pool_stats(&self) -> BufferPoolStats;
107}
108
109/// GPU buffer handle
110#[derive(Debug, Clone)]
111pub struct GpuBuffer {
112 /// Buffer ID
113 pub id: usize,
114 /// Buffer size in bytes
115 pub size: usize,
116 /// Buffer type
117 pub buffer_type: BufferType,
118}
119
120/// Buffer types
121#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
122pub enum BufferType {
123 /// Vector data buffer
124 Vector,
125 /// Metadata buffer
126 Metadata,
127 /// Temporary computation buffer
128 Temporary,
129 /// HNSW graph buffer
130 HnswGraph,
131}
132
133/// Buffer pool statistics
134#[derive(Debug, Clone)]
135pub struct BufferPoolStats {
136 /// Total buffers in pool
137 pub total_buffers: usize,
138 /// Available buffers
139 pub available_buffers: usize,
140 /// Pool utilization (0.0-1.0)
141 pub utilization: f32,
142 /// Total memory allocated
143 pub total_memory: usize,
144}
145
146/// GPU monitoring trait
147pub trait GpuMonitor {
148 /// Get VRAM usage statistics
149 fn get_vram_stats(&self) -> VramStats;
150
151 /// Validate VRAM-only operation
152 fn validate_all_vram(&self) -> Result<()>;
153
154 /// Generate VRAM report
155 fn generate_vram_report(&self) -> String;
156}
157
158/// VRAM statistics
159#[derive(Debug, Clone)]
160pub struct VramStats {
161 /// Total VRAM in bytes
162 pub total_vram: usize,
163 /// Allocated VRAM in bytes
164 pub allocated_vram: usize,
165 /// Available VRAM in bytes
166 pub available_vram: usize,
167 /// VRAM utilization (0.0-1.0)
168 pub utilization: f32,
169 /// Number of VRAM buffers
170 pub buffer_count: usize,
171}
172
173/// VRAM buffer information
174#[derive(Debug, Clone)]
175pub struct VramBufferInfo {
176 /// Buffer ID
177 pub buffer_id: usize,
178 /// Buffer size in bytes
179 pub size: usize,
180 /// Buffer type
181 pub buffer_type: BufferType,
182 /// Allocation timestamp
183 pub allocated_at: u64,
184}