rlx_runtime/backend/

mod.rs

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Backend trait — abstraction over CPU/GPU/CUDA execution.
//!
//! Each backend implements `Backend::compile(graph, &CompileOptions)` and
//! returns an `ExecutableGraph`. New compile knobs go in `CompileOptions`
//! rather than as new trait methods.

use crate::CompileOptions;
use rlx_ir::Graph;
use rlx_ir::hir::HirModule;
use rlx_ir::lir::LirModule;
// Used only by `cpu`-gated backend machinery (every real backend implies `cpu`);
// gated so the backend-less build stays warning-clean.
#[cfg(feature = "cpu")]
use std::collections::HashMap;
#[cfg(feature = "cpu")]
use std::sync::Arc;

use crate::cpu_low_precision;

// ── Typed I/O helpers (shared across f32-arena backends) ────────────────

/// Widen (or, for F64, narrow) a typed byte buffer to `Vec<f32>`. Used by
/// `set_param_typed` / `run_typed` overrides on backends whose internal
/// arena is f32-uniform (Vulkan / wgpu / CUDA / ROCm / oneAPI / TPU) so
/// callers can hand in F16/BF16/F64 without doing the host-side cast
/// themselves. F64 is narrowed to f32 to match the f32-uniform arena slot
/// (SPD-manifold ops re-widen to f64 inside their CPU host-fallback, so
/// the eigendecomposition itself stays f64). Panics on dtypes the f32
/// arena can't carry.
///
/// `data` may be **any** alignment — a `&buf[1..]` subslice or a tensor at an
/// odd offset inside an mmap is 1-aligned, and `data.len() % 4 == 0` says
/// nothing about the pointer. Decoding goes through [`rlx_ir::bytes::decode_le`],
/// which borrows zero-copy when the pointer happens to be aligned for the
/// element type and decodes elementwise when it is not. `data` is interpreted
/// as little-endian, matching the `to_le_bytes` convention
/// [`narrow_f32_to_bytes`] writes with.
///
/// Reinterpreting the bytes directly (`from_raw_parts(data.as_ptr() as
/// *const f32, n)`) was undefined behaviour for exactly this reason; see the
/// `alignment_tests` module below for the regression guard.
#[allow(dead_code)]
pub fn widen_bytes_to_f32(data: &[u8], dtype: rlx_ir::DType) -> Vec<f32> {
    use rlx_ir::DType;
    // An empty typed input (e.g. an unused KV-cache slot passed as `Vec::new()`)
    // has a dangling, 1-byte-aligned pointer. Nothing to widen — return empty.
    if data.is_empty() {
        return Vec::new();
    }
    // `data` is caller-supplied and carries no alignment guarantee — it can be
    // a subslice like `&buf[1..]` or a tensor at an odd offset inside an mmap.
    // `decode_le` borrows zero-copy when it *is* aligned and decodes
    // elementwise when it is not; see `rlx_ir::bytes`.
    match dtype {
        DType::F32 => rlx_ir::bytes::decode_le_vec::<f32>(data),
        DType::F64 => rlx_ir::bytes::decode_le::<f64>(data)
            .iter()
            .map(|&x| x as f32)
            .collect(),
        DType::F16 => {
            // `half::f16` is `repr(transparent)` over `u16`, so decoding the
            // raw bits keeps the aligned fast path zero-copy while leaving
            // `rlx-ir` free of a `half` dependency. `from_bits` is a no-op.
            let bits = rlx_ir::bytes::decode_le::<u16>(data);
            let bits = bits.as_ref();
            // Parallel for large buffers (e.g. the MXFP4 expert scales: ~124M
            // elements/layer — serial was the measured ~1.9s upload). Elementwise
            // + order-independent → bit-identical to the serial map.
            if bits.len() >= 1 << 20 {
                use rayon::prelude::*;
                bits.par_iter()
                    .map(|&b| half::f16::from_bits(b).to_f32())
                    .collect()
            } else {
                bits.iter()
                    .map(|&b| half::f16::from_bits(b).to_f32())
                    .collect()
            }
        }
        DType::BF16 => {
            let bits = rlx_ir::bytes::decode_le::<u16>(data);
            let bits = bits.as_ref();
            if bits.len() >= 1 << 20 {
                use rayon::prelude::*;
                bits.par_iter()
                    .map(|&b| half::bf16::from_bits(b).to_f32())
                    .collect()
            } else {
                bits.iter()
                    .map(|&b| half::bf16::from_bits(b).to_f32())
                    .collect()
            }
        }
        // Integer I/O (token ids, durations, masks) widened to f32 for the
        // f32-uniform arena — same convention wgpu already uses. Values are
        // small integers (<2^24), so f32 is exact. This lets TTS/LM graphs
        // that feed I64/I32 inputs run on CUDA/Vulkan (VITS token ids etc.).
        DType::I64 => rlx_ir::bytes::decode_le::<i64>(data)
            .iter()
            .map(|&x| x as f32)
            .collect(),
        DType::I32 => rlx_ir::bytes::decode_le::<i32>(data)
            .iter()
            .map(|&x| x as f32)
            .collect(),
        DType::U32 => rlx_ir::bytes::decode_le::<u32>(data)
            .iter()
            .map(|&x| x as f32)
            .collect(),
        DType::I8 => data.iter().map(|&b| b as i8 as f32).collect(),
        DType::U8 | DType::Bool => data.iter().map(|&b| b as f32).collect(),
        // ── Complex (shared f32-uniform GPU boundary) ───────────────────
        // C64 = 2 f32 lanes `[re, im]`; the host already stores it as
        // interleaved f32 pairs, so widening is a pure reinterpret — no
        // conversion, N complex elements → 2N f32 lanes.
        DType::C64 => rlx_ir::bytes::decode_le_vec::<f32>(data),
        // C128 = 4 f32 lanes df64 `[re_hi, re_lo, im_hi, im_lo]`. The host
        // stores it as 2×f64 (16 B/elem); this is the df64 SPLIT boundary:
        // each f64 component `v` becomes `hi=(f32)v` + `lo=(f32)(v-(f64)hi)`
        // so `(f64)hi + (f64)lo` reconstructs `v` to full double precision.
        // Read via `from_le_bytes` (host `Vec<u8>` is only 1-aligned).
        DType::C128 => {
            let split = |v: f64| -> (f32, f32) {
                let hi = v as f32;
                let lo = (v - hi as f64) as f32;
                (hi, lo)
            };
            let mut out = Vec::with_capacity((data.len() / 16) * 4);
            for elem in data.chunks_exact(16) {
                let re = f64::from_le_bytes(elem[0..8].try_into().unwrap());
                let im = f64::from_le_bytes(elem[8..16].try_into().unwrap());
                let (re_hi, re_lo) = split(re);
                let (im_hi, im_lo) = split(im);
                out.push(re_hi);
                out.push(re_lo);
                out.push(im_hi);
                out.push(im_lo);
            }
            out
        }
        other => panic!(
            "widen_bytes_to_f32: dtype {other:?} unsupported on f32-arena backends \
             (only F32/F64/F16/BF16/I64/I32/U32/I8/U8/Bool/C64/C128 are accepted on the host I/O surface)"
        ),
    }
}

