morok-device 0.1.0-alpha.2

Device abstraction layer for the Morok ML compiler
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

use crate::{DeviceSpec, registry::DeviceSpecExt};

#[test]
fn test_device_spec_parse() {
    assert_eq!(DeviceSpec::parse("CPU").unwrap(), DeviceSpec::Cpu);
    assert_eq!(DeviceSpec::parse("cpu").unwrap(), DeviceSpec::Cpu);

    #[cfg(feature = "cuda")]
    {
        assert_eq!(DeviceSpec::parse("CUDA:0").unwrap(), DeviceSpec::Cuda { device_id: 0 });
        assert_eq!(DeviceSpec::parse("cuda").unwrap(), DeviceSpec::Cuda { device_id: 0 });
        assert_eq!(DeviceSpec::parse("GPU:2").unwrap(), DeviceSpec::Cuda { device_id: 2 });
    }
}

#[test]
fn test_device_spec_canonicalize() {
    assert_eq!(DeviceSpec::Cpu.canonicalize(), "CPU");

    #[cfg(feature = "cuda")]
    {
        assert_eq!(DeviceSpec::Cuda { device_id: 1 }.canonicalize(), "CUDA:1");
    }
}