// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! Smoke tests: build a tiny graph, compile it through the CPU
//! backend, run it. This verifies the runtime path works end-to-end
//! before we drive any descriptor-specific parity.

use rlx_ir::infer::GraphExt;
use rlx_ir::{DType, Graph, Shape};
use rlx_runtime::{Device, Session};

#[test]
fn matmul_runs_on_cpu() {
    let mut g = Graph::new("smoke");
    let x = g.input("x", Shape::new(&[2, 3], DType::F32));
    let w = g.param("w", Shape::new(&[3, 4], DType::F32));
    let y = g.mm(x, w);
    g.set_outputs(vec![y]);

    let session = Session::new(Device::Cpu);
    let mut compiled = session.compile(g);

    // Set weights so y = x · I-like matrix (with mismatching dims) — pick a
    // pattern that's easy to verify by hand.
    let w_data: Vec<f32> = (0..12).map(|i| i as f32).collect();
    let w_bytes: Vec<u8> = w_data.iter().flat_map(|v| v.to_le_bytes()).collect();
    compiled.set_param_typed("w", &w_bytes, DType::F32);

    let x_data: Vec<f32> = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0];
    let out = compiled.run(&[("x", x_data.as_slice())]).pop().unwrap();
    assert_eq!(out.len(), 2 * 4);
    // First row of y = first row of W = [0,1,2,3]
    for j in 0..4 {
        assert!((out[j] - j as f32).abs() < 1e-6);
    }
    // Second row of y = second row of W = [4,5,6,7]
    for j in 0..4 {
        assert!((out[4 + j] - (4 + j) as f32).abs() < 1e-6);
    }
}
