Skip to main content

hive_gpu/utils/
memory.rs

1//! Memory management utilities
2
3/// Memory utility functions
4pub mod memory_utils {
5    /// Calculate memory size for vectors
6    pub fn calculate_vector_memory_size(dimension: usize, count: usize) -> usize {
7        count * dimension * std::mem::size_of::<f32>()
8    }
9
10    /// Calculate memory size for metadata
11    pub fn calculate_metadata_memory_size(count: usize) -> usize {
12        count * 256 // 256 bytes per vector metadata
13    }
14
15    /// Check if memory allocation is safe
16    pub fn is_memory_allocation_safe(
17        requested: usize,
18        available: usize,
19        safety_margin: f32,
20    ) -> bool {
21        let safe_limit = (available as f32 * safety_margin) as usize;
22        requested <= safe_limit
23    }
24
25    /// Calculate optimal buffer size
26    pub fn calculate_optimal_buffer_size(
27        current_size: usize,
28        growth_factor: f32,
29        max_size: usize,
30    ) -> usize {
31        let new_size = (current_size as f32 * growth_factor) as usize;
32        new_size.min(max_size)
33    }
34}
35
36/// Buffer utility functions
37pub mod buffer_utils {
38    use crate::error::{HiveGpuError, Result};
39
40    use super::memory_utils;
41
42    /// Buffer alignment for optimal GPU performance
43    pub const BUFFER_ALIGNMENT: usize = 256;
44
45    /// Align buffer size to GPU requirements
46    pub fn align_buffer_size(size: usize) -> usize {
47        size.div_ceil(BUFFER_ALIGNMENT) * BUFFER_ALIGNMENT
48    }
49
50    /// Calculate buffer pool size
51    pub fn calculate_buffer_pool_size(
52        vector_count: usize,
53        dimension: usize,
54        pool_factor: f32,
55    ) -> usize {
56        let base_size = memory_utils::calculate_vector_memory_size(dimension, vector_count);
57        (base_size as f32 * pool_factor) as usize
58    }
59
60    /// Validate buffer parameters
61    pub fn validate_buffer_parameters(
62        size: usize,
63        alignment: usize,
64        max_size: usize,
65    ) -> Result<()> {
66        if size == 0 {
67            return Err(HiveGpuError::BufferAllocationFailed(
68                "Zero buffer size".to_string(),
69            ));
70        }
71
72        if size > max_size {
73            return Err(HiveGpuError::BufferAllocationFailed(format!(
74                "Buffer size {} exceeds maximum {}",
75                size, max_size
76            )));
77        }
78
79        if !size.is_multiple_of(alignment) {
80            return Err(HiveGpuError::BufferAllocationFailed(format!(
81                "Buffer size {} not aligned to {}",
82                size, alignment
83            )));
84        }
85
86        Ok(())
87    }
88}