/// Narrow a `&[f32]` buffer down to the declared output dtype, returning
/// the corresponding little-endian byte stream. Mirrors the bytes a
/// backend that stores the native dtype would emit. Used by `run_typed`
/// to keep the byte-level output contract identical across backends.
#[allow(dead_code)]
pub fn narrow_f32_to_bytes(v: &[f32], dt: rlx_ir::DType) -> Vec<u8> {
    use rlx_ir::DType;
    match dt {
        DType::F32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&x.to_le_bytes());
            }
            bytes
        }
        DType::F16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&half::f16::from_f32(x).to_le_bytes());
            }
            bytes
        }
        DType::BF16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&half::bf16::from_f32(x).to_le_bytes());
            }
            bytes
        }
        DType::F64 => {
            let mut bytes = Vec::with_capacity(v.len() * 8);
            for &x in v {
                bytes.extend_from_slice(&(x as f64).to_le_bytes());
            }
            bytes
        }
        DType::I8 => v.iter().map(|&x| x as i8 as u8).collect(),
        DType::U8 => v.iter().map(|&x| x as u8).collect(),
        DType::I16 => {
            let mut bytes = Vec::with_capacity(v.len() * 2);
            for &x in v {
                bytes.extend_from_slice(&(x as i16).to_le_bytes());
            }
            bytes
        }
        DType::I32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&(x as i32).to_le_bytes());
            }
            bytes
        }
        DType::U32 => {
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&(x as u32).to_le_bytes());
            }
            bytes
        }
        DType::I64 => {
            let mut bytes = Vec::with_capacity(v.len() * 8);
            for &x in v {
                bytes.extend_from_slice(&(x as i64).to_le_bytes());
            }
            bytes
        }
        DType::Bool => v
            .iter()
            .map(|&x| if x != 0.0 { 1u8 } else { 0u8 })
            .collect(),
        DType::C64 => {
            // C64 = 2 f32 lanes `[re, im]`. The arena already holds a genuine
            // complex tensor as interleaved f32 pairs (2N lanes for N complex
            // elements), so narrowing is a pure reinterpret — emit the lanes as
            // 8-byte `[re, im]` pairs, no real→complex promotion.
            let mut bytes = Vec::with_capacity(v.len() * 4);
            for &x in v {
                bytes.extend_from_slice(&x.to_le_bytes());
            }
            bytes
        }
        DType::C128 => {
            // C128 = 4 f32 lanes df64 `[re_hi, re_lo, im_hi, im_lo]` per
            // complex element. This is the df64 COMBINE boundary: each host f64
            // component = `(f64)hi + (f64)lo`. Emit 16-byte `[re_f64, im_f64]`.
            let mut bytes = Vec::with_capacity((v.len() / 4) * 16);
            for lanes in v.chunks_exact(4) {
                let re = lanes[0] as f64 + lanes[1] as f64;
                let im = lanes[2] as f64 + lanes[3] as f64;
                bytes.extend_from_slice(&re.to_le_bytes());
                bytes.extend_from_slice(&im.to_le_bytes());
            }
            bytes
        }
    }
}

#[cfg(test)]
mod f64_boundary_tests {
    use super::{narrow_f32_to_bytes, widen_bytes_to_f32};
    use rlx_ir::DType;

    /// Regression: `widen_bytes_to_f32` used to `panic!` on F64, so uploading
    /// an F64 param/input (e.g. SPDNet's learnable `G`) on an f32-uniform
    /// backend (Vulkan/wgpu/CUDA/ROCm/oneAPI/TPU) crashed. It must now narrow
    /// f64→f32 for the f32 arena slot.
    #[test]
    fn widen_f64_narrows_without_panic() {
        let vals = [1.5f64, -2.25, 3.0, 4.5];
        let bytes: Vec<u8> = vals.iter().flat_map(|x| x.to_le_bytes()).collect();
        let out = widen_bytes_to_f32(&bytes, DType::F64);
        assert_eq!(out, vec![1.5f32, -2.25, 3.0, 4.5]);
    }

    /// F64 output round-trips: narrow f32→f64-bytes then widen back.
    #[test]
    fn f64_boundary_round_trips() {
        let v = [0.5f32, -1.0, 2.0];
        let bytes = narrow_f32_to_bytes(&v, DType::F64);
        assert_eq!(bytes.len(), v.len() * 8);
        assert_eq!(widen_bytes_to_f32(&bytes, DType::F64), v.to_vec());
    }

    /// C64 widen is a pure reinterpret: N complex host elements (`[re,im]`
    /// f32 pairs) → 2N f32 lanes, byte-identical.
    #[test]
    fn c64_widen_is_reinterpret() {
        // 2 complex elements: (1.5+2.5i), (-3.0+4.0i).
        let comps = [1.5f32, 2.5, -3.0, 4.0];
        let bytes: Vec<u8> = comps.iter().flat_map(|x| x.to_le_bytes()).collect();
        let lanes = widen_bytes_to_f32(&bytes, DType::C64);
        assert_eq!(lanes, comps.to_vec());
        // Narrow back → identical bytes.
        assert_eq!(narrow_f32_to_bytes(&lanes, DType::C64), bytes);
    }

