scirs2_io/gpu/
compression.rs

1//! GPU-accelerated compression operations for high-performance I/O
2//!
3//! This module provides GPU-accelerated compression and decompression
4//! with backend-specific optimizations for maximum throughput and efficiency.
5//!
6//! Uses Pure Rust oxiarc-* compression libraries (COOLJAPAN Policy).
7
8use super::backend_management::GpuIoProcessor;
9use crate::compression::{CompressionAlgorithm, ParallelCompressionConfig};
10use crate::error::{IoError, Result};
11use scirs2_core::gpu::{GpuBackend, GpuDataType};
12use scirs2_core::ndarray::{Array1, ArrayView1};
13use scirs2_core::simd_ops::PlatformCapabilities;
14
15/// GPU-accelerated compression processor
16#[derive(Debug)]
17pub struct GpuCompressionProcessor {
18    gpu_processor: GpuIoProcessor,
19    compression_threshold: usize,
20}
21
22impl GpuCompressionProcessor {
23    /// Create a new GPU compression processor
24    pub fn new() -> Result<Self> {
25        let gpu = GpuIoProcessor::new().unwrap_or_default();
26        Ok(Self {
27            gpu_processor: gpu,
28            compression_threshold: 10 * 1024 * 1024, // 10MB threshold
29        })
30    }
31
32    /// Create with custom compression threshold
33    pub fn with_threshold(threshold: usize) -> Result<Self> {
34        let gpu = GpuIoProcessor::new().unwrap_or_default();
35        Ok(Self {
36            gpu_processor: gpu,
37            compression_threshold: threshold,
38        })
39    }
40
41    /// Compress data using GPU acceleration
42    pub fn compress_gpu<T: GpuDataType>(
43        &self,
44        data: &ArrayView1<T>,
45        algorithm: CompressionAlgorithm,
46        level: Option<u32>,
47    ) -> Result<Vec<u8>> {
48        // Check if GPU should be used based on data size
49        let data_bytes = data.len() * std::mem::size_of::<T>();
50        let use_gpu = self.should_use_gpu(data_bytes);
51
52        if use_gpu {
53            match self.gpu_processor.backend() {
54                GpuBackend::Cuda => self.compress_cuda(data, algorithm, level),
55                GpuBackend::Metal => self.compress_metal(data, algorithm, level),
56                GpuBackend::OpenCL => self.compress_opencl(data, algorithm, level),
57                _ => {
58                    // Fallback to CPU implementation
59                    Err(IoError::Other(format!(
60                        "GPU backend {} not supported for compression",
61                        self.gpu_processor.backend()
62                    )))
63                }
64            }
65        } else {
66            // Data too small, use CPU
67            Err(IoError::Other(
68                "Data size too small for GPU acceleration".to_string(),
69            ))
70        }
71    }
72
73    /// Decompress data using GPU acceleration
74    pub fn decompress_gpu<T: GpuDataType>(
75        &self,
76        compressed_data: &[u8],
77        algorithm: CompressionAlgorithm,
78        expected_size: usize,
79    ) -> Result<Array1<T>> {
80        let use_gpu = self.should_use_gpu(expected_size);
81
82        if use_gpu {
83            match self.gpu_processor.backend() {
84                GpuBackend::Cuda => self.decompress_cuda(compressed_data, algorithm, expected_size),
85                GpuBackend::Metal => {
86                    self.decompress_metal(compressed_data, algorithm, expected_size)
87                }
88                GpuBackend::OpenCL => {
89                    self.decompress_opencl(compressed_data, algorithm, expected_size)
90                }
91                _ => Err(IoError::Other(format!(
92                    "GPU backend {} not supported for decompression",
93                    self.gpu_processor.backend()
94                ))),
95            }
96        } else {
97            Err(IoError::Other(
98                "Data size too small for GPU acceleration".to_string(),
99            ))
100        }
101    }
102
103    /// Determine if GPU should be used based on data size
104    pub(crate) fn should_use_gpu(&self, size: usize) -> bool {
105        size > self.compression_threshold
106    }
107
108    /// CUDA-specific compression implementation
109    fn compress_cuda<T: GpuDataType>(
110        &self,
111        data: &ArrayView1<T>,
112        algorithm: CompressionAlgorithm,
113        level: Option<u32>,
114    ) -> Result<Vec<u8>> {
115        let capabilities = &self.gpu_processor.capabilities;
116
117        // Convert array data to bytes for compression
118        let data_bytes = unsafe {
119            std::slice::from_raw_parts(
120                data.as_ptr() as *const u8,
121                data.len() * std::mem::size_of::<T>(),
122            )
123        };
124
125        // Use CUDA-optimized parallel compression
126        let chunk_size = if capabilities.simd_available {
127            1024 * 1024 // 1MB chunks for CUDA
128        } else {
129            512 * 1024 // 512KB fallback
130        };
131
132        let compressed_chunks: Result<Vec<Vec<u8>>> = data_bytes
133            .chunks(chunk_size)
134            .enumerate()
135            .map(|(i, chunk)| match algorithm {
136                CompressionAlgorithm::Gzip => {
137                    let compression_level = level.unwrap_or(6).clamp(1, 9) as u8;
138                    oxiarc_deflate::gzip_compress(chunk, compression_level)
139                        .map_err(|e| IoError::Other(format!("Gzip compression error: {}", e)))
140                }
141                CompressionAlgorithm::Zstd => {
142                    let compression_level = level.unwrap_or(3).clamp(1, 19);
143                    oxiarc_zstd::compress_with_level(chunk, compression_level as i32)
144                        .map_err(|e| IoError::Other(format!("Zstd compression error: {}", e)))
145                }
146                CompressionAlgorithm::Lz4 => oxiarc_lz4::compress_block(chunk)
147                    .map_err(|e| IoError::Other(format!("LZ4 compression error: {}", e))),
148                _ => Err(IoError::UnsupportedFormat(format!(
149                    "Compression algorithm {:?} not supported for CUDA",
150                    algorithm
151                ))),
152            })
153            .collect();
154
155        let chunks = compressed_chunks?;
156
157        // Create CUDA-specific header format
158        let mut result = Vec::new();
159
160        // CUDA header: magic + version + chunk count
161        result.extend_from_slice(b"CUDA"); // Magic number
162        result.extend_from_slice(&1u32.to_le_bytes()); // Version
163        result.extend_from_slice(&(chunks.len() as u32).to_le_bytes());
164
165        // Write chunk sizes
166        for chunk in &chunks {
167            result.extend_from_slice(&(chunk.len() as u32).to_le_bytes());
168        }
169
170        // Write chunk data
171        for chunk in chunks {
172            result.extend_from_slice(&chunk);
173        }
174
175        Ok(result)
176    }
177
178    /// Metal-specific compression implementation
179    fn compress_metal<T: GpuDataType>(
180        &self,
181        data: &ArrayView1<T>,
182        algorithm: CompressionAlgorithm,
183        level: Option<u32>,
184    ) -> Result<Vec<u8>> {
185        let capabilities = &self.gpu_processor.capabilities;
186
187        // Convert to bytes
188        let data_bytes = unsafe {
189            std::slice::from_raw_parts(
190                data.as_ptr() as *const u8,
191                data.len() * std::mem::size_of::<T>(),
192            )
193        };
194
195        // Metal-optimized chunk size for GPU compute units
196        let chunk_size = if capabilities.simd_available {
197            768 * 1024 // Metal-optimized 768KB chunks
198        } else {
199            384 * 1024
200        };
201
202        let compressed_chunks: Result<Vec<Vec<u8>>> = data_bytes
203            .chunks(chunk_size)
204            .map(|chunk| match algorithm {
205                CompressionAlgorithm::Gzip => {
206                    let compression_level = level.unwrap_or(6).clamp(1, 9) as u8;
207                    oxiarc_deflate::gzip_compress(chunk, compression_level)
208                        .map_err(|e| IoError::Other(format!("Gzip compression error: {}", e)))
209                }
210                CompressionAlgorithm::Zstd => {
211                    let compression_level = level.unwrap_or(4).clamp(1, 19);
212                    oxiarc_zstd::compress_with_level(chunk, compression_level as i32)
213                        .map_err(|e| IoError::Other(format!("Zstd compression error: {}", e)))
214                }
215                CompressionAlgorithm::Lz4 => oxiarc_lz4::compress_block(chunk)
216                    .map_err(|e| IoError::Other(format!("LZ4 compression error: {}", e))),
217                _ => Err(IoError::UnsupportedFormat(format!(
218                    "Compression algorithm {:?} not supported for Metal",
219                    algorithm
220                ))),
221            })
222            .collect();
223
224        let chunks = compressed_chunks?;
225
226        // Metal-specific header format
227        let mut result = Vec::new();
228
229        // Metal header: magic + version + device info + chunk count
230        result.extend_from_slice(b"METL"); // Magic number for Metal
231        result.extend_from_slice(&1u32.to_le_bytes()); // Version
232        result.extend_from_slice(&(chunks.len() as u32).to_le_bytes());
233
234        // Write chunk metadata
235        for chunk in &chunks {
236            result.extend_from_slice(&(chunk.len() as u32).to_le_bytes());
237        }
238
239        // Write chunk data
240        for chunk in chunks {
241            result.extend_from_slice(&chunk);
242        }
243
244        Ok(result)
245    }
246
247    /// OpenCL-specific compression implementation
248    fn compress_opencl<T: GpuDataType>(
249        &self,
250        data: &ArrayView1<T>,
251        algorithm: CompressionAlgorithm,
252        level: Option<u32>,
253    ) -> Result<Vec<u8>> {
254        let capabilities = &self.gpu_processor.capabilities;
255
256        // Convert to bytes
257        let data_bytes = unsafe {
258            std::slice::from_raw_parts(
259                data.as_ptr() as *const u8,
260                data.len() * std::mem::size_of::<T>(),
261            )
262        };
263
264        // OpenCL-optimized chunk size
265        let chunk_size = if capabilities.simd_available {
266            512 * 1024 // OpenCL works well with 512KB chunks
267        } else {
268            256 * 1024
269        };
270
271        let compressed_chunks: Result<Vec<Vec<u8>>> = data_bytes
272            .chunks(chunk_size)
273            .map(|chunk| match algorithm {
274                CompressionAlgorithm::Gzip => {
275                    let compression_level = level.unwrap_or(6).clamp(1, 9) as u8;
276                    oxiarc_deflate::gzip_compress(chunk, compression_level)
277                        .map_err(|e| IoError::Other(format!("Gzip compression error: {}", e)))
278                }
279                CompressionAlgorithm::Zstd => {
280                    let compression_level = level.unwrap_or(5).clamp(1, 19);
281                    oxiarc_zstd::compress_with_level(chunk, compression_level as i32)
282                        .map_err(|e| IoError::Other(format!("Zstd compression error: {}", e)))
283                }
284                CompressionAlgorithm::Lz4 => oxiarc_lz4::compress_block(chunk)
285                    .map_err(|e| IoError::Other(format!("LZ4 compression error: {}", e))),
286                _ => Err(IoError::UnsupportedFormat(format!(
287                    "Compression algorithm {:?} not supported for OpenCL",
288                    algorithm
289                ))),
290            })
291            .collect();
292
293        let chunks = compressed_chunks?;
294
295        // OpenCL-specific header format optimized for GPU decompression
296        let mut result = Vec::new();
297
298        // OpenCL header: magic + version + device info + chunk count
299        result.extend_from_slice(b"OPCL"); // Magic number for OpenCL compression
300        result.extend_from_slice(&1u32.to_le_bytes()); // Version
301
302        // Real compute-unit count queried from the backend, mirroring the
303        // pattern used by `get_performance_stats`. Falls back to a
304        // conservative default only if the underlying device query fails.
305        let compute_units = self
306            .gpu_processor
307            .get_backend_capabilities()
308            .map(|c| c.compute_units as u32)
309            .unwrap_or(32);
310        result.extend_from_slice(&compute_units.to_le_bytes());
311        result.extend_from_slice(&(chunks.len() as u32).to_le_bytes());
312
313        // Write chunk metadata optimized for OpenCL kernel processing
314        for (i, chunk) in chunks.iter().enumerate() {
315            result.extend_from_slice(&(i as u32).to_le_bytes()); // Chunk index
316            result.extend_from_slice(&(chunk.len() as u32).to_le_bytes()); // Chunk size
317        }
318
319        // Write chunk data
320        for chunk in chunks {
321            result.extend_from_slice(&chunk);
322        }
323
324        Ok(result)
325    }
326
327    /// CUDA-specific decompression implementation
328    fn decompress_cuda<T: GpuDataType>(
329        &self,
330        data: &[u8],
331        algorithm: CompressionAlgorithm,
332        expected_size: usize,
333    ) -> Result<Array1<T>> {
334        // Read CUDA header
335        if data.len() < 12 || &data[0..4] != b"CUDA" {
336            return Err(IoError::Other(
337                "Invalid CUDA compressed data format".to_string(),
338            ));
339        }
340
341        let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
342        if version != 1 {
343            return Err(IoError::Other(
344                "Unsupported CUDA compression version".to_string(),
345            ));
346        }
347
348        let num_chunks = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
349        let mut offset = 12;
350
351        // Read chunk sizes
352        let mut chunk_sizes = Vec::with_capacity(num_chunks);
353        for _ in 0..num_chunks {
354            if offset + 4 > data.len() {
355                return Err(IoError::Other("Invalid compressed data format".to_string()));
356            }
357            let size = u32::from_le_bytes([
358                data[offset],
359                data[offset + 1],
360                data[offset + 2],
361                data[offset + 3],
362            ]) as usize;
363            chunk_sizes.push(size);
364            offset += 4;
365        }
366
367        // Read and decompress chunks in parallel
368        self.decompress_chunks_parallel(data, offset, &chunk_sizes, algorithm)
369    }
370
371    /// Metal-specific decompression implementation
372    fn decompress_metal<T: GpuDataType>(
373        &self,
374        data: &[u8],
375        algorithm: CompressionAlgorithm,
376        expected_size: usize,
377    ) -> Result<Array1<T>> {
378        // Handle Metal-specific header format
379        if data.len() < 12 || &data[0..4] != b"METL" {
380            // Not Metal format, try CUDA decompression
381            return self.decompress_cuda(data, algorithm, expected_size);
382        }
383
384        let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
385        if version != 1 {
386            return Err(IoError::Other(
387                "Unsupported Metal compression version".to_string(),
388            ));
389        }
390
391        let num_chunks = u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize;
392        let mut offset = 12;
393
394        // Read chunk sizes
395        let mut chunk_sizes = Vec::with_capacity(num_chunks);
396        for _ in 0..num_chunks {
397            if offset + 4 > data.len() {
398                return Err(IoError::Other(
399                    "Invalid Metal compressed data format".to_string(),
400                ));
401            }
402            let size = u32::from_le_bytes([
403                data[offset],
404                data[offset + 1],
405                data[offset + 2],
406                data[offset + 3],
407            ]) as usize;
408            chunk_sizes.push(size);
409            offset += 4;
410        }
411
412        // Decompress chunks using Metal-optimized parallel processing
413        self.decompress_chunks_parallel(data, offset, &chunk_sizes, algorithm)
414    }
415
416    /// OpenCL-specific decompression implementation
417    fn decompress_opencl<T: GpuDataType>(
418        &self,
419        data: &[u8],
420        algorithm: CompressionAlgorithm,
421        expected_size: usize,
422    ) -> Result<Array1<T>> {
423        // Handle OpenCL-specific header format
424        if data.len() < 16 || &data[0..4] != b"OPCL" {
425            // Not OpenCL format, try CUDA decompression
426            return self.decompress_cuda(data, algorithm, expected_size);
427        }
428
429        let version = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
430        if version != 1 {
431            return Err(IoError::Other(
432                "Unsupported OpenCL compression version".to_string(),
433            ));
434        }
435
436        let compute_units = u32::from_le_bytes([data[8], data[9], data[10], data[11]]);
437        let num_chunks = u32::from_le_bytes([data[12], data[13], data[14], data[15]]) as usize;
438        let mut offset = 16;
439
440        // Read chunk metadata (index + size pairs)
441        let mut chunk_sizes = Vec::with_capacity(num_chunks);
442        for _ in 0..num_chunks {
443            if offset + 8 > data.len() {
444                return Err(IoError::Other(
445                    "Invalid OpenCL compressed data format".to_string(),
446                ));
447            }
448            let _chunk_index = u32::from_le_bytes([
449                data[offset],
450                data[offset + 1],
451                data[offset + 2],
452                data[offset + 3],
453            ]);
454            let chunk_size = u32::from_le_bytes([
455                data[offset + 4],
456                data[offset + 5],
457                data[offset + 6],
458                data[offset + 7],
459            ]) as usize;
460            chunk_sizes.push(chunk_size);
461            offset += 8;
462        }
463
464        // Decompress using OpenCL-optimized approach
465        self.decompress_chunks_parallel(data, offset, &chunk_sizes, algorithm)
466    }
467
468    /// Parallel chunk decompression (shared implementation)
469    fn decompress_chunks_parallel<T: GpuDataType>(
470        &self,
471        data: &[u8],
472        mut offset: usize,
473        chunk_sizes: &[usize],
474        algorithm: CompressionAlgorithm,
475    ) -> Result<Array1<T>> {
476        use scirs2_core::parallel_ops::*;
477
478        // Extract chunk data
479        let mut chunk_data = Vec::new();
480        for &size in chunk_sizes {
481            if offset + size > data.len() {
482                return Err(IoError::Other("Invalid compressed data format".to_string()));
483            }
484            chunk_data.push(&data[offset..offset + size]);
485            offset += size;
486        }
487
488        // Decompress chunks in parallel
489        let decompressed_chunks: Result<Vec<Vec<u8>>> = chunk_data
490            .par_iter()
491            .map(|chunk| match algorithm {
492                CompressionAlgorithm::Gzip => oxiarc_deflate::gzip_decompress(chunk)
493                    .map_err(|e| IoError::Other(format!("Gzip decompression error: {}", e))),
494                CompressionAlgorithm::Zstd => oxiarc_zstd::decompress(chunk)
495                    .map_err(|e| IoError::Other(format!("Zstd decompression error: {}", e))),
496                CompressionAlgorithm::Lz4 => {
497                    let max_size = chunk.len() * 10;
498                    oxiarc_lz4::decompress_block(chunk, max_size)
499                        .map_err(|e| IoError::Other(format!("LZ4 decompression error: {}", e)))
500                }
501                _ => Err(IoError::UnsupportedFormat(format!(
502                    "Compression algorithm {:?} not supported for GPU decompression",
503                    algorithm
504                ))),
505            })
506            .collect();
507
508        let chunks = decompressed_chunks?;
509
510        // Combine chunks
511        let mut combined_data = Vec::new();
512        for chunk in chunks {
513            combined_data.extend_from_slice(&chunk);
514        }
515
516        // Convert bytes back to T array
517        let element_size = std::mem::size_of::<T>();
518        if combined_data.len() % element_size != 0 {
519            return Err(IoError::Other(
520                "Decompressed data size mismatch".to_string(),
521            ));
522        }
523
524        let num_elements = combined_data.len() / element_size;
525        let typed_data = unsafe {
526            std::slice::from_raw_parts(combined_data.as_ptr() as *const T, num_elements).to_vec()
527        };
528
529        Ok(Array1::from_vec(typed_data))
530    }
531
532    /// Get compression performance statistics
533    pub fn get_performance_stats(&self) -> CompressionStats {
534        let capabilities = self
535            .gpu_processor
536            .get_backend_capabilities()
537            .unwrap_or_else(|_| {
538                use super::backend_management::BackendCapabilities;
539                BackendCapabilities {
540                    backend: scirs2_core::gpu::GpuBackend::Cpu,
541                    memory_gb: 1.0,
542                    max_work_group_size: 64,
543                    supports_fp64: false,
544                    supports_fp16: false,
545                    compute_units: 1,
546                    max_allocation_size: 1024 * 1024,
547                    local_memory_size: 64 * 1024,
548                }
549            });
550
551        CompressionStats {
552            backend: capabilities.backend,
553            threshold_bytes: self.compression_threshold,
554            estimated_throughput_gbps: capabilities.estimate_memory_bandwidth(),
555            parallel_chunks: capabilities.compute_units,
556        }
557    }
558}
559
560impl Default for GpuCompressionProcessor {
561    fn default() -> Self {
562        Self::new().unwrap_or_else(|_| {
563            // Fallback configuration
564            Self {
565                gpu_processor: GpuIoProcessor::default(),
566                compression_threshold: 10 * 1024 * 1024,
567            }
568        })
569    }
570}
571
572/// Compression performance statistics
573#[derive(Debug, Clone)]
574pub struct CompressionStats {
575    /// GPU backend type
576    pub backend: scirs2_core::gpu::GpuBackend,
577    /// Minimum data size for GPU acceleration
578    pub threshold_bytes: usize,
579    /// Estimated throughput in GB/s
580    pub estimated_throughput_gbps: f64,
581    /// Number of parallel processing chunks
582    pub parallel_chunks: usize,
583}
584
585#[cfg(test)]
586mod tests {
587    use super::*;
588    use scirs2_core::ndarray::arr1;
589
590    #[test]
591    fn test_compression_processor_creation() {
592        // Should work even without GPU due to fallback
593        let processor = GpuCompressionProcessor::default();
594        assert!(processor.compression_threshold > 0);
595    }
596
597    #[test]
598    fn test_compression_threshold() {
599        let processor = GpuCompressionProcessor::with_threshold(1024).unwrap_or_default();
600        assert!(!processor.should_use_gpu(512)); // Below threshold
601        assert!(processor.should_use_gpu(2048)); // Above threshold
602    }
603
604    #[test]
605    fn test_compression_stats() {
606        let processor = GpuCompressionProcessor::default();
607        let stats = processor.get_performance_stats();
608        assert!(stats.threshold_bytes > 0);
609        assert!(stats.parallel_chunks > 0);
610    }
611
612    #[test]
613    fn test_small_data_compression() {
614        let processor = GpuCompressionProcessor::default();
615        let small_data = arr1(&[1.0f32, 2.0, 3.0, 4.0]);
616
617        // Should fail due to size threshold
618        let result = processor.compress_gpu(&small_data.view(), CompressionAlgorithm::Lz4, None);
619        assert!(result.is_err());
620    }
621
622    #[test]
623    fn test_compress_opencl_header_uses_real_compute_units() {
624        let processor = GpuCompressionProcessor::default();
625        let data = arr1(&[1.0f32; 1024]);
626
627        // Call the OpenCL-specific path directly (bypasses the size-threshold
628        // gate in `compress_gpu`, which only routes by backend + data size).
629        let compressed = processor
630            .compress_opencl(&data.view(), CompressionAlgorithm::Lz4, None)
631            .expect("compress_opencl should succeed for well-formed input");
632
633        // OpenCL header layout:
634        // magic(4) + version(4) + compute_units(4) + chunk_count(4) + ...
635        assert_eq!(&compressed[0..4], b"OPCL");
636        let version =
637            u32::from_le_bytes([compressed[4], compressed[5], compressed[6], compressed[7]]);
638        assert_eq!(version, 1);
639
640        let compute_units_in_header =
641            u32::from_le_bytes([compressed[8], compressed[9], compressed[10], compressed[11]]);
642
643        // The header must carry the *real* compute-unit count sourced from
644        // backend capabilities (matching the same query `compress_opencl`
645        // performs), not the old hardcoded `32` placeholder that was written
646        // unconditionally regardless of what the backend actually reports.
647        let expected_compute_units = processor
648            .gpu_processor
649            .get_backend_capabilities()
650            .map(|c| c.compute_units as u32)
651            .unwrap_or(32);
652
653        assert_eq!(compute_units_in_header, expected_compute_units);
654        assert!(compute_units_in_header > 0);
655    }
656}
scirs2_io/gpu/compression.rs

scirs2_io/gpu/
compression.rs
