oxirs_vec/gpu/
device.rs

1//! GPU device information and management
2
3#[cfg(feature = "cuda")]
4use anyhow::anyhow;
5use anyhow::Result;
6
7/// GPU device information
8#[derive(Debug, Clone)]
9pub struct GpuDevice {
10    pub device_id: i32,
11    pub name: String,
12    pub compute_capability: (i32, i32),
13    pub total_memory: usize,
14    pub free_memory: usize,
15    pub max_threads_per_block: i32,
16    pub max_blocks_per_grid: i32,
17    pub warp_size: i32,
18    pub memory_bandwidth: f32,
19    pub peak_flops: f64,
20}
21
22impl GpuDevice {
23    /// Get information about a specific GPU device
24    pub fn get_device_info(device_id: i32) -> Result<Self> {
25        #[cfg(feature = "cuda")]
26        {
27            use cuda_runtime_sys::*;
28            unsafe {
29                let result = cudaSetDevice(device_id);
30                if result != cudaError_t::cudaSuccess {
31                    return Err(anyhow!("Failed to set CUDA device {}", device_id));
32                }
33
34                let mut props: cudaDeviceProp = std::mem::zeroed();
35                let result = cudaGetDeviceProperties(&mut props, device_id);
36                if result != cudaError_t::cudaSuccess {
37                    return Err(anyhow!("Failed to get device properties"));
38                }
39
40                let mut free_mem: usize = 0;
41                let mut total_mem: usize = 0;
42                let result = cudaMemGetInfo(&mut free_mem, &mut total_mem);
43                if result != cudaError_t::cudaSuccess {
44                    return Err(anyhow!("Failed to get memory info"));
45                }
46
47                Ok(Self {
48                    device_id,
49                    name: std::ffi::CStr::from_ptr(props.name.as_ptr())
50                        .to_string_lossy()
51                        .to_string(),
52                    compute_capability: (props.major, props.minor),
53                    total_memory: total_mem,
54                    free_memory: free_mem,
55                    max_threads_per_block: props.maxThreadsPerBlock,
56                    max_blocks_per_grid: props.maxGridSize[0],
57                    warp_size: props.warpSize,
58                    memory_bandwidth: props.memoryBusWidth as f32
59                        * props.memoryClockRate as f32
60                        * 2.0
61                        / 8.0
62                        / 1e6,
63                    peak_flops: props.clockRate as f64
64                        * props.multiProcessorCount as f64
65                        * props.maxThreadsPerMultiProcessor as f64
66                        / 1e6,
67                })
68            }
69        }
70
71        #[cfg(not(feature = "cuda"))]
72        {
73            // Fallback for testing without CUDA
74            Ok(Self {
75                device_id,
76                name: format!("Simulated GPU {device_id}"),
77                compute_capability: (7, 5), // Simulate modern GPU
78                total_memory: 8 * 1024 * 1024 * 1024, // 8GB
79                free_memory: 6 * 1024 * 1024 * 1024, // 6GB free
80                max_threads_per_block: 1024,
81                max_blocks_per_grid: 65535,
82                warp_size: 32,
83                memory_bandwidth: 900.0, // GB/s
84                peak_flops: 14000.0,     // GFLOPS
85            })
86        }
87    }
88
89    /// Get information about all available GPU devices
90    pub fn get_all_devices() -> Result<Vec<Self>> {
91        #[cfg(feature = "cuda")]
92        {
93            use cuda_runtime_sys::*;
94            unsafe {
95                let mut device_count: i32 = 0;
96                let result = cudaGetDeviceCount(&mut device_count);
97                if result != cudaError_t::cudaSuccess {
98                    return Err(anyhow!("Failed to get device count"));
99                }
100
101                let mut devices = Vec::new();
102                for i in 0..device_count {
103                    if let Ok(device) = Self::get_device_info(i) {
104                        devices.push(device);
105                    }
106                }
107                Ok(devices)
108            }
109        }
110
111        #[cfg(not(feature = "cuda"))]
112        {
113            // Fallback: simulate 2 GPUs for testing
114            Ok(vec![Self::get_device_info(0)?, Self::get_device_info(1)?])
115        }
116    }
117
118    /// Check if this device supports a specific compute capability
119    pub fn supports_compute_capability(&self, major: i32, minor: i32) -> bool {
120        self.compute_capability.0 > major
121            || (self.compute_capability.0 == major && self.compute_capability.1 >= minor)
122    }
123
124    /// Get theoretical peak memory bandwidth in GB/s
125    pub fn peak_memory_bandwidth(&self) -> f32 {
126        self.memory_bandwidth
127    }
128
129    /// Get theoretical peak compute performance in GFLOPS
130    pub fn peak_compute_performance(&self) -> f64 {
131        self.peak_flops
132    }
133
134    /// Calculate optimal thread block configuration for given problem size
135    pub fn calculate_optimal_block_config(&self, problem_size: usize) -> (i32, i32) {
136        let optimal_threads = (self.max_threads_per_block as f32 * 0.75) as i32; // Use 75% of max
137        let blocks_needed = ((problem_size as f32) / (optimal_threads as f32)).ceil() as i32;
138        let blocks = blocks_needed.min(self.max_blocks_per_grid);
139        (blocks, optimal_threads)
140    }
141}
oxirs_vec/gpu/device.rs

oxirs_vec/gpu/
device.rs