    /// df64 C128 boundary: host f64 → widen(SPLIT) → 4 f32 lanes → narrow
    /// (COMBINE) → host f64. Sweeps π, 1/3, and an exactly-f32-representable
    /// value (round-trips bit-exact); rel err ≤ 1e-14.
    #[test]
    fn c128_df64_boundary_round_trips() {
        let reals = [
            std::f64::consts::PI,
            1.0 / 3.0,
            0.5,   // exactly f32-representable → bit-exact
            -2.25, // exactly f32-representable → bit-exact
            1.0e-12,
            123456.789012345,
        ];
        // Build a C128 host buffer: element k = reals[k] + reals[k]*2 i.
        let mut host = Vec::new();
        for &r in &reals {
            host.extend_from_slice(&r.to_le_bytes());
            host.extend_from_slice(&(r * 2.0).to_le_bytes());
        }
        let lanes = widen_bytes_to_f32(&host, DType::C128);
        assert_eq!(lanes.len(), reals.len() * 4, "4 f32 lanes per C128 element");
        let back = narrow_f32_to_bytes(&lanes, DType::C128);
        assert_eq!(back.len(), host.len());
        for (k, &r) in reals.iter().enumerate() {
            let re = f64::from_le_bytes(back[k * 16..k * 16 + 8].try_into().unwrap());
            let im = f64::from_le_bytes(back[k * 16 + 8..k * 16 + 16].try_into().unwrap());
            let rel = |a: f64, b: f64| {
                if b == 0.0 {
                    a.abs()
                } else {
                    (a - b).abs() / b.abs()
                }
            };
            assert!(rel(re, r) <= 1e-14, "re[{k}]={re} vs {r}");
            assert!(rel(im, r * 2.0) <= 1e-14, "im[{k}]={im} vs {}", r * 2.0);
        }
        // f32-exact values round-trip bit-exact.
        for &exact in &[0.5f64, -2.25] {
            let mut h = Vec::new();
            h.extend_from_slice(&exact.to_le_bytes());
            h.extend_from_slice(&0.0f64.to_le_bytes());
            let l = widen_bytes_to_f32(&h, DType::C128);
            let b = narrow_f32_to_bytes(&l, DType::C128);
            assert_eq!(
                b, h,
                "f32-exact C128 value {exact} must round-trip bit-exact"
            );
        }
    }
}

#[cfg(all(test, feature = "cpu"))]
mod capabilities_tests {
    use super::*;
    use rlx_ir::{DType, Graph, Shape};

    #[test]
    fn cpu_executable_declares_core_capabilities() {
        let mut g = Graph::new("caps");
        let x = g.input("x", Shape::new(&[2], DType::F32));
        g.set_outputs(vec![x]);
        let be = cpu_backend::CpuBackend;
        let exec = be.compile(g, &CompileOptions::default());
        let caps = exec.capabilities();
        assert!(caps.clone);
        assert!(caps.moe);
        assert!(caps.persistent_handles);
        assert!(caps.typed_io);
        assert!(caps.active_extent);
        assert!(!caps.gpu_handles);
        assert_eq!(
            caps.enabled_names(),
            vec![
                "clone",
                "moe",
                "persistent_handles",
                "typed_io",
                "active_extent"
            ]
        );
    }
}

/// A compiled, ready-to-execute graph on a specific backend.
///
/// # Method groups
///
/// See [`crate::ExecutableCapabilities`] for the optional feature bits and a
/// table of which trait methods belong to which group. Core execution is
/// `set_param` / `run` (and friends); everything else defaults to unsupported.
pub trait ExecutableGraph: Send {
    /// Advisory optional-feature bitmask. Default: all clear. Override when
    /// you implement MoE / GPU handles / typed I/O / etc. Callers must still
    /// treat individual method return values as authoritative.
    fn capabilities(&self) -> crate::ExecutableCapabilities {
        crate::ExecutableCapabilities::NONE
    }

    /// Set a named parameter (weight) buffer.
    fn set_param(&mut self, name: &str, data: &[f32]);

    /// Incrementally write raw `data` into a named param starting `byte_offset`
    /// bytes into its storage. Returns true if the backend performed the partial
    /// write; false means "unsupported / not applicable" and the caller must fall
    /// back to a whole-buffer [`Self::set_param_typed`]. Lets a large packed
    /// residency buffer (paged MoE experts) upload only the ONE changed slot
    /// instead of re-copying the whole buffer each step. Default: unsupported.
    fn set_param_range(&mut self, _name: &str, _byte_offset: usize, _data: &[u8]) -> bool {
        false
    }

    /// Called after all params are uploaded (`set_param` / `set_param_typed`).
    /// Backends may warm caches (e.g. Metal QMatMul weight dequant).
    fn finalize_params(&mut self) {}

    /// Deep-clone this executable into a fresh `Box`. Lets
    /// `CompiledGraph` implement `Clone` so callers (e.g. eda-mna's
    /// `SensitivityContext`) can spin up N independent executor
    /// copies for thread-parallel dispatch without paying the full
    /// graph-compile cost N times. Default implementation panics;
    /// backends that support cloning override.
    fn clone_box(&self) -> Box<dyn ExecutableGraph> {
        panic!("clone_box not implemented for this backend");
    }

    /// Execute the graph with named inputs. Returns output data (copies from arena).
    fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>>;

    /// Like [`Self::run`] but only read back outputs at `read_indices`.
    /// GPU handle feeds still update for every output. Default: all outputs.
    fn run_read_outputs(
        &mut self,
        inputs: &[(&str, &[f32])],
        read_indices: Option<&[usize]>,
    ) -> Vec<Vec<f32>> {
        match read_indices {
            None => self.run(inputs),
            Some(ix) => {
                // Backends without a native partial-read path still run the full
                // graph; only clone the requested outputs on the host.
                let all = self.run(inputs);
                ix.iter().filter_map(|&i| all.get(i).cloned()).collect()
            }
        }
    }

