-
Notifications
You must be signed in to change notification settings - Fork 136
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
A new FFI call (`util::get_gpu_devices()`) is introduced. It returns the available GPUs as an array of strings.
- Loading branch information
Showing
5 changed files
with
87 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -8,3 +8,4 @@ extern crate log; | |
|
||
pub mod bls; | ||
pub mod proofs; | ||
pub mod util; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,47 @@ | ||
use std::ffi::CString; | ||
|
||
use bellperson::GPU_NVIDIA_DEVICES; | ||
use ffi_toolkit::{catch_panic_response, raw_ptr}; | ||
|
||
use super::types::GpuDeviceResponse; | ||
|
||
/// Returns an array of strings containing the device names that can be used. | ||
#[no_mangle] | ||
pub unsafe extern "C" fn get_gpu_devices() -> *mut GpuDeviceResponse { | ||
catch_panic_response(|| { | ||
let devices: Vec<*const i8> = GPU_NVIDIA_DEVICES | ||
.iter() | ||
.map(|device| { | ||
let name = device.name().unwrap_or("Unknown".to_string()); | ||
CString::new(&name[..]).unwrap().as_ptr() | ||
}) | ||
.collect(); | ||
let mut response = GpuDeviceResponse::default(); | ||
response.devices_len = devices.len(); | ||
response.devices_ptr = devices.as_ptr(); | ||
|
||
raw_ptr(response) | ||
}) | ||
} | ||
|
||
#[cfg(test)] | ||
mod tests { | ||
use std::ffi::CStr; | ||
use std::slice::from_raw_parts; | ||
|
||
use crate::util::api::get_gpu_devices; | ||
use crate::util::types::destroy_gpu_device_response; | ||
|
||
#[test] | ||
fn test_get_gpu_devices() { | ||
unsafe { | ||
let resp = get_gpu_devices(); | ||
let devices: Vec<&str> = from_raw_parts((*resp).devices_ptr, (*resp).devices_len) | ||
.iter() | ||
.map(|name_ptr| CStr::from_ptr(*name_ptr).to_str().unwrap()) | ||
.collect(); | ||
assert_eq!(devices.len(), (*resp).devices_len); | ||
destroy_gpu_device_response(resp); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
pub mod api; | ||
pub mod types; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,32 @@ | ||
use std::ptr; | ||
|
||
use drop_struct_macro_derive::DropStructMacro; | ||
// `CodeAndMessage` is the trait implemented by `code_and_message_impl | ||
use ffi_toolkit::{code_and_message_impl, free_c_str, CodeAndMessage, FCPResponseStatus}; | ||
|
||
#[repr(C)] | ||
#[derive(DropStructMacro)] | ||
pub struct GpuDeviceResponse { | ||
pub status_code: FCPResponseStatus, | ||
pub error_msg: *const libc::c_char, | ||
pub devices_len: libc::size_t, | ||
pub devices_ptr: *const *const i8, | ||
} | ||
|
||
impl Default for GpuDeviceResponse { | ||
fn default() -> Self { | ||
Self { | ||
error_msg: ptr::null(), | ||
status_code: FCPResponseStatus::FCPNoError, | ||
devices_len: 0, | ||
devices_ptr: ptr::null(), | ||
} | ||
} | ||
} | ||
|
||
code_and_message_impl!(GpuDeviceResponse); | ||
|
||
#[no_mangle] | ||
pub unsafe extern "C" fn destroy_gpu_device_response(ptr: *mut GpuDeviceResponse) { | ||
let _ = Box::from_raw(ptr); | ||
} |