hive_gpu/
traits.rs

1//! Core traits for Hive GPU
2
3use crate::types::{GpuVector, GpuSearchResult, GpuDistanceMetric, GpuDeviceInfo, GpuCapabilities, GpuMemoryStats, HnswConfig};
4use crate::error::Result;
5
6/// Core GPU backend trait
7pub trait GpuBackend {
8    /// Get device information
9    fn device_info(&self) -> GpuDeviceInfo;
10    
11    /// Get device capabilities
12    fn supports_operations(&self) -> GpuCapabilities;
13    
14    /// Get memory statistics
15    fn memory_stats(&self) -> GpuMemoryStats;
16}
17
18/// GPU vector storage trait
19pub trait GpuVectorStorage {
20    /// Add vectors to storage
21    fn add_vectors(&mut self, vectors: &[GpuVector]) -> Result<Vec<usize>>;
22    
23    /// Search for similar vectors
24    fn search(&self, query: &[f32], limit: usize) -> Result<Vec<GpuSearchResult>>;
25    
26    /// Remove vectors by IDs
27    fn remove_vectors(&mut self, ids: &[String]) -> Result<()>;
28    
29    /// Get total vector count
30    fn vector_count(&self) -> usize;
31    
32    /// Get vector dimension
33    fn dimension(&self) -> usize;
34    
35    /// Get vector by ID
36    fn get_vector(&self, id: &str) -> Result<Option<GpuVector>>;
37    
38    /// Clear all vectors
39    fn clear(&mut self) -> Result<()>;
40}
41
42/// GPU context trait for creating storage
43pub trait GpuContext {
44    /// Create vector storage with default configuration
45    fn create_storage(&self, dimension: usize, metric: GpuDistanceMetric) -> Result<Box<dyn GpuVectorStorage>>;
46    
47    /// Create vector storage with HNSW configuration
48    fn create_storage_with_config(&self, dimension: usize, metric: GpuDistanceMetric, config: HnswConfig) -> Result<Box<dyn GpuVectorStorage>>;
49    
50    /// Get memory statistics
51    fn memory_stats(&self) -> GpuMemoryStats;
52    
53    /// Get device information
54    fn device_info(&self) -> GpuDeviceInfo;
55}
56
57/// GPU buffer management trait
58pub trait GpuBufferManager {
59    /// Allocate buffer with specified size
60    fn allocate_buffer(&mut self, size: usize) -> Result<GpuBuffer>;
61    
62    /// Deallocate buffer
63    fn deallocate_buffer(&mut self, buffer: GpuBuffer) -> Result<()>;
64    
65    /// Get buffer pool statistics
66    fn pool_stats(&self) -> BufferPoolStats;
67}
68
69/// GPU buffer handle
70#[derive(Debug, Clone)]
71pub struct GpuBuffer {
72    /// Buffer ID
73    pub id: usize,
74    /// Buffer size in bytes
75    pub size: usize,
76    /// Buffer type
77    pub buffer_type: BufferType,
78}
79
80/// Buffer types
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub enum BufferType {
83    /// Vector data buffer
84    Vector,
85    /// Metadata buffer
86    Metadata,
87    /// Temporary computation buffer
88    Temporary,
89    /// HNSW graph buffer
90    HnswGraph,
91}
92
93/// Buffer pool statistics
94#[derive(Debug, Clone)]
95pub struct BufferPoolStats {
96    /// Total buffers in pool
97    pub total_buffers: usize,
98    /// Available buffers
99    pub available_buffers: usize,
100    /// Pool utilization (0.0-1.0)
101    pub utilization: f32,
102    /// Total memory allocated
103    pub total_memory: usize,
104}
105
106/// GPU monitoring trait
107pub trait GpuMonitor {
108    /// Get VRAM usage statistics
109    fn get_vram_stats(&self) -> VramStats;
110    
111    /// Validate VRAM-only operation
112    fn validate_all_vram(&self) -> Result<()>;
113    
114    /// Generate VRAM report
115    fn generate_vram_report(&self) -> String;
116}
117
118/// VRAM statistics
119#[derive(Debug, Clone)]
120pub struct VramStats {
121    /// Total VRAM in bytes
122    pub total_vram: usize,
123    /// Allocated VRAM in bytes
124    pub allocated_vram: usize,
125    /// Available VRAM in bytes
126    pub available_vram: usize,
127    /// VRAM utilization (0.0-1.0)
128    pub utilization: f32,
129    /// Number of VRAM buffers
130    pub buffer_count: usize,
131}
132
133/// VRAM buffer information
134#[derive(Debug, Clone)]
135pub struct VramBufferInfo {
136    /// Buffer ID
137    pub buffer_id: usize,
138    /// Buffer size in bytes
139    pub size: usize,
140    /// Buffer type
141    pub buffer_type: BufferType,
142    /// Allocation timestamp
143    pub allocated_at: u64,
144}