himada-integrations 0.1.0

//! Burn backend — Himada dispatches as a Burn compute backend.
//!
//! Burn is a Rust-native deep learning framework with a pluggable
//! Backend trait. This module provides operations backed by
//! Himada's runtime-optimized SIMD dispatch.

use himada_dispatch::*;

/// Himada-backed Burn compute backend.
pub struct HimadaBurnBackend {
    dot_dispatch: Dispatch<DotKernel>,
    dot_f32_dispatch: Dispatch<DotF32Kernel>,
    matmul_dispatch: Dispatch<MatMulKernel>,
    matmul_f32_dispatch: Dispatch<MatMulF32Kernel>,
}

impl HimadaBurnBackend {
    pub fn new() -> Self {
        Self {
            dot_dispatch: Dispatch::new("burn_dot", vec![]),
            dot_f32_dispatch: Dispatch::new("burn_dot_f32", vec![]),
            matmul_dispatch: Dispatch::new("burn_matmul", vec![]),
            matmul_f32_dispatch: Dispatch::new("burn_matmul_f32", vec![]),
        }
    }

    pub fn register_kernels(&mut self, kernels: Vec<KernelInfo<DotKernel>>) {
        self.dot_dispatch = Dispatch::new("burn_dot", kernels);
    }

    pub fn register_kernels_f32(&mut self, kernels: Vec<KernelInfo<DotF32Kernel>>) {
        self.dot_f32_dispatch = Dispatch::new("burn_dot_f32", kernels);
    }

    pub fn register_matmul(&mut self, kernels: Vec<KernelInfo<MatMulKernel>>) {
        self.matmul_dispatch = Dispatch::new("burn_matmul", kernels);
    }

    /// Dot product (f64) via Himada dispatch.
    pub fn dot(&mut self, a: &[f64], b: &[f64]) -> f64 {
        self.dot_dispatch.compute(a, b)
    }

    /// Dot product (f32) via Himada dispatch.
    pub fn dot_f32(&mut self, a: &[f32], b: &[f32]) -> f32 {
        self.dot_f32_dispatch.compute(a, b)
    }

    /// Matrix multiply via Himada dispatch.
    pub fn matmul(&mut self, a: &[f64], b: &[f64], c: &mut [f64], n: usize) {
        self.matmul_dispatch.compute(a, b, c, n)
    }

    /// Optimize all dispatches (select best kernels).
    pub fn optimize(&mut self) -> Result<usize, HwdnaError> {
        let mut count = 0;
        if self.dot_dispatch.select().is_ok() { count += 1; }
        if self.dot_f32_dispatch.select().is_ok() { count += 1; }
        if self.matmul_dispatch.select().is_ok() { count += 1; }
        if self.matmul_f32_dispatch.select().is_ok() { count += 1; }
        Ok(count)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use himada_dispatch::kernels;
    use himada_core::HardwareDNA;

    #[test]
    fn test_burn_dot_consistency() {
        let mut backend = HimadaBurnBackend::new();
        backend.register_kernels(vec![KernelInfo {
            name: "scalar",
            func: kernels::dot_scalar as DotKernel,
            is_supported: |_: &HardwareDNA| true,
            thermal_priority: 0,
        }]);
        backend.register_matmul(vec![KernelInfo {
            name: "scalar",
            func: kernels::matmul_scalar as MatMulKernel,
            is_supported: |_: &HardwareDNA| true,
            thermal_priority: 0,
        }]);
        backend.optimize().unwrap();

        let a: Vec<f64> = (0..200).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..200).map(|i| (i * 5) as f64).collect();

        let himada_result = backend.dot(&a, &b);
        let expected: f64 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
        assert!((himada_result - expected).abs() < 1e-10);

        // Matmul test
        let n = 16;
        let ma: Vec<f64> = (0..n*n).map(|i| (i % 10) as f64).collect();
        let mb: Vec<f64> = (0..n*n).map(|i| ((i * 2) % 10) as f64).collect();
        let mut mc = vec![0.0; n*n];
        backend.matmul(&ma, &mb, &mut mc, n);

        let mut expected_c = vec![0.0; n*n];
        for i in 0..n {
            for j in 0..n {
                let mut s = 0.0;
                for k in 0..n {
                    s += ma[i * n + k] * mb[k * n + j];
                }
                expected_c[i * n + j] = s;
            }
        }
        for i in 0..n*n {
            assert!((mc[i] - expected_c[i]).abs() < 1e-10);
        }
    }

    #[test]
    fn test_burn_f32() {
        let mut backend = HimadaBurnBackend::new();
        backend.register_kernels_f32(vec![KernelInfo {
            name: "scalar",
            func: kernels::dot_f32_scalar as DotF32Kernel,
            is_supported: |_: &HardwareDNA| true,
            thermal_priority: 0,
        }]);
        backend.optimize().unwrap();

        let a: Vec<f32> = (0..100).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..100).map(|i| (i * 2) as f32).collect();
        let result = backend.dot_f32(&a, &b);
        let expected: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
        assert!((result - expected).abs() < 1e-6);
    }
}