    /// ZERO-COPY in-place optimizer step over GPU-resident weights. For each
    /// `trainable[i] = (weight_input_name, shape)` — a resident arena Input (see
    /// [`Self::bind_gpu_handle`]) whose gradient is at output slot `1 + i`
    /// (backward outputs `[loss, grad0, …]`) — `step` is invoked with the param
    /// `&mut [f32]` and grad `&[f32]` ALIASED into the unified-memory arena, so
    /// the optimizer writes the weight in place: no host `Vec`, no D2H/H2D copy.
    /// The updated weight stays resident, so the next forward reads it with no
    /// re-upload — killing the GPU→host→optimizer→host→GPU roundtrip. Returns
    /// false on backends without GPU-resident support (caller falls back to a
    /// host-copy optimizer step). Default: unsupported.
    fn optimizer_step_resident(
        &mut self,
        _trainable: &[(String, Vec<usize>)],
        _step: &mut dyn FnMut(&str, &[usize], &mut [f32], &[f32]),
    ) -> bool {
        false
    }

    /// Execute and return raw pointers to output data in arena (zero-copy).
    fn run_raw(&mut self, inputs: &[(&str, &[f32])]) -> Vec<(*const f32, usize)> {
        let vecs = self.run(inputs);
        vecs.iter().map(|v| (v.as_ptr(), v.len())).collect()
    }

    /// Fastest: inputs by slot index, returns output (offset, len) pairs.
    /// Read output from arena via `arena_ptr().add(offset)`.
    fn run_slots(&mut self, _inputs: &[&[f32]]) -> &[(usize, usize)] {
        &[] // default: not supported
    }

    /// Get the raw arena buffer pointer for reading outputs after run_slots.
    fn arena_ptr(&self) -> *const u8 {
        std::ptr::null()
    }

    /// Hint the executor that subsequent `run` calls should process
    /// only the first `actual` rows along the bucket axis (out of
    /// `upper`, the extent the graph was compiled at). Backends that
    /// support per-kernel active-extent dispatch honor this; others
    /// ignore it and process the full compiled extent.
    ///
    /// Pass `None` to clear the hint. The hint is sticky — set it
    /// before each `run` and clear it after, or maintain it across
    /// runs at your discretion.
    ///
    /// Even when honored, callers must not rely on the contents of the
    /// output past `actual` rows — that region may contain stale data
    /// from earlier runs (kernels skip it).
    ///
    /// Default: no-op. See `BucketedCompileCache::run_padded` for the
    /// canonical caller; backends opt in by overriding this method.
    fn set_active_extent(&mut self, extent: Option<(usize, usize)>) {
        let _ = extent;
    }

    /// Return this graph's idle scratch (activation) memory to the OS while its
    /// compiled form stays cached — for a multi-subgraph pipeline (e.g. a TTS
    /// duration→encoder→flow→vocoder chain) whose stages run sequentially and
    /// would otherwise each pin a multi-hundred-MB activation arena resident at
    /// once. Params/weights stay resident; the next run transparently refaults +
    /// re-zeros scratch. Purely advisory — safe to leave as a no-op.
    /// Call only when the graph won't run again immediately (not inside a hot
    /// loop). Default: no-op; CPU overrides.
    fn release_scratch(&mut self) {}

    /// Override RNG policy for in-graph random ops without recompiling.
    fn set_rng(&mut self, rng: rlx_ir::RngOptions) {
        let _ = rng;
    }

    /// Current RNG policy (default when the backend does not override).
    fn rng(&self) -> rlx_ir::RngOptions {
        rlx_ir::RngOptions::default()
    }

    /// TIDE merged placement mask (union across MoE layers). CPU: stats + host path.
    fn set_moe_resident_experts(&mut self, _mask: &[bool]) {}

    /// Per MoE layer placement (`masks[layer][expert]`). Preferred over merged on CPU.
    fn set_moe_resident_experts_per_layer(&mut self, _masks: &[&[bool]]) {}

    /// Capture MoE router TopK indices on the next CPU forward (TIDE refresh).
    fn enable_moe_topk_capture(&mut self, _num_experts: usize) -> bool {
        false
    }

    /// Take captured per-layer expert indices (one vec per MoE TopK in order).
    fn take_moe_topk_capture(&mut self) -> Option<Vec<Vec<u32>>> {
        None
    }

    /// MoE GroupedMatMul residency accounting from the last forward (CPU).
    fn take_moe_residency_stats(&mut self) -> Option<crate::MoeResidencyStats> {
        None
    }

    /// Bind a persistent buffer handle (KV-cache, training state, etc.).
    /// The buffer lives across run() calls and is not in the arena.
    /// Returns true if the backend supports persistent handles.
    fn bind_handle(&mut self, _name: &str, _data: &[f32]) -> bool {
        false
    }

    /// Read a persistent buffer's current contents.
    fn read_handle(&self, _name: &str) -> Option<Vec<f32>> {
        None
    }

    /// GPU-resident input (MLX): upload once, reuse across runs.
    fn bind_gpu_handle(&mut self, _name: &str, _data: &[f32]) -> bool {
        false
    }

    fn has_gpu_handle(&self, _name: &str) -> bool {
        false
    }

    fn set_gpu_handle_feed(&mut self, _handle_name: &str, _output_index: usize) -> bool {
        false
    }

    fn read_gpu_handle(&self, _name: &str) -> Option<Vec<f32>> {
        None
    }

    /// Read one row from a resident GPU input handle without full-tensor D2H.
    fn read_gpu_handle_row(&self, _name: &str, _row: usize, _row_inner: usize) -> Option<Vec<f32>> {
        None
    }

    /// Register a targeted *row* feed for resident KV decode (graphs that emit
    /// the new token at the last bucket-padded output row). Returns false when
    /// the backend has no GPU-resident handle support. See `feed_kv_row`.
    fn register_kv_row_feed(&mut self, _handle_name: &str, _output_index: usize) -> bool {
        false
    }

