neuronx 0.1.0

Bindings for the AWS Neuron runtime (libnrt).
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

use std::env;
use std::path::PathBuf;

fn env_var_rerun(name: &str) -> Option<String> {
    println!("cargo:rerun-if-env-changed={name}");
    env::var(name).ok()
}

fn main() {
    let neuron_dir =
        env_var_rerun("NEURON_DIR").map_or_else(|| PathBuf::from("/opt/aws/neuron"), PathBuf::from);
    let include_dir = neuron_dir.join("include");
    if !include_dir.join("nrt/nrt.h").exists() {
        panic!(
            "cannot find nrt/nrt.h in {}, install aws-neuronx-runtime-lib or set NEURON_DIR",
            include_dir.display()
        )
    }

    println!("cargo:rerun-if-changed=wrapper.h");
    let bindings = bindgen::Builder::default()
        .header("wrapper.h")
        .clang_arg(format!("-I{}", include_dir.display()))
        .allowlist_function("nrta?_.*")
        .allowlist_type("nrta?_.*|NRTA?_.*")
        .allowlist_var("NRTA?_.*")
        // C enums are kept as plain integer constants: libnrt may return values that
        // were added after these bindings were generated.
        .prepend_enum_name(false)
        .derive_default(true)
        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
        .generate()
        .expect("Unable to generate bindings");
    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
    bindings.write_to_file(out_path.join("nrt.rs")).expect("Couldn't write bindings!");

    let lib_dir = neuron_dir.join("lib");
    println!("cargo:rustc-link-search=native={}", lib_dir.display());
    println!("cargo:rustc-link-arg=-Wl,-rpath={}", lib_dir.display());
    println!("cargo:rustc-link-lib=nrt");
}