native_neural_network_std 0.2.1

use native_neural_network_std::std::activations_std::ActivationKind;
use native_neural_network_std::std::engine_std::forward_plan;
use native_neural_network_std::std::layers_std::{DenseLayerDesc, LayerPlanStd, LayerSpec};
use native_neural_network_std::std::model_std::ModelStd;
use native_neural_network_std::std::rnn_std::rnn_required_infer_scratch_from_specs_std;
use std::sync::atomic::{AtomicBool, Ordering};

static GPU_DENSE_HANDLER_CALLED: AtomicBool = AtomicBool::new(false);

fn gpu_identity_handler(
    src: &[f32],
    dst: &mut [f32],
    batch_size: usize,
    stride: usize,
    in_size: usize,
    out_size: usize,
    weights: &[f32],
    biases: &[f32],
    _activation: ActivationKind,
) -> bool {
    GPU_DENSE_HANDLER_CALLED.store(true, Ordering::SeqCst);

    if in_size != 2 || out_size != 2 || stride < 2 {
        return false;
    }
    if weights.len() < 4
        || biases.len() < 2
        || src.len() < batch_size * stride
        || dst.len() < batch_size * stride
    {
        return false;
    }

    for b in 0..batch_size {
        let base = b * stride;
        let x0 = src[base];
        let x1 = src[base + 1];
        dst[base] = x0 * weights[0] + x1 * weights[1] + biases[0];
        dst[base + 1] = x0 * weights[2] + x1 * weights[3] + biases[1];
    }
    true
}

struct BackendStateGuard {
    previous: native_neural_network_std::std::engine_std::ComputeBackend,
}

impl BackendStateGuard {
    fn new() -> Self {
        let previous = native_neural_network_std::std::engine_std::get_compute_backend();
        Self { previous }
    }
}

impl Drop for BackendStateGuard {
    fn drop(&mut self) {
        native_neural_network_std::std::engine_std::set_compute_backend(self.previous);
    }
}

fn make_identity_2x2() -> ModelStd {
    let layers = vec![LayerSpec::Dense(DenseLayerDesc {
        input_size: 2,
        output_size: 2,
        weight_offset: 0,
        bias_offset: 0,
        activation: ActivationKind::Identity,
    })];
    let weights = vec![1.0f32, 0.0f32, 0.0f32, 1.0f32];
    let biases = vec![0.0f32, 0.0f32];
    let infer_scratch_len = rnn_required_infer_scratch_from_specs_std(&layers).unwrap_or(0);
    ModelStd {
        layers,
        weights,
        biases,
        input_size: 2,
        output_size: 2,
        infer_scratch_len,
    }
}

#[test]
fn cpu_backend_abstraction_real() {
    let previous = native_neural_network_std::std::engine_std::get_compute_backend();
    native_neural_network_std::std::engine_std::set_compute_backend(
        native_neural_network_std::std::engine_std::ComputeBackend::Cpu,
    );
    assert_eq!(
        native_neural_network_std::std::engine_std::get_compute_backend(),
        native_neural_network_std::std::engine_std::ComputeBackend::Cpu
    );
    native_neural_network_std::std::engine_std::set_compute_backend(previous);
}

#[test]
fn batch_matches_single_and_deterministic() {
    let model = make_identity_2x2();
    let input = vec![1.0f32, 2.0f32, 3.0f32, 4.0f32];
    let batch_out = model.run_batch(&input, 2).expect("batch");
    let mut singles = Vec::new();
    for r in 0..2 {
        let off = r * 2;
        let s = model.run_single(&input[off..off + 2]).expect("single");
        singles.extend_from_slice(&s);
    }
    assert_eq!(batch_out, singles);
    let first = model.run_batch(&input, 2).unwrap();
    for _ in 0..50 {
        assert_eq!(model.run_batch(&input, 2).unwrap(), first);
    }
    let plan = LayerPlanStd::new(
        model.layers.clone(),
        model.weights.clone(),
        model.biases.clone(),
    );
    let req = native_neural_network_std::std::engine_std::required_single_infer_scratch(&plan)
        .unwrap_or(0);
    assert!(req <= model.infer_scratch_len + 1024);
}

#[test]
fn relu_activation_runs_without_regression() {
    let layers = vec![LayerSpec::Dense(DenseLayerDesc {
        input_size: 2,
        output_size: 2,
        weight_offset: 0,
        bias_offset: 0,
        activation: ActivationKind::Relu,
    })];
    let weights = vec![1.0f32, -1.0f32, -1.0f32, 1.0f32];
    let biases = vec![0.0f32, 0.0f32];
    let infer_scratch_len = rnn_required_infer_scratch_from_specs_std(&layers).unwrap_or(0);
    let model = ModelStd {
        layers,
        weights,
        biases,
        input_size: 2,
        output_size: 2,
        infer_scratch_len,
    };
    let out = model.run_single(&[2.0f32, 1.0f32]).expect("run");
    assert_eq!(out.len(), 2);
    assert!(out.iter().all(|v| v.is_finite()));
}

#[test]
fn forward_plan_ok_for_matching_sizes() {
    let model = make_identity_2x2();
    let plan = LayerPlanStd::new(
        model.layers.clone(),
        model.weights.clone(),
        model.biases.clone(),
    );
    let mut out = vec![0f32; model.output_size];
    let mut scratch = vec![0f32; model.infer_scratch_len.max(16)];
    let res = forward_plan(&plan, &[1.0f32, 2.0f32], &mut out, &mut scratch);
    assert!(res.is_ok());
}

#[test]
fn gpu_backend_dispatches_native_handler_without_extra_dependencies() {
    let _guard = BackendStateGuard::new();

    GPU_DENSE_HANDLER_CALLED.store(false, Ordering::SeqCst);
    native_neural_network_std::std::engine_std::set_compute_backend(
        native_neural_network_std::std::engine_std::ComputeBackend::Gpu,
    );
    native_neural_network_std::std::engine_std::register_gpu_kernel_f32(gpu_identity_handler);

    let model = make_identity_2x2();
    let plan = LayerPlanStd::new(
        model.layers.clone(),
        model.weights.clone(),
        model.biases.clone(),
    );
    let mut out = vec![0f32; model.output_size];
    let mut scratch = vec![0f32; model.infer_scratch_len.max(16)];

    let res = forward_plan(&plan, &[3.0f32, 4.0f32], &mut out, &mut scratch);
    assert!(res.is_ok());
    assert!(GPU_DENSE_HANDLER_CALLED.load(Ordering::SeqCst));
    assert_eq!(out, vec![3.0f32, 4.0f32]);
}