    /// Fold each registered row feed's new-token row (`src_row` of its output)
    /// into the resident handle slot at `dst_row` (`row_elems` = kv_dim),
    /// in-place on device. Call after a logits-only run. Returns false when
    /// unsupported (caller keeps the host KV path).
    fn feed_kv_row(&mut self, _src_row: usize, _dst_row: usize, _row_elems: usize) -> bool {
        false
    }

    /// Batch-major resident KV: `new` is `[B,1,row_elems]`, past is
    /// `[B,seq_cap,row_elems]`. Writes each batch's new row into
    /// `past[b, dst_row, :]`. Default falls back to contiguous `feed_kv_row`
    /// when `batch == 1`.
    fn feed_kv_batch_major(
        &mut self,
        dst_row: usize,
        batch: usize,
        seq_cap: usize,
        row_elems: usize,
    ) -> bool {
        if batch == 1 {
            return self.feed_kv_row(0, dst_row, row_elems);
        }
        let _ = (dst_row, seq_cap, row_elems);
        false
    }

    /// Strided-source batch-major KV feed (folds the new token from a
    /// `[B, src_seq_cap, row_elems]` output at `src_row` into a `[B, dst_seq_cap,
    /// row_elems]` resident handle at `dst_row`). Default: unsupported.
    #[allow(clippy::too_many_arguments)]
    fn feed_kv_batch_major_src(
        &mut self,
        src_row: usize,
        src_seq_cap: usize,
        dst_row: usize,
        batch: usize,
        dst_seq_cap: usize,
        row_elems: usize,
    ) -> bool {
        let _ = (src_row, src_seq_cap, dst_row, batch, dst_seq_cap, row_elems);
        false
    }

    /// Ragged variant of [`Self::feed_kv_batch_major_src`]: each batch element
    /// folds into its own `dst_rows[b]`. Default: unsupported.
    fn feed_kv_batch_major_ragged(
        &mut self,
        src_row: usize,
        src_seq_cap: usize,
        dst_rows: &[usize],
        dst_seq_cap: usize,
        row_elems: usize,
    ) -> bool {
        let _ = (src_row, src_seq_cap, dst_rows, dst_seq_cap, row_elems);
        false
    }

    /// Mark a graph input as a device-resident handle with no host mirror.
    fn prepare_resident_gpu_handle(&mut self, _name: &str) -> bool {
        false
    }

    /// Upload bound (non-resident) GPU handle mirrors into the arena.
    fn stage_bound_gpu_handles_to_arena(&mut self) {}

    /// D2D seed of resident `past_k_*` / `past_v_*` from another executable's
    /// resident prefix (bucket rollover without host DRAM round-trip).
    fn seed_resident_kv_prefix_from(
        &mut self,
        _src: &dyn ExecutableGraph,
        _prefix_tokens: usize,
        _outgoing_upper: usize,
        _kv_dim: usize,
        _n_layers: usize,
    ) -> bool {
        false
    }

    /// D2D copy resident KV rows `[from_row..to_row)` from another executable.
    fn copy_resident_kv_rows_from(
        &mut self,
        _src: &dyn ExecutableGraph,
        _from_row: usize,
        _to_row: usize,
        _outgoing_upper: usize,
        _kv_dim: usize,
        _n_layers: usize,
    ) -> bool {
        false
    }

    /// Copy named parameter storage from another executable on the same backend.
    /// Used to avoid re-uploading packed U8 weights when compiling decode buckets.
    fn copy_params_from(&mut self, src: &dyn ExecutableGraph) -> bool {
        let _ = src;
        false
    }

    /// Share (retain, don't copy) the packed weight storage of another executable
    /// on the same backend when the weight layout matches EXACTLY — one GPU copy
    /// of the (large, read-only) weights backs both executables, instead of a
    /// per-executable duplicate. Returns false (caller must fall back to a full
    /// upload/copy) when unsupported or the layouts differ. `src` must already be
    /// fully populated. See `rlx_metal::MetalExecutable::share_weights_from`.
    fn share_params_from(&mut self, src: &dyn ExecutableGraph) -> bool {
        let _ = src;
        false
    }

    /// Downcast hook for [`Self::copy_params_from`]. Backends override when supported.
    fn executable_as_any(&self) -> Option<&dyn std::any::Any> {
        None
    }

