ai-hwaccel 1.1.0

Universal AI hardware accelerator detection, capability querying, and workload planning 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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

//! AcceleratorRequirement tests.

use crate::*;

#[test]
fn requirement_satisfied_by_all_types() {
    let cuda = AcceleratorProfile::cuda(0, 24 * 1024 * 1024 * 1024);
    let tpu = AcceleratorProfile::tpu(0, 4, TpuVersion::V5p);
    let cpu = AcceleratorProfile::cpu(16 * 1024 * 1024 * 1024);
    let gaudi = AcceleratorProfile::gaudi(0, GaudiGeneration::Gaudi3);
    let neuron = AcceleratorProfile {
        accelerator: AcceleratorType::AwsNeuron {
            device_id: 0,
            chip_type: NeuronChipType::Trainium,
            core_count: 2,
        },
        available: true,
        memory_bytes: 32 * 1024 * 1024 * 1024,
        compute_capability: None,
        driver_version: None,
        device_name: None,
        memory_bandwidth_gbps: None,
        memory_used_bytes: None,
        memory_free_bytes: None,
        pcie_bandwidth_gbps: None,
        numa_node: None,
        temperature_c: None,
        power_watts: None,
        gpu_utilization_percent: None,
    };
    let npu = AcceleratorProfile {
        accelerator: AcceleratorType::IntelNpu,
        available: true,
        memory_bytes: 2 * 1024 * 1024 * 1024,
        compute_capability: None,
        driver_version: None,
        device_name: None,
        memory_bandwidth_gbps: None,
        memory_used_bytes: None,
        memory_free_bytes: None,
        pcie_bandwidth_gbps: None,
        numa_node: None,
        temperature_c: None,
        power_watts: None,
        gpu_utilization_percent: None,
    };

    // None satisfied by anything
    assert!(AcceleratorRequirement::None.satisfied_by(&cuda));
    assert!(AcceleratorRequirement::None.satisfied_by(&cpu));
    assert!(AcceleratorRequirement::None.satisfied_by(&gaudi));
    assert!(AcceleratorRequirement::None.satisfied_by(&neuron));

    // GPU requirement
    assert!(AcceleratorRequirement::Gpu.satisfied_by(&cuda));
    assert!(!AcceleratorRequirement::Gpu.satisfied_by(&tpu));
    assert!(!AcceleratorRequirement::Gpu.satisfied_by(&cpu));
    assert!(!AcceleratorRequirement::Gpu.satisfied_by(&gaudi));
    assert!(!AcceleratorRequirement::Gpu.satisfied_by(&npu));

    // TPU requirement
    assert!(AcceleratorRequirement::Tpu { min_chips: 2 }.satisfied_by(&tpu));
    assert!(!AcceleratorRequirement::Tpu { min_chips: 8 }.satisfied_by(&tpu));
    assert!(!AcceleratorRequirement::Tpu { min_chips: 1 }.satisfied_by(&cuda));

    // Gaudi requirement
    assert!(AcceleratorRequirement::Gaudi.satisfied_by(&gaudi));
    assert!(!AcceleratorRequirement::Gaudi.satisfied_by(&cuda));
    assert!(!AcceleratorRequirement::Gaudi.satisfied_by(&tpu));
    assert!(!AcceleratorRequirement::Gaudi.satisfied_by(&cpu));

    // AwsNeuron requirement
    assert!(AcceleratorRequirement::AwsNeuron.satisfied_by(&neuron));
    assert!(!AcceleratorRequirement::AwsNeuron.satisfied_by(&cuda));
    assert!(!AcceleratorRequirement::AwsNeuron.satisfied_by(&gaudi));

    // GpuOrTpu
    assert!(AcceleratorRequirement::GpuOrTpu.satisfied_by(&cuda));
    assert!(AcceleratorRequirement::GpuOrTpu.satisfied_by(&tpu));
    assert!(!AcceleratorRequirement::GpuOrTpu.satisfied_by(&cpu));
    assert!(!AcceleratorRequirement::GpuOrTpu.satisfied_by(&gaudi));

    // AnyAccelerator
    assert!(AcceleratorRequirement::AnyAccelerator.satisfied_by(&cuda));
    assert!(AcceleratorRequirement::AnyAccelerator.satisfied_by(&tpu));
    assert!(AcceleratorRequirement::AnyAccelerator.satisfied_by(&gaudi));
    assert!(AcceleratorRequirement::AnyAccelerator.satisfied_by(&neuron));
    assert!(AcceleratorRequirement::AnyAccelerator.satisfied_by(&npu));
    assert!(!AcceleratorRequirement::AnyAccelerator.satisfied_by(&cpu));
}

#[test]
fn requirement_unavailable_device_never_satisfies() {
    let mut unavailable = AcceleratorProfile::cuda(0, 24 * 1024 * 1024 * 1024);
    unavailable.available = false;
    assert!(!AcceleratorRequirement::Gpu.satisfied_by(&unavailable));
    assert!(!AcceleratorRequirement::AnyAccelerator.satisfied_by(&unavailable));
    assert!(!AcceleratorRequirement::GpuOrTpu.satisfied_by(&unavailable));
    // None still requires available
    assert!(!AcceleratorRequirement::None.satisfied_by(&unavailable));
}

#[test]
fn requirement_default_is_none() {
    let req: AcceleratorRequirement = Default::default();
    assert_eq!(req, AcceleratorRequirement::None);
}