ironaccelerator-neuron 2.1.0

AWS Neuron backend for IronAccelerator — Trainium (trn1 / trn2) and Inferentia (inf1 / inf2) via libnrt.
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

//! AWS Neuron `Backend` impl. One descriptor per NeuronCore.

use ironaccelerator_core::{
    Backend, BackendKind, Capability, CapabilityFlags, ComputeTier, DeviceDescriptor, DeviceId,
    Result, Vendor,
};

use crate::drv::NeuronGen;

pub struct NeuronBackend;
pub static NEURON_BACKEND: NeuronBackend = NeuronBackend;

impl Backend for NeuronBackend {
    fn kind(&self) -> BackendKind {
        BackendKind::Neuron
    }

    fn is_available(&self) -> bool {
        crate::drv::is_available() && crate::drv::total_cores() > 0
    }

    fn enumerate(&self) -> Result<Vec<DeviceDescriptor>> {
        let cores = crate::drv::total_cores();
        if cores == 0 {
            return Ok(Vec::new());
        }
        let gen = crate::drv::detect_generation();
        let flags = capability_flags(gen);
        let arch = arch_string(gen);
        let hbm = hbm_bytes_per_core(gen);
        Ok((0..cores)
            .map(|ord| DeviceDescriptor {
                id: DeviceId {
                    backend: BackendKind::Neuron,
                    ordinal: ord,
                },
                vendor: Vendor::Aws,
                name: format!("AWS NeuronCore ({})", arch),
                arch: arch.to_string(),
                total_memory_bytes: hbm,
                multiprocessor_count: 0,
                clock_khz: 0,
                capability: Capability {
                    flags,
                    tier: ComputeTier::Datacenter,
                    fp16_tflops: None,
                    fp8_tflops: None,
                    mem_bandwidth_gbs: None,
                },
            })
            .collect())
    }

    fn capabilities(&self, device: u32) -> Result<CapabilityFlags> {
        // All NeuronCores on an instance are the same generation; validate the
        // ordinal so this agrees with `enumerate` on what exists.
        if !self.enumerate()?.iter().any(|d| d.id.ordinal == device) {
            return Err(ironaccelerator_core::Error::InvalidArgument(
                "neuron core ordinal out of range",
            ));
        }
        Ok(capability_flags(crate::drv::detect_generation()))
    }
}

fn capability_flags(gen: NeuronGen) -> CapabilityFlags {
    let base = CapabilityFlags::FP32
        | CapabilityFlags::BF16
        | CapabilityFlags::INT8
        | CapabilityFlags::TENSOR_CORES
        | CapabilityFlags::HBM
        | CapabilityFlags::MULTI_STREAM;
    match gen {
        NeuronGen::Trn2 => {
            base | CapabilityFlags::FP8_E4M3 | CapabilityFlags::FP8_E5M2 | CapabilityFlags::INT4
        }
        NeuronGen::Trn1 => base | CapabilityFlags::FP8_E4M3 | CapabilityFlags::FP8_E5M2,
        NeuronGen::Inf1 => {
            CapabilityFlags::FP32
                | CapabilityFlags::BF16
                | CapabilityFlags::INT8
                | CapabilityFlags::TENSOR_CORES
                | CapabilityFlags::MULTI_STREAM
        }
        NeuronGen::Unknown => base,
    }
}

fn arch_string(gen: NeuronGen) -> &'static str {
    match gen {
        NeuronGen::Inf1 => "neuron-v1-inferentia",
        NeuronGen::Trn1 => "neuron-v2-trainium",
        NeuronGen::Trn2 => "neuron-v3-trainium2",
        NeuronGen::Unknown => "neuron",
    }
}

fn hbm_bytes_per_core(gen: NeuronGen) -> u64 {
    const GIB: u64 = 1024 * 1024 * 1024;
    match gen {
        NeuronGen::Inf1 => 8 * GIB,
        NeuronGen::Trn1 => 16 * GIB,
        NeuronGen::Trn2 => 24 * GIB,
        NeuronGen::Unknown => 0,
    }
}