    /// Mutable downcast hook for [`Self::copy_params_from`].
    fn executable_as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
        None
    }

    /// CUDA-only: mutable access for device KV seeding. Default `None`.
    #[cfg(feature = "cuda")]
    fn cuda_executable_for_kv_seed(&mut self) -> Option<&mut rlx_cuda::backend::CudaExecutable> {
        let _ = self;
        None
    }

    /// CUDA-only: immutable access for device KV seeding. Default `None`.
    #[cfg(feature = "cuda")]
    fn cuda_executable_for_kv_seed_ref(&self) -> Option<&rlx_cuda::backend::CudaExecutable> {
        None
    }

    /// Read one row from a row-major graph output after `run` / `run_read_outputs`.
    /// Metal reads a single row from the arena; default returns `None` (caller falls back).
    fn read_output_row(&self, _out_idx: usize, _row: usize, _row_inner: usize) -> Option<Vec<f32>> {
        None
    }

    /// Run and refresh a GPU handle from `output_index`; returns that output on host.
    fn run_feed_gpu_handle(
        &mut self,
        inputs: &[(&str, &[f32])],
        _handle_name: &str,
        _output_index: usize,
    ) -> Option<Vec<f32>> {
        let _ = inputs;
        None
    }

    // ── Pipelined / async execution (Phase C) ─────────────────────────
    //
    // These allow callers to amortize per-run sync latency on backends
    // where it matters (Metal: ~150 µs `wait_until_completed` per commit).
    // CPU has no such cost, so the default impls just call `run` serially.

    /// Encode + commit a forward pass without waiting for completion.
    ///
    /// Outputs of intermediate calls are stomped — use `run_pipelined` if
    /// you need outputs from each individual commit. Pair with
    /// `sync_pending` to drain.
    ///
    /// Default: synchronous fallback (calls `run`, discards output). CPU
    /// uses this default since BLAS is synchronous anyway.
    fn commit_no_wait(&mut self, inputs: &[(&str, &[f32])]) {
        let _ = self.run(inputs);
    }

    /// Wait for every command queued by `commit_no_wait`.
    /// Default: no-op (synchronous backends have nothing pending).
    fn sync_pending(&mut self) {}

    /// Issue a batch of forward passes pipelined, returning per-run outputs.
    ///
    /// The Metal impl encodes a per-commit blit so each in-flight run's
    /// outputs survive subsequent commits stomping the shared arena. The
    /// CPU default is just sequential `run`s — equally correct, no perf
    /// penalty (CPU has no GPU sync cost to amortize).
    ///
    /// Returns `out[run_idx][output_idx][element_idx]`.
    fn run_pipelined(&mut self, input_sets: &[Vec<(&str, &[f32])>]) -> Vec<Vec<Vec<f32>>> {
        input_sets.iter().map(|inputs| self.run(inputs)).collect()
    }

    // ── Typed (non-F32) host I/O ──────────────────────────────────
    //
    // `set_param` and `run` are F32 by contract. The typed entry
    // points let callers pass and receive raw bytes in any rlx-ir
    // dtype, avoiding the f32 widen/narrow round-trip that's
    // wasteful for F16/BF16 weights and activations.
    //
    // The default impls only handle F32 — any other dtype panics.
    // Backends that support typed I/O natively (e.g. MLX via
    // Array::from_bytes/to_bytes) override these.

    /// Set a named parameter from raw bytes in the given dtype.
    fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: rlx_ir::DType) {
        if dtype != rlx_ir::DType::F32 {
            panic!(
                "backend's default set_param_typed only handles F32; \
                    got {dtype:?}. Override on the backend for typed support."
            );
        }
        if !data.len().is_multiple_of(4) {
            panic!(
                "set_param_typed F32: data length {} not a multiple of 4",
                data.len()
            );
        }
        // `data` is caller-supplied, so it carries no alignment guarantee —
        // the `% 4` length check above says nothing about the pointer. Borrow
        // zero-copy when it is 4-aligned, decode elementwise when it is not.
        let f32_slice = rlx_ir::bytes::decode_le::<f32>(data);
        self.set_param(name, f32_slice.as_ref());
    }

    /// Run with typed inputs and typed outputs. Returns
    /// `(bytes, dtype)` per output; the dtype is whatever the
    /// graph's output node was declared as.
    fn run_typed(
        &mut self,
        inputs: &[(&str, &[u8], rlx_ir::DType)],
    ) -> Vec<(Vec<u8>, rlx_ir::DType)> {
        // Default impl: convert each typed input to f32 (F32-only),
        // run, then re-emit outputs as F32 bytes.
        let mut owned: Vec<(String, Vec<f32>)> = Vec::with_capacity(inputs.len());
        for (name, data, dt) in inputs {
            if *dt != rlx_ir::DType::F32 {
                panic!(
                    "backend's default run_typed only handles F32 inputs; \
                        got {dt:?} for input '{name}'"
                );
            }
            if data.len() % 4 != 0 {
                panic!(
                    "run_typed F32 input '{name}': len {} not multiple of 4",
                    data.len()
                );
            }
            owned.push((name.to_string(), rlx_ir::bytes::decode_le_vec::<f32>(data)));
        }
        let refs: Vec<(&str, &[f32])> = owned
            .iter()
            .map(|(n, d)| (n.as_str(), d.as_slice()))
            .collect();
        let outs = self.run(&refs);
        outs.into_iter()
            .map(|v| {
                let bytes =
                    unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * 4) }
                        .to_vec();
                (bytes, rlx_ir::DType::F32)
            })
            .collect()
    }

    /// True when the lowered wgpu schedule contains host-executed steps
    /// (unsupported on browser WebGPU). Default: false.
    #[cfg(target_arch = "wasm32")]
    fn wgpu_requires_host(&self) -> bool {
        false
    }

    /// Async WebGPU execution for wasm (non-blocking readback). Default: error.
    ///
    /// Futures are **not** `Send` on wasm — WebGPU buffer mapping uses
    /// `Rc`/`RefCell` and browser callbacks that are single-threaded.
    #[cfg(target_arch = "wasm32")]
    fn wgpu_run_async<'a>(
        &'a mut self,
        _inputs: &'a [(&'a str, &'a [f32])],
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<Vec<f32>>, String>> + 'a>>
    {
        Box::pin(async move { Err("wgpu_run_async not implemented for this backend".into()) })
    }
}

/// Backend implementation trait.
///
/// Single compile entry point. New compile-time knobs are added to
/// `CompileOptions`, not as new trait methods.
///
/// `Send + Sync` because backends are stateless factories — multiple
/// threads can call `compile` concurrently. The returned
/// `Box<dyn ExecutableGraph>` is `Send` (moveable to a worker thread)
/// but **not** `Sync` (`run`/`run_slots` take `&mut self`).
pub trait Backend: Send + Sync {
    /// Compile a graph for this backend with the given options.
    fn compile(&self, graph: Graph, options: &CompileOptions) -> Box<dyn ExecutableGraph>;

    /// Compile pre-optimized LIR (HIR → MIR → LIR pipeline output).
    /// Default re-enters [`Self::compile`] — backends should override
    /// when they can reuse the embedded buffer plan.
    fn compile_lir(&self, lir: LirModule, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
        self.compile(lir.into_graph(), options)
    }

    /// HIR-first compile: lower blocks, run fusion pipeline, emit executable.
    fn compile_hir(
        &self,
        hir: HirModule,
        device: rlx_driver::Device,
        options: &CompileOptions,
    ) -> Result<Box<dyn ExecutableGraph>, rlx_ir::hir::LowerError> {
        let result = crate::stages::compile_hir_stages(device, hir, options)?;
        crate::stages::maybe_log_fusion(&result.fusion);
        Ok(self.compile_lir(result.lir, options))
    }

