tensor_compute 0.1.0

A Tensor computing library for Rust
Documentation 
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

use crate::gpu_internals::{GpuInfo, GpuInstance};

/// Should be used for querying the available GPUs and instantiating a [GpuInstance] in
/// order to be able to interact with them. Does not need to be kept alive after a [GpuInstance].
/// ** Having multiple GpuInstances referring to the same physical device might have unexpected
/// consequences! **
pub struct GpuFactory {
    adapters: Vec<wgpu::Adapter>,
}

impl GpuFactory {
    pub async fn new() -> Self {
        let wgpu_instance = wgpu::Instance::new(wgpu::BackendBit::PRIMARY);
        let adapters: Vec<wgpu::Adapter> = wgpu_instance
            .enumerate_adapters(wgpu::BackendBit::PRIMARY)
            .collect();
        GpuFactory { adapters }
    }

    pub async fn list_gpus(&self) -> Vec<GpuInfo> {
        self.adapters
            .iter()
            .map(|adapter| adapter.get_info())
            .collect()
    }

    /// ** Having multiple GpuInstances referring to the same physical device might have unexpected
    /// consequences! **
    pub async fn request_gpu(&self, gpu_info: &GpuInfo) -> GpuInstance {
        let adapter = self
            .adapters
            .iter()
            .find(|adapter| &adapter.get_info() == gpu_info)
            .expect("Adapter does not exist for given GpuInfo");
        let (device, queue) = adapter
            .request_device(
                &wgpu::DeviceDescriptor {
                    features: wgpu::Features::empty(),
                    limits: wgpu::Limits::default(),
                    shader_validation: true,
                },
                Some(std::path::Path::new("traces")),
            )
            .await
            .unwrap();
        GpuInstance {
            device,
            queue,
            info: gpu_info.clone(),
        }
    }
}