hive_gpu/
types.rs

1//! Core types for Hive GPU
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6/// A GPU vector with its associated data
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct GpuVector {
9    /// Unique identifier for the vector
10    pub id: String,
11    /// The vector data (always f32 for compatibility)
12    pub data: Vec<f32>,
13    /// Optional metadata associated with the vector
14    pub metadata: HashMap<String, String>,
15}
16
17impl GpuVector {
18    /// Create a new GPU vector
19    pub fn new(id: String, data: Vec<f32>) -> Self {
20        Self {
21            id,
22            data,
23            metadata: HashMap::new(),
24        }
25    }
26
27    /// Create a new GPU vector with metadata
28    pub fn with_metadata(id: String, data: Vec<f32>, metadata: HashMap<String, String>) -> Self {
29        Self { id, data, metadata }
30    }
31
32    /// Get the dimension of the vector
33    pub fn dimension(&self) -> usize {
34        self.data.len()
35    }
36
37    /// Get memory usage in bytes
38    pub fn memory_size(&self) -> usize {
39        self.data.len() * std::mem::size_of::<f32>() + self.id.len() + self.metadata.len() * 32 // rough estimate
40    }
41}
42
43impl From<&GpuVector> for Vec<f32> {
44    fn from(v: &GpuVector) -> Self {
45        v.data.clone()
46    }
47}
48
49/// Distance metrics for vector similarity
50#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
51#[serde(rename_all = "lowercase")]
52pub enum GpuDistanceMetric {
53    /// Cosine similarity
54    Cosine,
55    /// Euclidean distance
56    Euclidean,
57    /// Dot product
58    DotProduct,
59}
60
61impl std::fmt::Display for GpuDistanceMetric {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        match self {
64            GpuDistanceMetric::Cosine => write!(f, "cosine"),
65            GpuDistanceMetric::Euclidean => write!(f, "euclidean"),
66            GpuDistanceMetric::DotProduct => write!(f, "dot_product"),
67        }
68    }
69}
70
71/// Search result from GPU operations
72#[derive(Debug, Clone, Serialize, Deserialize)]
73pub struct GpuSearchResult {
74    /// Vector ID
75    pub id: String,
76    /// Similarity score
77    pub score: f32,
78    /// Vector index in storage
79    pub index: usize,
80}
81
82/// GPU device information
83#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct GpuDeviceInfo {
85    /// Device name
86    pub name: String,
87    /// Device type (Metal, CUDA, etc.)
88    pub device_type: String,
89    /// Available memory in bytes
90    pub memory_bytes: u64,
91    /// Maximum buffer size
92    pub max_buffer_size: u64,
93    /// Compute capability
94    pub compute_capability: Option<String>,
95}
96
97/// GPU capabilities
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct GpuCapabilities {
100    /// Supports HNSW operations
101    pub supports_hnsw: bool,
102    /// Supports batch operations
103    pub supports_batch: bool,
104    /// Maximum vector dimension
105    pub max_dimension: usize,
106    /// Maximum vectors per batch
107    pub max_batch_size: usize,
108}
109
110/// GPU memory statistics
111#[derive(Debug, Clone, Serialize, Deserialize)]
112pub struct GpuMemoryStats {
113    /// Total allocated memory in bytes
114    pub total_allocated: usize,
115    /// Available memory in bytes
116    pub available: usize,
117    /// Memory utilization percentage (0.0-1.0)
118    pub utilization: f32,
119    /// Number of active buffers
120    pub buffer_count: usize,
121}
122
123/// HNSW configuration for GPU operations
124#[derive(Debug, Clone, Serialize, Deserialize)]
125pub struct HnswConfig {
126    /// Number of bidirectional links created for each node
127    pub max_connections: usize,
128    /// Size of the dynamic list for nearest neighbors (construction)
129    pub ef_construction: usize,
130    /// Size of the dynamic list for nearest neighbors (search)
131    pub ef_search: usize,
132    /// Maximum level in the hierarchy
133    pub max_level: usize,
134    /// Level assignment multiplier
135    pub level_multiplier: f32,
136    /// Random seed for level assignment
137    pub seed: Option<u64>,
138}
139
140impl Default for HnswConfig {
141    fn default() -> Self {
142        Self {
143            max_connections: 16,
144            ef_construction: 100,
145            ef_search: 50,
146            max_level: 8,
147            level_multiplier: 0.5,
148            seed: None,
149        }
150    }
151}
152
153/// Vector metadata for GPU operations
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct VectorMetadata {
156    /// Original vector ID
157    pub original_id: String,
158    /// Index in storage
159    pub index: usize,
160    /// Timestamp of creation
161    pub timestamp: u64,
162}