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}