92 lines
2.8 KiB
Rust
92 lines
2.8 KiB
Rust
use crate::error::{self, Result};
|
|
use crate::ffi;
|
|
use std::ffi::CStr;
|
|
use std::os::raw::c_char;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct DeviceInfo {
|
|
pub index: u32,
|
|
pub name: String,
|
|
pub total_memory: usize,
|
|
pub free_memory: usize,
|
|
pub compute_major: i32,
|
|
pub compute_minor: i32,
|
|
pub sm_count: i32,
|
|
pub shared_mem_per_block: usize,
|
|
pub warp_size: i32,
|
|
pub max_threads_per_block: i32,
|
|
}
|
|
|
|
unsafe extern "C" {
|
|
fn cudaDeviceGetAttribute(value: *mut i32, attr: i32, device: i32) -> i32;
|
|
fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32;
|
|
}
|
|
|
|
fn get_attr(attr: i32, device: u32) -> Result<i32> {
|
|
let mut value = 0;
|
|
error::check(unsafe { cudaDeviceGetAttribute(&mut value, attr, device as i32) })?;
|
|
Ok(value)
|
|
}
|
|
|
|
pub fn device_count() -> Result<i32> {
|
|
let mut count = 0;
|
|
error::check(unsafe { ffi::cudaGetDeviceCount(&mut count) })?;
|
|
Ok(count)
|
|
}
|
|
|
|
pub fn set_device(device: u32) -> Result<()> {
|
|
error::check(unsafe { ffi::cudaSetDevice(device as i32) })
|
|
}
|
|
|
|
pub fn current_device() -> Result<u32> {
|
|
let mut dev = 0;
|
|
error::check(unsafe { ffi::cudaGetDevice(&mut dev) })?;
|
|
Ok(dev as u32)
|
|
}
|
|
|
|
pub fn device_info(device: u32) -> Result<DeviceInfo> {
|
|
// Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
|
|
// CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
|
|
let mut prop_buf = vec![0u8; 32768];
|
|
error::check(unsafe { ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32) })?;
|
|
// Name is always the first field: char[256].
|
|
let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
|
|
.to_string_lossy()
|
|
.into_owned();
|
|
|
|
// Get memory info via cudaMemGetInfo (layout-independent).
|
|
let prev = current_device()?;
|
|
set_device(device)?;
|
|
let mut free = 0usize;
|
|
let mut total = 0usize;
|
|
error::check(unsafe { cudaMemGetInfo(&mut free, &mut total) })?;
|
|
if prev != device {
|
|
set_device(prev)?;
|
|
}
|
|
|
|
// Attribute IDs from cuda_runtime_api.h
|
|
const SHARED_MEM_PER_BLOCK: i32 = 8;
|
|
const WARP_SIZE: i32 = 10;
|
|
const MAX_THREADS_PER_BLOCK: i32 = 1;
|
|
const MULTI_PROCESSOR_COUNT: i32 = 16;
|
|
const COMPUTE_MAJOR: i32 = 75;
|
|
const COMPUTE_MINOR: i32 = 76;
|
|
|
|
Ok(DeviceInfo {
|
|
index: device,
|
|
name,
|
|
total_memory: total,
|
|
free_memory: free,
|
|
compute_major: get_attr(COMPUTE_MAJOR, device)?,
|
|
compute_minor: get_attr(COMPUTE_MINOR, device)?,
|
|
sm_count: get_attr(MULTI_PROCESSOR_COUNT, device)?,
|
|
shared_mem_per_block: get_attr(SHARED_MEM_PER_BLOCK, device)? as usize,
|
|
warp_size: get_attr(WARP_SIZE, device)?,
|
|
max_threads_per_block: get_attr(MAX_THREADS_PER_BLOCK, device)?,
|
|
})
|
|
}
|
|
|
|
pub fn synchronize() -> Result<()> {
|
|
error::check(unsafe { ffi::cudaDeviceSynchronize() })
|
|
}
|