furiosa-opt-std 0.4.0

Standard library for Furiosa NPU TCP Virtual ISA programming.
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

#[expect(dead_code, missing_docs, reason = "bindgen output; bindgen does not emit docs")]
mod bindings {
    include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
}

pub use bindings::NpuDesc;
pub(crate) use bindings::*;

use std::sync::OnceLock;

use crate::context::Device;

pub(crate) fn rt() -> *const Runtime {
    struct Handle(*mut Runtime);
    unsafe impl Send for Handle {}
    unsafe impl Sync for Handle {}

    static RT: OnceLock<Handle> = OnceLock::new();
    RT.get_or_init(|| {
        let ver = unsafe { std::ffi::CStr::from_ptr(furiosa_version()) };
        log::info!("device-runtime: {}", ver.to_string_lossy());
        let Device { chip: chips, pe: pes } = *DEVICE.get_or_init(|| Device { chip: 1, pe: 8 });

        // Optimistic acquisition: `furiosa_runtime_init`'s exclusive open (null when a
        // chip is busy) is the atomic arbiter, so try each candidate set of `chips`
        // and take the first that acquires. No check-then-use, hence no TOCTOU race;
        // the chips are whatever the host exposes and need not be contiguous.
        let available = available_chips();
        for group in available.windows(chips as usize) {
            let npus: Vec<NpuDesc> = group
                .iter()
                .flat_map(|&chip| match pes {
                    8 => vec![
                        NpuDesc {
                            chip,
                            pe_start: 0,
                            pe_end: 3,
                        },
                        NpuDesc {
                            chip,
                            pe_start: 4,
                            pe_end: 7,
                        },
                    ],
                    n => vec![NpuDesc {
                        chip,
                        pe_start: 0,
                        pe_end: n - 1,
                    }],
                })
                .collect();
            let ptr = unsafe { furiosa_runtime_init(npus.as_ptr(), npus.len()) };
            if !ptr.is_null() {
                return Handle(ptr);
            }
        }
        panic!("no NPU available: {} chip(s) present, all in use", available.len());
    })
    .0
}

/// Chip ids the host exposes, read from the `/dev/rngd/npu<N>mgmt` device nodes.
fn available_chips() -> Vec<u8> {
    let mut chips: Vec<u8> = std::fs::read_dir("/dev/rngd")
        .into_iter()
        .flatten()
        .flatten()
        .filter_map(|entry| {
            let name = entry.file_name();
            name.to_str()?.strip_prefix("npu")?.strip_suffix("mgmt")?.parse().ok()
        })
        .collect();
    chips.sort_unstable();
    chips
}

static DEVICE: OnceLock<Device> = OnceLock::new();

/// Binds the `#[device(chip, pe)]` the process runs on, via `Context::acquire().bind(..)`.
/// First bind wins; a conflicting one panics.
pub(crate) fn bind_device(device: Device) {
    let cur = *DEVICE.get_or_init(|| device);
    assert_eq!(
        cur, device,
        "conflicting NPU device in one process: {cur:?} vs {device:?}"
    );
}