1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
use crate::gpu_internals::{GpuInfo, GpuInstance}; use once_cell::sync::Lazy; use std::sync::RwLock; static DEVICES: Lazy<GpuStore> = Lazy::new(|| { let (s, r) = std::sync::mpsc::channel(); std::thread::spawn(move || s.send(futures::executor::block_on(GpuStore::new()))); r.recv().unwrap() }); pub struct GpuStore { current: RwLock<usize>, available_devices: Vec<GpuInstance>, } impl GpuStore { pub fn get_default() -> &'static GpuInstance { let current_idx = DEVICES.current.read().unwrap(); &DEVICES.available_devices[*current_idx] } pub fn get(gpu_info: &GpuInfo) -> &'static GpuInstance { DEVICES .available_devices .iter() .find(|dev| dev.info() == gpu_info) .unwrap() } pub async fn select_gpu(gpu_info: &GpuInfo) { let idx = DEVICES .available_devices .iter() .position(|dev| dev.info() == gpu_info) .unwrap(); *(&DEVICES.current).write().unwrap() = idx; } pub fn list_gpus() -> Vec<&'static GpuInfo> { (&DEVICES) .available_devices .iter() .map(|dev| dev.info()) .collect() } async fn new() -> Self { let gpu_factory = crate::gpu_internals::gpu_factory::GpuFactory::new().await; let gpu_list = gpu_factory.list_gpus().await; let mut gpu_instances = vec![]; for gpu_info in &gpu_list { gpu_instances.push(gpu_factory.request_gpu(&gpu_info).await); } assert!(!gpu_instances.is_empty(), "No GPU detected!"); Self { current: RwLock::new(0), available_devices: gpu_instances, } } }