use native_neural_network_std as nn;
use nn::modules::activations::ActivationKind;
use nn::modules::engine::{get_compute_backend, set_compute_backend, ComputeBackend};
use nn::std::layers_std::{DenseLayerDesc, LayerSpec};
use nn::std::rnn_std as rstd;
use std::path::{Path, PathBuf};
use std::process::Command;
mod utils {
include!("test_utils.rs");
}
fn make_input_csv(count: usize, seed: u32) -> String {
let mut rng = utils::XorShift32::new(seed);
let mut values = Vec::with_capacity(count);
for i in 0..count {
let base = (((i * 31) % 1000) as f32) / 500.0 - 1.0;
let noise = rng.next_f32() * 0.05;
let value = utils::saturating_add_f32(base, noise);
values.push(format!("{value:.6}"));
}
values.join(",")
}
fn run_example(name: &str, envs: &[(&str, String)]) {
let mut cmd = Command::new("cargo");
cmd.args(["run", "--quiet", "--example", name]);
cmd.env("NNS_OUTPUT_DIR", "rnn_examples/__tmp_templates");
for (k, v) in envs {
cmd.env(k, v);
}
let out = cmd
.output()
.unwrap_or_else(|e| panic!("failed to run example {name}: {e}"));
if !out.status.success() {
panic!(
"example {name} failed\nstdout:\n{}\nstderr:\n{}",
String::from_utf8_lossy(&out.stdout),
String::from_utf8_lossy(&out.stderr)
);
}
}
fn prepare_layout() {
let root = Path::new("rnn_examples");
let _ = std::fs::remove_dir_all(root);
let _ = std::fs::remove_dir_all("examples/generated_models");
std::fs::create_dir_all(root.join("f32").join("rmd")).expect("create rnn_examples/f32/rmd");
std::fs::create_dir_all(root.join("f32").join("rnn")).expect("create rnn_examples/f32/rnn");
std::fs::create_dir_all(root.join("f64").join("rmd")).expect("create rnn_examples/f64/rmd");
std::fs::create_dir_all(root.join("f64").join("rnn")).expect("create rnn_examples/f64/rnn");
}
fn run_templates() {
run_example(
"simple",
&[("NNS_SIMPLE_INPUT_VALUES", make_input_csv(24, 0x1001))],
);
run_example(
"medium",
&[("NNS_MEDIUM_INPUT_VALUES", make_input_csv(128, 0x2002))],
);
run_example(
"complex",
&[
(
"NNS_COMPLEX_NET_A_INPUT_VALUES",
make_input_csv(144, 0x3001),
),
(
"NNS_COMPLEX_NET_B_INPUT_VALUES",
make_input_csv(144, 0x3002),
),
(
"NNS_COMPLEX_NET_C_INPUT_VALUES",
make_input_csv(144, 0x3003),
),
],
);
run_example(
"very_complex",
&[("NNS_VERY_COMPLEX_INPUT_VALUES", make_input_csv(192, 0x4004))],
);
let generated_root = Path::new("rnn_examples/__tmp_templates");
assert!(
generated_root.join("simple.rnn").exists(),
"missing simple.rnn"
);
assert!(
generated_root.join("medium.rnn").exists(),
"missing medium.rnn"
);
assert!(
generated_root.join("complex.rnn").exists(),
"missing complex.rnn"
);
assert!(
generated_root.join("very_complex.rnn").exists(),
"missing very_complex.rnn"
);
}
fn topology_from_layers(layers: &[LayerSpec]) -> Vec<usize> {
let mut topology = Vec::new();
if let Some(LayerSpec::Dense(first)) = layers.first() {
topology.push(first.input_size);
}
for layer in layers {
let LayerSpec::Dense(d) = layer;
topology.push(d.output_size);
}
topology
}
fn write_variant_artifacts(name: &str) {
let source = PathBuf::from("rnn_examples/__tmp_templates").join(format!("{name}.rnn"));
let container_bytes = std::fs::read(&source)
.unwrap_or_else(|e| panic!("read template artifact {} failed: {e}", source.display()));
let decode_any = |bytes: &[u8]| -> Option<(Vec<LayerSpec>, Vec<f32>, Vec<f32>)> {
if let Ok((layers, weights, biases)) = rstd::rnn_unpack_v1_std(bytes) {
return Some((layers, weights, biases));
}
if let Ok((layers, weights, biases)) = rstd::rnn_unpack_v1_f64_std(bytes) {
let w32: Vec<f32> = weights.iter().map(|v| *v as f32).collect();
let b32: Vec<f32> = biases.iter().map(|v| *v as f32).collect();
return Some((layers, w32, b32));
}
None
};
let extracted_model = rstd::rnn_extract_blob_bytes_std(&container_bytes, "model").ok();
let (payload_f32, layers, weights_f32, biases_f32) =
if let Some((layers, weights, biases)) = decode_any(&container_bytes) {
(container_bytes.clone(), layers, weights, biases)
} else if let Some(payload) = extracted_model {
if let Some((layers, weights, biases)) = decode_any(&payload) {
(payload, layers, weights, biases)
} else {
let root = Path::new("rnn_examples");
let fallback = container_bytes.clone();
std::fs::write(
root.join("f32").join("rmd").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f32 rmd artifact");
std::fs::write(
root.join("f32").join("rnn").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f32 rnn artifact");
std::fs::write(
root.join("f64").join("rmd").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f64 rmd artifact");
std::fs::write(
root.join("f64").join("rnn").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f64 rnn artifact");
return;
}
} else {
let root = Path::new("rnn_examples");
let fallback = container_bytes.clone();
std::fs::write(
root.join("f32").join("rmd").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f32 rmd artifact");
std::fs::write(
root.join("f32").join("rnn").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f32 rnn artifact");
std::fs::write(
root.join("f64").join("rmd").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f64 rmd artifact");
std::fs::write(
root.join("f64").join("rnn").join(format!("{name}.rnn")),
&fallback,
)
.expect("write fallback f64 rnn artifact");
return;
};
let topology = topology_from_layers(&layers);
let root = Path::new("rnn_examples");
std::fs::write(
root.join("f32").join("rmd").join(format!("{name}.rnn")),
&payload_f32,
)
.expect("write f32 rmd artifact");
let container_f32 = rstd::rnn_wrap_payload_in_container_std(&payload_f32, "model");
std::fs::write(
root.join("f32").join("rnn").join(format!("{name}.rnn")),
&container_f32,
)
.expect("write f32 rnn artifact");
let weights_f64: Vec<f64> = weights_f32.iter().map(|v| *v as f64).collect();
let biases_f64: Vec<f64> = biases_f32.iter().map(|v| *v as f64).collect();
let counts =
rstd::rnn_required_from_topology(&topology).expect("required counts from topology");
let mut layer_scratch = vec![
LayerSpec::Dense(DenseLayerDesc {
input_size: 0,
output_size: 0,
weight_offset: 0,
bias_offset: 0,
activation: ActivationKind::Identity,
});
counts.layers
];
let mut out_bytes64 =
vec![0u8; (weights_f64.len() + biases_f64.len()) * 8 + counts.layers * 64 + 4096];
let written64 = loop {
match rstd::rnn_pack_v1_f64(
&topology,
ActivationKind::Identity,
ActivationKind::Identity,
&weights_f64,
&biases_f64,
&mut layer_scratch,
&mut out_bytes64,
) {
Ok(w) => break w,
Err(native_neural_network_std::native::rnn_api::RnnApiError::CapacityTooSmall) => {
let new_len = out_bytes64.len().saturating_mul(2).max(8192);
out_bytes64.resize(new_len, 0);
}
Err(e) => panic!("pack f64 payload: {:?}", e),
}
};
let payload_f64 = &out_bytes64[..written64];
std::fs::write(
root.join("f64").join("rmd").join(format!("{name}.rnn")),
payload_f64,
)
.expect("write f64 rmd artifact");
let container_f64 = rstd::rnn_wrap_payload_in_container_std(payload_f64, "model");
std::fs::write(
root.join("f64").join("rnn").join(format!("{name}.rnn")),
&container_f64,
)
.expect("write f64 rnn artifact");
}
fn count_rnn_files(root: &Path) -> usize {
fn visit(p: &Path, total: &mut usize) {
if p.is_dir() {
for entry in std::fs::read_dir(p)
.unwrap_or_else(|e| panic!("read_dir {} failed: {e}", p.display()))
{
let entry =
entry.unwrap_or_else(|e| panic!("dir entry {} failed: {e}", p.display()));
visit(&entry.path(), total);
}
} else if p.extension().and_then(|s| s.to_str()) == Some("rnn") {
*total += 1;
}
}
let mut total = 0usize;
visit(root, &mut total);
total
}
#[test]
fn cpu_backend_abstraction_real() {
let previous = get_compute_backend();
set_compute_backend(ComputeBackend::Cpu);
assert_eq!(get_compute_backend(), ComputeBackend::Cpu);
set_compute_backend(previous);
}
#[test]
fn generate_rnn_examples_from_templates() {
prepare_layout();
run_templates();
for name in ["simple", "medium", "complex", "very_complex"] {
write_variant_artifacts(name);
}
let _ = std::fs::remove_dir_all("rnn_examples/__tmp_templates");
let root = Path::new("rnn_examples");
let total = count_rnn_files(root);
assert_eq!(
total, 16,
"expected 16 .rnn artifacts in rnn_examples, got {total}"
);
}