Skip to main content

scirs2_io/gpu/
mod.rs

1//! GPU-accelerated I/O operations with comprehensive backend support
2//!
3//! This module provides a complete GPU acceleration framework for I/O operations
4//! including backend management, compression, memory management, and performance
5//! optimization across CUDA, Metal, and OpenCL backends.
6
7pub mod backend_management;
8pub mod compression;
9pub mod memory_management;
10
11// Re-export key types for easy access
12pub use backend_management::{
13    BackendCapabilities, BackendPerformanceProfile, GpuIoProcessor, GpuWorkloadType,
14};
15
16pub use compression::{CompressionStats, GpuCompressionProcessor};
17
18pub use memory_management::{
19    AdvancedGpuMemoryPool, AllocationStats, BufferMetadata, FragmentationManager, GlobalPoolStats,
20    GpuMemoryPoolManager, MemoryType, PoolConfig, PoolStats, PooledBuffer,
21};
22
23use crate::error::Result;
24use scirs2_core::gpu::{GpuBackend, GpuDataType, GpuDevice};
25use scirs2_core::ndarray::{Array1, ArrayView1};
26
27/// Unified GPU I/O interface combining all GPU acceleration capabilities
28#[derive(Debug)]
29pub struct UnifiedGpuProcessor {
30    io_processor: GpuIoProcessor,
31    compression_processor: GpuCompressionProcessor,
32    memory_manager: GpuMemoryPoolManager,
33}
34
35impl UnifiedGpuProcessor {
36    /// Create a new unified GPU processor with optimal configuration
37    pub fn new() -> Result<Self> {
38        let io_processor = GpuIoProcessor::new()?;
39        let compression_processor = GpuCompressionProcessor::new()?;
40        let memory_manager = GpuMemoryPoolManager::new(io_processor.device.clone())?;
41
42        Ok(Self {
43            io_processor,
44            compression_processor,
45            memory_manager,
46        })
47    }
48
49    /// Create with specific backend
50    pub fn with_backend(backend: GpuBackend) -> Result<Self> {
51        let io_processor = GpuIoProcessor::with_backend(backend)?;
52        let compression_processor = GpuCompressionProcessor::new()?;
53        let memory_manager = GpuMemoryPoolManager::new(io_processor.device.clone())?;
54
55        Ok(Self {
56            io_processor,
57            compression_processor,
58            memory_manager,
59        })
60    }
61
62    /// Get the current GPU backend
63    pub fn backend(&self) -> GpuBackend {
64        self.io_processor.backend()
65    }
66
67    /// Get comprehensive GPU capabilities
68    pub fn get_capabilities(&self) -> Result<GpuCapabilities> {
69        let backend_caps = self.io_processor.get_backend_capabilities()?;
70        let compression_stats = self.compression_processor.get_performance_stats();
71        let memory_stats = self.memory_manager.get_global_stats();
72
73        Ok(GpuCapabilities {
74            backend: backend_caps.backend,
75            memory_gb: backend_caps.memory_gb,
76            compute_units: backend_caps.compute_units,
77            supports_fp64: backend_caps.supports_fp64,
78            supports_fp16: backend_caps.supports_fp16,
79            compression_throughput_gbps: compression_stats.estimated_throughput_gbps,
80            memory_pools: memory_stats.pool_count,
81            total_pool_size: memory_stats.total_pool_size,
82            performance_score: self.calculate_performance_score(&backend_caps),
83        })
84    }
85
86    /// Compress data with automatic backend optimization
87    pub fn compress<T: GpuDataType>(
88        &self,
89        data: &ArrayView1<T>,
90        algorithm: crate::compression::CompressionAlgorithm,
91        level: Option<u32>,
92    ) -> Result<Vec<u8>> {
93        self.compression_processor
94            .compress_gpu(data, algorithm, level)
95    }
96
97    /// Decompress data with automatic backend optimization
98    pub fn decompress<T: GpuDataType>(
99        &self,
100        compressed_data: &[u8],
101        algorithm: crate::compression::CompressionAlgorithm,
102        expected_size: usize,
103    ) -> Result<Array1<T>> {
104        self.compression_processor
105            .decompress_gpu(compressed_data, algorithm, expected_size)
106    }
107
108    /// Allocate GPU memory buffer
109    pub fn allocate_buffer(
110        &mut self,
111        size: usize,
112        memory_type: MemoryType,
113    ) -> Result<PooledBuffer> {
114        self.memory_manager.allocate(size, memory_type)
115    }
116
117    /// Return buffer to memory pool
118    pub fn deallocate_buffer(
119        &mut self,
120        buffer: PooledBuffer,
121        memory_type: MemoryType,
122    ) -> Result<()> {
123        self.memory_manager.deallocate(buffer, memory_type)
124    }
125
126    /// Get comprehensive performance statistics
127    pub fn get_performance_stats(&self) -> UnifiedGpuStats {
128        let backend_caps = self
129            .io_processor
130            .get_backend_capabilities()
131            .unwrap_or_else(|_| BackendCapabilities {
132                backend: GpuBackend::Cpu,
133                memory_gb: 1.0,
134                max_work_group_size: 64,
135                supports_fp64: false,
136                supports_fp16: false,
137                compute_units: 1,
138                max_allocation_size: 1024 * 1024,
139                local_memory_size: 64 * 1024,
140            });
141
142        let compression_stats = self.compression_processor.get_performance_stats();
143        let memory_stats = self.memory_manager.get_global_stats();
144
145        UnifiedGpuStats {
146            backend: backend_caps.backend,
147            compression_stats,
148            memory_stats: memory_stats.clone(),
149            overall_efficiency: self.calculate_efficiency_score(&memory_stats),
150        }
151    }
152
153    /// Perform maintenance operations (garbage collection, compaction)
154    pub fn maintenance(&mut self) -> Result<MaintenanceReport> {
155        let freed_buffers = self.memory_manager.garbage_collect_all()?;
156
157        Ok(MaintenanceReport {
158            freed_buffers,
159            timestamp: std::time::Instant::now(),
160        })
161    }
162
163    /// Optimize processor for specific workload
164    pub fn optimize_for_workload(&mut self, workload: GpuWorkloadType) -> Result<()> {
165        // Adjust memory pool configurations based on workload
166        match workload {
167            GpuWorkloadType::MachineLearning => {
168                // ML workloads benefit from larger buffers and device memory
169                let pool_size = 512 * 1024 * 1024; // 512MB
170                self.memory_manager
171                    .create_pool(pool_size, MemoryType::Device)?;
172            }
173            GpuWorkloadType::ImageProcessing => {
174                // Image processing needs unified memory for CPU/GPU transfers
175                let pool_size = 256 * 1024 * 1024; // 256MB
176                self.memory_manager
177                    .create_pool(pool_size, MemoryType::Unified)?;
178            }
179            GpuWorkloadType::Compression => {
180                // Compression benefits from pinned memory for fast transfers
181                let pool_size = 128 * 1024 * 1024; // 128MB
182                self.memory_manager
183                    .create_pool(pool_size, MemoryType::Pinned)?;
184            }
185            GpuWorkloadType::GeneralCompute => {
186                // Balanced approach for general compute
187                let pool_size = 256 * 1024 * 1024; // 256MB
188                self.memory_manager
189                    .create_pool(pool_size, MemoryType::Device)?;
190            }
191        }
192
193        Ok(())
194    }
195
196    // Private helper methods
197    fn calculate_performance_score(&self, caps: &BackendCapabilities) -> f64 {
198        let memory_score = (caps.memory_gb / 16.0).min(1.0); // Normalize to 16GB max
199        let compute_score = (caps.compute_units as f64 / 64.0).min(1.0); // Normalize to 64 units
200        let feature_score = if caps.supports_fp64 && caps.supports_fp16 {
201            1.0
202        } else {
203            0.7
204        };
205
206        (memory_score + compute_score + feature_score) / 3.0
207    }
208
209    fn calculate_efficiency_score(&self, memory_stats: &GlobalPoolStats) -> f64 {
210        let allocation_efficiency = memory_stats.global_allocation_stats.get_cache_hit_rate();
211        let fragmentation_penalty = 1.0 - memory_stats.average_fragmentation.min(1.0);
212
213        (allocation_efficiency + fragmentation_penalty) / 2.0
214    }
215}
216
217impl Default for UnifiedGpuProcessor {
218    fn default() -> Self {
219        Self::new().unwrap_or_else(|_| {
220            // Fallback to CPU-only processor
221            let device = GpuDevice::new(GpuBackend::Cpu, 0);
222            Self {
223                io_processor: GpuIoProcessor::default(),
224                compression_processor: GpuCompressionProcessor::default(),
225                memory_manager: GpuMemoryPoolManager::new(device)
226                    .unwrap_or_else(|_| panic!("Failed to create fallback GPU memory manager")),
227            }
228        })
229    }
230}
231
232/// Comprehensive GPU capabilities information
233#[derive(Debug, Clone)]
234pub struct GpuCapabilities {
235    /// GPU backend type
236    pub backend: GpuBackend,
237    /// Total memory in gigabytes
238    pub memory_gb: f64,
239    /// Number of compute units
240    pub compute_units: usize,
241    /// Whether FP64 operations are supported
242    pub supports_fp64: bool,
243    /// Whether FP16 operations are supported
244    pub supports_fp16: bool,
245    /// Compression throughput in GB/s
246    pub compression_throughput_gbps: f64,
247    /// Number of memory pools
248    pub memory_pools: usize,
249    /// Total size of all pools
250    pub total_pool_size: usize,
251    /// Overall performance score (0.0-1.0)
252    pub performance_score: f64,
253}
254
255impl GpuCapabilities {
256    /// Check if GPU is suitable for high-performance computing
257    pub fn is_hpc_capable(&self) -> bool {
258        self.memory_gb >= 4.0
259            && self.compute_units >= 16
260            && self.supports_fp64
261            && self.performance_score >= 0.7
262    }
263
264    /// Check if GPU is suitable for AI/ML workloads
265    pub fn is_ml_capable(&self) -> bool {
266        self.memory_gb >= 6.0
267            && self.compute_units >= 32
268            && (self.supports_fp16 || self.supports_fp64)
269            && self.performance_score >= 0.6
270    }
271
272    /// Get recommended workload types for this GPU
273    pub fn get_recommended_workloads(&self) -> Vec<GpuWorkloadType> {
274        let mut workloads = Vec::new();
275
276        if self.is_ml_capable() {
277            workloads.push(GpuWorkloadType::MachineLearning);
278        }
279
280        if self.memory_gb >= 2.0 && self.compute_units >= 8 {
281            workloads.push(GpuWorkloadType::ImageProcessing);
282        }
283
284        if self.compression_throughput_gbps >= 1.0 {
285            workloads.push(GpuWorkloadType::Compression);
286        }
287
288        workloads.push(GpuWorkloadType::GeneralCompute);
289        workloads
290    }
291}
292
293/// Unified GPU performance statistics
294#[derive(Debug, Clone)]
295pub struct UnifiedGpuStats {
296    /// GPU backend type
297    pub backend: GpuBackend,
298    /// Compression performance statistics
299    pub compression_stats: CompressionStats,
300    /// Memory pool statistics
301    pub memory_stats: GlobalPoolStats,
302    /// Overall efficiency score (0.0-1.0)
303    pub overall_efficiency: f64,
304}
305
306/// Maintenance operation report
307#[derive(Debug)]
308pub struct MaintenanceReport {
309    /// Number of buffers freed during maintenance
310    pub freed_buffers: usize,
311    /// Timestamp of maintenance operation
312    pub timestamp: std::time::Instant,
313}
314
315/// Convenience functions for common GPU operations
316pub mod utils {
317    use super::*;
318
319    /// Check if any GPU backend is available
320    pub fn is_gpu_available() -> bool {
321        GpuIoProcessor::list_available_backends()
322            .iter()
323            .any(|&backend| backend != GpuBackend::Cpu)
324    }
325
326    /// Get the best available GPU backend for a workload
327    pub fn get_best_backend_for_workload(workload: GpuWorkloadType) -> Result<GpuBackend> {
328        GpuIoProcessor::get_optimal_backend_for_workload(workload)
329    }
330
331    /// Create optimized GPU processor for specific workload
332    pub fn create_optimized_processor(workload: GpuWorkloadType) -> Result<UnifiedGpuProcessor> {
333        let backend = get_best_backend_for_workload(workload)?;
334        let mut processor = UnifiedGpuProcessor::with_backend(backend)?;
335        processor.optimize_for_workload(workload)?;
336        Ok(processor)
337    }
338
339    /// Benchmark GPU performance for different operations
340    pub fn benchmark_gpu_performance(processor: &UnifiedGpuProcessor) -> GpuBenchmarkResults {
341        let capabilities = processor
342            .get_capabilities()
343            .unwrap_or_else(|_| GpuCapabilities {
344                backend: GpuBackend::Cpu,
345                memory_gb: 1.0,
346                compute_units: 1,
347                supports_fp64: false,
348                supports_fp16: false,
349                compression_throughput_gbps: 0.1,
350                memory_pools: 1,
351                total_pool_size: 1024 * 1024,
352                performance_score: 0.1,
353            });
354
355        GpuBenchmarkResults {
356            backend: capabilities.backend,
357            memory_bandwidth_gbps: capabilities.memory_gb * 0.8, // Estimate 80% peak
358            compute_throughput: capabilities.compute_units as f64 * 100.0, // Arbitrary units
359            compression_throughput: capabilities.compression_throughput_gbps,
360            overall_score: capabilities.performance_score,
361        }
362    }
363}
364
365/// GPU benchmark results
366#[derive(Debug, Clone)]
367pub struct GpuBenchmarkResults {
368    /// GPU backend type
369    pub backend: GpuBackend,
370    /// Memory bandwidth in GB/s
371    pub memory_bandwidth_gbps: f64,
372    /// Compute throughput (arbitrary units)
373    pub compute_throughput: f64,
374    /// Compression throughput in GB/s
375    pub compression_throughput: f64,
376    /// Overall performance score (0.0-1.0)
377    pub overall_score: f64,
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    #[test]
385    fn test_unified_processor_creation() {
386        // Should construct successfully and select any supported backend (CPU or GPU)
387        let processor = UnifiedGpuProcessor::default();
388        let backend = processor.backend();
389        match backend {
390            GpuBackend::Cpu
391            | GpuBackend::Metal
392            | GpuBackend::OpenCL
393            | GpuBackend::Cuda
394            | GpuBackend::Rocm
395            | GpuBackend::Wgpu => {}
396        }
397    }
398
399    #[test]
400    fn test_gpu_availability_check() {
401        // Should always return true due to CPU fallback
402        // Test passes if it doesn't panic (CPU is always available)
403    }
404
405    #[test]
406    fn test_gpu_capabilities() {
407        let processor = UnifiedGpuProcessor::default();
408        let capabilities = processor.get_capabilities();
409        assert!(capabilities.is_ok());
410
411        let caps = capabilities.expect("Operation failed");
412        assert!(caps.memory_gb > 0.0);
413        assert!(caps.compute_units > 0);
414    }
415
416    #[test]
417    fn test_workload_optimization() {
418        let mut processor = UnifiedGpuProcessor::default();
419        let result = processor.optimize_for_workload(GpuWorkloadType::MachineLearning);
420        assert!(result.is_ok());
421    }
422
423    #[test]
424    fn test_performance_benchmarking() {
425        let processor = UnifiedGpuProcessor::default();
426        let benchmark = utils::benchmark_gpu_performance(&processor);
427        assert!(benchmark.overall_score >= 0.0);
428        assert!(benchmark.memory_bandwidth_gbps >= 0.0);
429    }
430
431    #[test]
432    fn test_maintenance_operations() {
433        let mut processor = UnifiedGpuProcessor::default();
434        let report = processor.maintenance();
435        assert!(report.is_ok());
436    }
437}