ironaccelerator_neuron/
backend.rs

1//! AWS Neuron `Backend` impl. One descriptor per NeuronCore.
2
3use ironaccelerator_core::{
4    Backend, BackendKind, Capability, CapabilityFlags, ComputeTier, DeviceDescriptor, DeviceId,
5    Result, Vendor,
6};
7
8use crate::drv::NeuronGen;
9
10/// AWS Neuron backend. Reports one device per NeuronCore.
11pub struct NeuronBackend;
12/// Process-wide Neuron backend instance.
13pub static NEURON_BACKEND: NeuronBackend = NeuronBackend;
14
15impl Backend for NeuronBackend {
16    fn kind(&self) -> BackendKind {
17        BackendKind::Neuron
18    }
19
20    fn is_available(&self) -> bool {
21        crate::drv::is_available() && crate::drv::total_cores() > 0
22    }
23
24    fn enumerate(&self) -> Result<Vec<DeviceDescriptor>> {
25        let cores = crate::drv::total_cores();
26        if cores == 0 {
27            return Ok(Vec::new());
28        }
29        let gen = crate::drv::detect_generation();
30        let flags = capability_flags(gen);
31        let arch = arch_string(gen);
32        let hbm = hbm_bytes_per_core(gen);
33        Ok((0..cores)
34            .map(|ord| DeviceDescriptor {
35                id: DeviceId {
36                    backend: BackendKind::Neuron,
37                    ordinal: ord,
38                },
39                vendor: Vendor::Aws,
40                name: format!("AWS NeuronCore ({})", arch),
41                arch: arch.to_string(),
42                total_memory_bytes: hbm,
43                multiprocessor_count: 0,
44                clock_khz: 0,
45                capability: Capability {
46                    flags,
47                    tier: ComputeTier::Datacenter,
48                    fp16_tflops: None,
49                    fp8_tflops: None,
50                    mem_bandwidth_gbs: None,
51                },
52            })
53            .collect())
54    }
55
56    fn capabilities(&self, device: u32) -> Result<CapabilityFlags> {
57        // All NeuronCores on an instance are the same generation; validate the
58        // ordinal so this agrees with `enumerate` on what exists.
59        if !self.enumerate()?.iter().any(|d| d.id.ordinal == device) {
60            return Err(ironaccelerator_core::Error::InvalidArgument(
61                "neuron core ordinal out of range",
62            ));
63        }
64        Ok(capability_flags(crate::drv::detect_generation()))
65    }
66}
67
68fn capability_flags(gen: NeuronGen) -> CapabilityFlags {
69    let base = CapabilityFlags::FP32
70        | CapabilityFlags::BF16
71        | CapabilityFlags::INT8
72        | CapabilityFlags::TENSOR_CORES
73        | CapabilityFlags::HBM
74        | CapabilityFlags::MULTI_STREAM;
75    match gen {
76        NeuronGen::Trn2 => {
77            base | CapabilityFlags::FP8_E4M3 | CapabilityFlags::FP8_E5M2 | CapabilityFlags::INT4
78        }
79        NeuronGen::Trn1 => base | CapabilityFlags::FP8_E4M3 | CapabilityFlags::FP8_E5M2,
80        NeuronGen::Inf1 => {
81            CapabilityFlags::FP32
82                | CapabilityFlags::BF16
83                | CapabilityFlags::INT8
84                | CapabilityFlags::TENSOR_CORES
85                | CapabilityFlags::MULTI_STREAM
86        }
87        NeuronGen::Unknown => base,
88    }
89}
90
91fn arch_string(gen: NeuronGen) -> &'static str {
92    match gen {
93        NeuronGen::Inf1 => "neuron-v1-inferentia",
94        NeuronGen::Trn1 => "neuron-v2-trainium",
95        NeuronGen::Trn2 => "neuron-v3-trainium2",
96        NeuronGen::Unknown => "neuron",
97    }
98}
99
100fn hbm_bytes_per_core(gen: NeuronGen) -> u64 {
101    const GIB: u64 = 1024 * 1024 * 1024;
102    match gen {
103        NeuronGen::Inf1 => 8 * GIB,
104        NeuronGen::Trn1 => 16 * GIB,
105        NeuronGen::Trn2 => 24 * GIB,
106        NeuronGen::Unknown => 0,
107    }
108}
ironaccelerator_neuron/backend.rs

ironaccelerator_neuron/
backend.rs