    /// [`GraphModule`](rlx_ir::GraphModule) compile — unified HIR/MIR/LIR entry.
    fn compile_module(
        &self,
        module: rlx_ir::GraphModule,
        device: rlx_driver::Device,
        options: &CompileOptions,
    ) -> Result<Box<dyn ExecutableGraph>, rlx_ir::hir::LowerError> {
        let result = crate::stages::compile_module_stages(device, module, options)?;
        crate::stages::maybe_log_fusion(&result.fusion);
        Ok(self.compile_lir(result.lir, options))
    }

    /// PLAN L4: declare which `OpKind`s this backend can lower.
    /// Default: empty slice = "no claim made — accept everything"
    /// (preserves existing behavior; backends opt in by overriding).
    /// When non-empty, the `LegalizeForBackend` pass will refuse to
    /// compile a graph that contains an op outside this set, instead
    /// of silently falling through to slower / wrong dispatch.
    fn supported_ops(&self) -> &'static [rlx_ir::OpKind] {
        &[]
    }
}

/// Prepare a fused MIR graph from LIR for backend executable construction.
/// Skips the fusion pipeline — LIR must come from `compile_*_stages`.
#[allow(dead_code)]
fn prepare_fused_graph(
    graph: Graph,
    options: &CompileOptions,
    supported_ops: &[rlx_ir::OpKind],
    backend_name: &str,
) -> Graph {
    let (mut graph, report) = rlx_opt::prepare_graph_for_backend_with_report(
        graph,
        backend_name,
        supported_ops,
        options.kernel_dispatch,
    );
    rlx_opt::maybe_log_dispatch_report(&report);
    if !report.compile_ready {
        panic!(
            "{}\n{}",
            rlx_opt::format_legalize_error(backend_name, &report.still_unsupported),
            rlx_opt::format_dispatch_report(&report)
        );
    }
    // Backends that claim `OpKind::Scan` keep long Scans for host-fallback;
    // short ones IR-unroll so the body runs as ordinary device ops.
    if supported_ops.contains(&rlx_ir::OpKind::Scan) {
        graph = apply_scan_device_preference(graph, options);
    }
    graph = crate::precompile::post_fusion_cleanup(graph, options);
    if let Some(p) = options.policy.clone() {
        use rlx_opt::pass::Pass as _;
        graph = rlx_opt::AutoMixedPrecision::new(p).run(graph);
    }
    graph
}

/// Short-length + small-budget Scan unroll so GPU backends run the body on-device.
#[cfg_attr(not(feature = "cpu"), allow(unused_variables))]
fn apply_scan_device_preference(graph: Graph, options: &CompileOptions) -> Graph {
    // The short-length unroll macro lives in rlx-cpu; without it, fall through
    // to the budget-only unroll (both are pure IR rewrites).
    #[cfg(feature = "cpu")]
    let graph = rlx_cpu::rlx_maybe_unroll_scans!(graph, options.scan_unroll_max_length);
    rlx_opt::maybe_unroll_scans_budget(graph, 4096)
}

#[allow(dead_code)]
fn declared_output_dtypes(
    manifest: &cpu_low_precision::IoDtypeManifest,
    exec_dtypes: Vec<rlx_ir::DType>,
) -> Vec<rlx_ir::DType> {
    exec_dtypes
        .into_iter()
        .enumerate()
        .map(|(i, exec)| manifest.output_dtype(i, exec))
        .collect()
}

// ── Convenience helpers preserved from older API ──────────────────────
//
// These let existing call sites keep working unchanged while the new
// trait is the canonical one. We provide free functions rather than
// trait methods so adding them doesn't grow the trait surface.

/// Compile at default options (F32, no policy).
pub fn compile(backend: &dyn Backend, graph: Graph) -> Box<dyn ExecutableGraph> {
    backend.compile(graph, &CompileOptions::default())
}

/// Compile HIR through the fusion-first pipeline.
pub fn compile_hir(
    backend: &dyn Backend,
    hir: HirModule,
    device: rlx_driver::Device,
    options: &CompileOptions,
) -> Result<Box<dyn ExecutableGraph>, rlx_ir::hir::LowerError> {
    backend.compile_hir(hir, device, options)
}

/// Compile a [`GraphModule`](rlx_ir::GraphModule) through the fusion-first pipeline.
pub fn compile_module(
    backend: &dyn Backend,
    module: rlx_ir::GraphModule,
    device: rlx_driver::Device,
    options: &CompileOptions,
) -> Result<Box<dyn ExecutableGraph>, rlx_ir::hir::LowerError> {
    backend.compile_module(module, device, options)
}

/// Compile at a specific precision (default policy = none).
pub fn compile_with_precision(
    backend: &dyn Backend,
    graph: Graph,
    precision: crate::Precision,
) -> Box<dyn ExecutableGraph> {
    backend.compile(graph, &CompileOptions::new().precision(precision))
}

/// Helper retained for backward compatibility — applies the precision
/// rewrite at the runtime layer if backends don't override their
/// pipeline placement. Modern code: pass the policy via CompileOptions
/// and let the backend handle ordering.
fn _legacy_apply_policy(graph: Graph, policy: Option<rlx_opt::PrecisionPolicy>) -> Graph {
    use rlx_opt::pass::Pass as _;
    match policy {
        Some(p) => rlx_opt::AutoMixedPrecision::new(p).run(graph),
        None => graph,
    }
}

// ── CPU Backend ─────────────────────────────────────────────────────────

// Per-device Backend / ExecutableGraph wiring (feature-gated).

#[cfg(feature = "cpu")]
pub mod cpu_backend;

#[cfg(feature = "gpu")]
pub mod wgpu_backend;

#[cfg(feature = "opengl")]
pub mod webgl_backend;

#[cfg(feature = "vulkan")]
pub mod vulkan_backend;

#[cfg(feature = "oneapi")]
pub mod oneapi_backend;

#[cfg(all(feature = "mlx", rlx_mlx_host))]
pub mod mlx_backend;

