himada-dispatch 0.1.1

//! JIT compilation of specialized kernels for specific input sizes.
//!
//! For a known shape (e.g., dot 1024, matmul 512×512), we can generate
//! optimized code at runtime instead of using the generic SIMD path.
//! This avoids the dispatch overhead and loop-trip overhead for fixed N.
//!
//! # Strategy (Cranelift)
//!
//! 1. Build a Cranelift `FunctionBuilder` with the target ISA (native CPU
//!    features detected via `std::arch` / `raw_cpuid`).
//! 2. For **dot product**: generate a single basic block that loops over
//!    `[0..N)` with explicit SIMD loads (e.g., `vld1q_f64` on NEON or
//!    `vmovupd` + `vmulpd` + `vaddpd` on AVX2). Unroll factor = 4 or 8
//!    depending on N.
//! 3. For **matmul**: generate a tiled triple-loop with tile sizes chosen
//!    for L1 cache (e.g., 64×64 tiles). Use SIMD for the inner micro-kernel.
//! 4. Emit the function pointer via `JITModule::finalize` / `finalize_definitions`.
//! 5. Cache by `(op_name, N)` — reuse on subsequent calls.
//!
//! # Status
//!
//! Stub — the actual Cranelift integration requires adding `cranelift-jit`
//! and `cranelift-module` as dependencies. The methods below outline the
//! intended structure.

use crate::HwdnaError;

/// A compiled JIT kernel.
pub struct JitKernel {
    name: String,
    #[allow(dead_code)]
    code: Vec<u8>,
}

impl JitKernel {
    pub fn new(name: &str, code: Vec<u8>) -> Self {
        Self { name: name.into(), code }
    }

    pub fn name(&self) -> &str {
        &self.name
    }
}

/// JIT engine — compiles and caches kernels.
pub struct JitEngine {
    cache: Vec<JitKernel>,
}

impl JitEngine {
    pub fn new() -> Self {
        Self { cache: Vec::new() }
    }

    /// Compile a specialized dot product kernel for the given size.
    ///
    /// # Cranelift IR outline (future implementation)
    ///
    /// ```ignore
    /// // 1. Create `Builder` with `isa::lookup(...).finish()`
    /// // 2. Define function signature: `(a: *const f64, b: *const f64) -> f64`
    /// // 3. Build loop:
    /// //    block0(v0: i64 = 0):
    /// //      if v0 >= N: br block_exit
    /// //      va = load.f64x4 a[v0..v0+4]
    /// //      vb = load.f64x4 b[v0..v0+4]
    /// //      vprod = fmul va, vb
    /// //      vacc = fadd vacc, vprod
    /// //      v0 = iadd v0, 4
    /// //      br block0
/// //    block_exit:
    /// //      return horizontal_add(vacc)
    /// // 4. Compile → `*const fn(a: *const f64, b: *const f64) -> f64`
    /// // 5. Store pointer in `code` (as bytes of the function).
    /// ```
    ///
    /// For now this is a no-op stub that records the cache entry.
    /// Returns the index of the kernel in the cache.
    pub fn compile_dot(&mut self, n: usize) -> Result<usize, HwdnaError> {
        let name = format!("dot_{}", n);
        if let Some(idx) = self.cache.iter().position(|k| k.name() == name) {
            return Ok(idx);
        }
        let code = Vec::new();
        self.cache.push(JitKernel::new(&name, code));
        Ok(self.cache.len() - 1)
    }

    /// Compile a specialized matmul kernel for N×N.
    ///
    /// # Cranelift IR outline (future implementation)
    ///
    /// ```ignore
    /// // Signature: `(a: *const f64, b: *const f64, c: *mut f64)`
    /// // Tiled triple-loop:
    /// //   for ti in 0..N step 64:
    /// //     for tj in 0..N step 64:
    /// //       for tk in 0..N step 64:
    /// //         microkernel(a[ti*N+tk], b[tk*N+tj], c[ti*N+tj], 64)
    /// // microkernel is a 4×4 or 8×8 SIMD kernel fully unrolled.
    /// ```
    ///
    /// Returns the index of the kernel in the cache.
    pub fn compile_matmul(&mut self, n: usize) -> Result<usize, HwdnaError> {
        let name = format!("matmul_{}x{}", n, n);
        if let Some(idx) = self.cache.iter().position(|k| k.name() == name) {
            return Ok(idx);
        }
        let code = Vec::new();
        self.cache.push(JitKernel::new(&name, code));
        Ok(self.cache.len() - 1)
    }

    pub fn cache_size(&self) -> usize {
        self.cache.len()
    }
}