trueno_graph/gpu/
memory.rs

1//! GPU memory management and VRAM detection
2//!
3//! Detects available GPU memory and manages memory limits for paging.
4//! Based on Funke et al. (2018) "Pipelined Query Processing in Coprocessor Environments"
5
6use super::GpuDevice;
7use anyhow::Result;
8
9/// Default morsel size (128MB chunks as per DuckDB/Umbra)
10pub const DEFAULT_MORSEL_SIZE: usize = 128 * 1024 * 1024; // 128 MB
11
12/// GPU memory limits and statistics
13#[derive(Debug, Clone)]
14pub struct GpuMemoryLimits {
15    /// Total VRAM available (bytes)
16    pub total_vram: u64,
17
18    /// Maximum memory to use for graph data (70% of total to leave headroom)
19    pub usable_vram: u64,
20
21    /// Morsel size for chunking (default 128MB)
22    pub morsel_size: usize,
23
24    /// Maximum number of morsels that fit in VRAM
25    pub max_morsels: usize,
26}
27
28impl GpuMemoryLimits {
29    /// Detect GPU memory limits
30    ///
31    /// # Errors
32    ///
33    /// Returns error if GPU device is not available
34    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss, clippy::cast_precision_loss)]
35    pub fn detect(device: &GpuDevice) -> Result<Self> {
36        // Get adapter limits from wgpu
37        let limits = device.device().limits();
38
39        // Estimate VRAM from buffer size limit (conservative estimate)
40        // Most GPUs have max_buffer_size = VRAM size or similar
41        let total_vram = limits.max_buffer_size;
42
43        // Use 70% of VRAM for graph data (30% for shaders, temporaries, OS)
44        let usable_vram = (total_vram as f64 * 0.7) as u64;
45
46        let morsel_size = DEFAULT_MORSEL_SIZE;
47        let max_morsels = (usable_vram as usize / morsel_size).max(1);
48
49        Ok(Self { total_vram, usable_vram, morsel_size, max_morsels })
50    }
51
52    /// Check if graph fits entirely in VRAM
53    #[must_use]
54    #[allow(clippy::cast_possible_truncation)]
55    pub fn fits_in_vram(&self, graph_size_bytes: usize) -> bool {
56        graph_size_bytes <= self.usable_vram as usize
57    }
58
59    /// Calculate number of morsels needed for given size
60    #[must_use]
61    pub fn morsels_needed(&self, size_bytes: usize) -> usize {
62        size_bytes.div_ceil(self.morsel_size)
63    }
64
65    /// Get recommended tile size in nodes for chunking
66    #[must_use]
67    pub fn recommended_tile_size(&self, node_size_bytes: usize) -> usize {
68        if node_size_bytes == 0 {
69            return 1000; // Default fallback
70        }
71        (self.morsel_size / node_size_bytes).max(100)
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[tokio::test]
80    async fn test_memory_limits_detection() {
81        if !GpuDevice::is_gpu_available().await {
82            eprintln!("⚠️  Skipping test_memory_limits_detection: GPU not available");
83            return;
84        }
85
86        let device = GpuDevice::new().await.unwrap();
87        let limits = GpuMemoryLimits::detect(&device).unwrap();
88
89        println!("Total VRAM: {} bytes", limits.total_vram);
90        println!("Usable VRAM: {} bytes", limits.usable_vram);
91        println!("Max morsels: {}", limits.max_morsels);
92
93        assert!(limits.total_vram > 0);
94        assert!(limits.usable_vram > 0);
95        assert!(limits.usable_vram <= limits.total_vram);
96        assert!(limits.max_morsels > 0);
97    }
98
99    #[test]
100    fn test_fits_in_vram() {
101        let limits = GpuMemoryLimits {
102            total_vram: 8 * 1024 * 1024 * 1024,  // 8GB
103            usable_vram: 5 * 1024 * 1024 * 1024, // 5GB usable
104            morsel_size: DEFAULT_MORSEL_SIZE,
105            max_morsels: 40,
106        };
107
108        assert!(limits.fits_in_vram(100 * 1024 * 1024)); // 100MB fits
109        assert!(limits.fits_in_vram(4 * 1024 * 1024 * 1024)); // 4GB fits
110        assert!(!limits.fits_in_vram(6 * 1024 * 1024 * 1024)); // 6GB doesn't fit
111    }
112
113    #[test]
114    fn test_morsels_needed() {
115        let limits = GpuMemoryLimits {
116            total_vram: 8 * 1024 * 1024 * 1024,
117            usable_vram: 5 * 1024 * 1024 * 1024,
118            morsel_size: DEFAULT_MORSEL_SIZE,
119            max_morsels: 40,
120        };
121
122        assert_eq!(limits.morsels_needed(128 * 1024 * 1024), 1);
123        assert_eq!(limits.morsels_needed(256 * 1024 * 1024), 2);
124        assert_eq!(limits.morsels_needed(200 * 1024 * 1024), 2);
125    }
126}
trueno_graph/gpu/memory.rs

trueno_graph/gpu/
memory.rs