/// Apple CoreML / Neural Engine (ANE) backend wiring.
///
/// Unlike the GPU backends, CoreML compiles a *static* model with weights
/// baked in, so we hand the raw graph (Param nodes intact) straight to
/// `rlx_coreml` and skip the fusion/LIR arena pipeline.
#[cfg(all(
    feature = "coreml",
    target_vendor = "apple",
    not(target_os = "watchos")
))]
pub mod coreml_backend;

#[cfg(all(feature = "metal", target_vendor = "apple", not(target_os = "watchos")))]
pub mod metal_backend;

#[cfg(feature = "cuda")]
pub mod cuda_backend;

#[cfg(feature = "rocm")]
pub mod rocm_backend;

#[cfg(feature = "tpu")]
pub mod tpu_backend;

/// QNN (Hexagon NPU) backend adapter — wraps the `rlx-qnn` FFI runtime
/// (`Device::Hexagon`). Lowers a general supported subgraph in-process via
/// `libQnnCpu.so` / `libQnnHtp.so`. Claims `FusedAttentionBlock` and
/// decomposes it before FFI lower (same pattern as CoreML / Vulkan).
#[cfg(feature = "qnn")]
pub mod qnn_backend;

/// AMD XDNA / Ryzen AI NPU backend adapter (`Device::Xdna`) — wraps `rlx-xdna`.
/// Detection + execution seam only: `compile` surfaces rlx_xdna's honest
/// diagnostic until the AIE kernel path lands (no CPU masquerade).
#[cfg(feature = "xdna")]
pub mod xdna_backend;

/// External GPU over a PCIe tunnel (`Device::Egpu`) — wraps `rlx-egpu`.
/// Discovery + transport seam only: `compile` surfaces rlx_egpu's diagnostic
/// until device bring-up lands (no CPU masquerade).
#[cfg(feature = "egpu")]
pub mod egpu_backend;

#[cfg(test)]
mod alignment_tests {
    use super::widen_bytes_to_f32;
    use rlx_ir::DType;

    /// Build `n_elems` elements of `dtype` as bytes, offset by `skew` filler
    /// bytes so the tensor data starts at `base + skew`. Returns the whole
    /// buffer; the caller slices `[skew..]` to get a deliberately misaligned
    /// `&[u8]`.
    fn skewed(dtype: DType, n_elems: usize, skew: usize) -> Vec<u8> {
        let mut out = vec![0xAAu8; skew];
        for i in 0..n_elems {
            let x = i as f64 - 3.5;
            match dtype {
                DType::F32 => out.extend_from_slice(&(x as f32).to_le_bytes()),
                DType::F64 => out.extend_from_slice(&x.to_le_bytes()),
                DType::F16 => out.extend_from_slice(&half::f16::from_f64(x).to_le_bytes()),
                DType::BF16 => out.extend_from_slice(&half::bf16::from_f64(x).to_le_bytes()),
                DType::I64 => out.extend_from_slice(&(i as i64).to_le_bytes()),
                DType::I32 => out.extend_from_slice(&(i as i32).to_le_bytes()),
                DType::U32 => out.extend_from_slice(&(i as u32).to_le_bytes()),
                DType::C64 => out.extend_from_slice(&(x as f32).to_le_bytes()),
                DType::C128 => {
                    out.extend_from_slice(&x.to_le_bytes());
                    out.extend_from_slice(&(-x).to_le_bytes());
                }
                other => panic!("skewed: unhandled {other:?}"),
            }
        }
        out
    }

    /// Guard for the `&[u8]` → `&[f32]` unsoundness (issue #1): the host I/O
    /// boundary takes caller-supplied byte slices, which carry no alignment
    /// guarantee — `&buf[1..]`, or a tensor at an odd offset inside an mmap,
    /// is 1-aligned. The old `from_raw_parts(data.as_ptr() as *const f32, n)`
    /// was UB there (Miri: "encountered an unaligned reference"), and a
    /// `len % 4 == 0` check says nothing about the pointer.
    ///
    /// Skews 0..8 cover every residue class mod 8 relative to the allocator's
    /// base, so every dtype sees both aligned and misaligned input and all of
    /// them must decode bit-identically.
    #[test]
    fn widen_is_bit_identical_on_misaligned_input() {
        const DTYPES: &[DType] = &[
            DType::F32,
            DType::F64,
            DType::F16,
            DType::BF16,
            DType::I64,
            DType::I32,
            DType::U32,
            DType::C64,
            DType::C128,
        ];
        for &dt in DTYPES {
            let n = 37;
            let reference = {
                let b = skewed(dt, n, 0);
                widen_bytes_to_f32(&b, dt)
            };
            assert!(!reference.is_empty(), "{dt:?}: empty reference");
            for skew in 1..8usize {
                let buf = skewed(dt, n, skew);
                let got = widen_bytes_to_f32(&buf[skew..], dt);
                assert_eq!(
                    got.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
                    reference.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
                    "{dt:?} at skew {skew} disagrees with the aligned decode"
                );
            }
        }
    }

    /// Byte-granular dtypes have no alignment requirement, but they share the
    /// same match arm, so pin them too.
    #[test]
    fn widen_byte_dtypes_ignore_skew() {
        for dt in [DType::U8, DType::I8, DType::Bool] {
            let mut buf = vec![0xAAu8; 3];
            buf.extend_from_slice(&[0u8, 1, 2, 3, 250]);
            assert_eq!(
                widen_bytes_to_f32(&buf[3..], dt),
                widen_bytes_to_f32(&[0u8, 1, 2, 3, 250], dt),
                "{dt:?}"
            );
        }
    }

    /// An empty typed input (an unused KV-cache slot passed as `Vec::new()`)
    /// has a dangling pointer; it must not reach a reinterpret at all.
    #[test]
    fn widen_empty_is_empty() {
        for dt in [DType::F32, DType::F16, DType::I64, DType::C128] {
            assert!(widen_bytes_to_f32(&[], dt).is_empty(), "{dt:?}");
        }
    }
}