environment/training-source/source-snapshot/src/inference.rs

//! The inference worker: a thread that owns the meganeura session and
//! publishes coloured feature grids for the renderer to display.
//!
//! # Why a thread at all
//!
//! The encoder is expected to take 40–60 ms on a Quest 3 while the
//! compositor wants a frame every 13.9 ms. Running inference inline would
//! put the render loop on the encoder's cadence, which is uncomfortable at
//! best. Instead the renderer runs free, always drawing the most recent
//! completed result, and inference lands whenever it lands.
//!
//! # Why the worker builds its own session
//!
//! `meganeura::Session` is not `Send` — it owns a `CommandEncoder` holding
//! raw Vulkan handles. It cannot be constructed here and moved. What *is*
//! shareable is `Arc<blade_graphics::Context>`, so the worker receives the
//! context and builds the session in place. See `tests/threading.rs`.
//!
//! # The caveat this design cannot fix
//!
//! Both threads submit to the same Vulkan queue, which blade guards with a
//! mutex. A long compute submission can still delay the render submission
//! queued behind it — threading decouples *CPU* orchestration, not GPU
//! occupancy. Whether that shows up as dropped frames is exactly what the
//! on-device numbers will reveal; if it does, the fix is splitting the plan
//! across frames, which needs meganeura to grow partial execution
//! (`step()` currently submits every dispatch in one go).

use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::time::Instant;

use crate::dinov3::Config;
use crate::pca::{self, Basis};

/// A completed inference result: a square RGB image, either the coarse
/// feature colouring or a full reconstruction.
#[derive(Debug, Clone)]
pub struct FeatureGrid {
    /// Row-major RGB. For [`Display::PcaColour`] this is `[tokens, 3]`
    /// including the prefix tokens; for [`Display::Reconstruction`] it is
    /// already `[grid * grid, 3]` with no prefix.
    pub colors: Vec<f32>,
    /// Side length of the displayable square: the patch grid for PCA
    /// colouring, the image size for a reconstruction.
    pub grid: usize,
    pub prefix_tokens: usize,
    /// Encoder wall time, for the on-screen/logged rate readout.
    pub latency_ms: f64,
}

impl FeatureGrid {
    /// The displayable `grid × grid × 3` image, with any prefix tokens
    /// dropped.
    pub fn patch_rgb(&self) -> &[f32] {
        &self.colors[self.prefix_tokens * pca::COMPONENTS..]
    }
}

/// What the encoder's output is turned into for display.
#[derive(Clone, Debug)]
pub enum Display {
    /// Project patch features onto their top three principal components and
    /// read those off as RGB. No training, one `[hidden, 3]` matmul, and a
    /// `grid × grid` result. Shows what the features *distinguish*.
    PcaColour,
    /// Reconstruct RGB through the trained decoder — the actual roundtrip.
    /// Produces a full `image_size × image_size` picture, so the readback
    /// per frame is ~600 KB rather than ~3 KB. Shows what the features
    /// *retain*.
    Reconstruction(std::path::PathBuf),
}

/// What the worker is told to do next.
enum Command {
    /// Encode this image. `[num_patches, patch_dim]`, from
    /// [`crate::preprocess`].
    Submit(Vec<f32>),
    /// Refit the PCA basis from the next frame's features.
    Refit,
    Stop,
}

/// Handle to a running worker.
pub struct Worker {
    tx: mpsc::Sender<Command>,
    latest: Arc<Mutex<Option<FeatureGrid>>>,
    generation: Arc<AtomicU64>,
    ready: Arc<AtomicBool>,
    handle: Option<std::thread::JoinHandle<()>>,
}

impl Worker {
    /// Queue a frame for encoding.
    ///
    /// Returns `false` if the worker is busy or gone. Dropping frames is
    /// the correct behaviour: the camera produces 60 Hz and the encoder
    /// cannot keep up, so anything queued would be stale by the time it ran.
    pub fn submit(&self, patches: Vec<f32>) -> bool {
        if !self.ready.load(Ordering::Acquire) {
            return false;
        }
        self.ready.store(false, Ordering::Release);
        self.tx.send(Command::Submit(patches)).is_ok()
    }

    /// Ask for the colour basis to be refitted from the next frame.
    pub fn request_refit(&self) {
        let _ = self.tx.send(Command::Refit);
    }

    /// The most recent completed result, if any.
    pub fn latest(&self) -> Option<FeatureGrid> {
        self.latest.lock().ok().and_then(|g| g.clone())
    }

    /// Monotonic count of completed encodes, so the renderer can tell a
    /// fresh result from a repeat without comparing pixel data.
    pub fn generation(&self) -> u64 {
        self.generation.load(Ordering::Acquire)
    }

    /// True when the worker is idle and would accept a frame.
    pub fn is_ready(&self) -> bool {
        self.ready.load(Ordering::Acquire)
    }
}

impl Drop for Worker {
    fn drop(&mut self) {
        let _ = self.tx.send(Command::Stop);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

/// How the worker should obtain its weights.
pub enum Weights {
    /// Load a checkpoint from disk. On device this is a path the app has
    /// pushed or unpacked from its assets.
    SafeTensors(std::path::PathBuf),
    /// Deterministic synthetic weights. Produces meaningless features, but
    /// exercises the entire pipeline end to end — useful for bringing the
    /// render path up before the real checkpoint is on the device.
    Synthetic,
}

/// Start the worker.
///
/// Returns immediately; the session is built on the worker thread, which
/// takes a few seconds on a mobile CPU. [`Worker::is_ready`] reports false
/// until it is done, and [`Worker::latest`] yields `None`.
pub fn spawn(
    gpu: Arc<blade_graphics::Context>,
    config: Config,
    weights: Weights,
    plan_cache: Option<std::path::PathBuf>,
    submission_chunks: usize,
    display: Display,
) -> Worker {
    let (tx, rx) = mpsc::channel();
    let latest = Arc::new(Mutex::new(None));
    let generation = Arc::new(AtomicU64::new(0));
    let ready = Arc::new(AtomicBool::new(false));

    let handle = {
        let latest = Arc::clone(&latest);
        let generation = Arc::clone(&generation);
        let ready = Arc::clone(&ready);
        std::thread::Builder::new()
            .name("dinovision-inference".into())
            .spawn(move || {
                worker_main(
                    gpu,
                    config,
                    weights,
                    plan_cache,
                    submission_chunks,
                    display,
                    rx,
                    latest,
                    generation,
                    ready,
                );
            })
            .expect("failed to spawn inference thread")
    };

    Worker {
        tx,
        latest,
        generation,
        ready,
        handle: Some(handle),
    }
}

#[allow(clippy::too_many_arguments)]
fn worker_main(
    gpu: Arc<blade_graphics::Context>,
    config: Config,
    weights: Weights,
    plan_cache: Option<std::path::PathBuf>,
    submission_chunks: usize,
    display: Display,
    rx: mpsc::Receiver<Command>,
    latest: Arc<Mutex<Option<FeatureGrid>>>,
    generation: Arc<AtomicU64>,
    ready: Arc<AtomicBool>,
) {
    // For PCA colouring both the colours and the features are outputs:
    // colours every frame (3 KB), features only when refitting the basis
    // (300 KB), which is why they are separate. A reconstruction needs
    // neither — the decoder consumes the features inside the graph.
    let reconstructing = matches!(display, Display::Reconstruction(_));
    let mut g = meganeura::Graph::new();
    let features = crate::dinov3::build_encoder(&mut g, &config);
    if reconstructing {
        let rgb = crate::decoder::attach_to_encoder(&mut g, &config, features);
        g.set_outputs(vec![rgb]);
    } else {
        let colors = pca::add_projection(&mut g, features, config.hidden_size);
        g.set_outputs(vec![colors, features]);
    }

    let build_start = Instant::now();
    let (mut session, _) = meganeura::train::build(&g, meganeura::train::SessionConfig {
        mode: meganeura::train::Mode::Inference,
        gpu: Some(gpu),
        cache: plan_cache.as_deref(),
        ..Default::default()
    });
    log::info!(
        "inference session ready in {:.1} s",
        build_start.elapsed().as_secs_f64()
    );

    // Hand the queue back periodically so the renderer can get a frame in.
    // One long submission is faster in isolation but starves everything
    // else sharing the device.
    session.set_submission_chunks(submission_chunks);

    match weights {
        Weights::SafeTensors(path) => {
            match meganeura::data::safetensors::SafeTensorsModel::load(path.clone()) {
                Ok(model) => {
                    if let Err(e) = crate::weights::load_encoder(&mut session, &model, &config) {
                        log::error!("failed to bind weights from {}: {e}", path.display());
                        return;
                    }
                }
                Err(e) => {
                    log::error!("failed to read {}: {e}", path.display());
                    return;
                }
            }
        }
        Weights::Synthetic => {
            log::warn!("running with synthetic weights — features are meaningless");
            crate::bench::fill_parameters(&mut session, &g);
        }
    }

    let tokens = config.num_tokens();
    let mut basis = Basis::placeholder(config.hidden_size);
    let mut refit_pending = false;
    let mut colors_out;
    let mut features_out = vec![0.0f32; tokens * config.hidden_size];
    // Reconstruction reads a planar [3, H, W] image; the renderer wants it
    // interleaved, so keep a staging buffer for the transpose.
    let mut planar = Vec::new();

    match &display {
        Display::PcaColour => {
            // Start with a placeholder basis so the first frame displays
            // something, then refit from real features as soon as one lands.
            refit_pending = true;
            apply_basis(&mut session, &basis);
            colors_out = vec![0.0f32; tokens * pca::COMPONENTS];
        }
        Display::Reconstruction(path) => {
            if let Err(e) = crate::decoder::load_parameters(&mut session, &g, path) {
                log::error!("failed to load the decoder: {e}");
                return;
            }
            let pixels = config.image_size * config.image_size;
            planar = vec![0.0f32; 3 * pixels];
            colors_out = vec![0.0f32; pixels * 3];
        }
    }

    ready.store(true, Ordering::Release);

    while let Ok(cmd) = rx.recv() {
        let patches = match cmd {
            Command::Submit(p) => p,
            Command::Refit => {
                refit_pending = true;
                continue;
            }
            Command::Stop => break,
        };

        let start = Instant::now();
        session.set_input("patches", &patches);
        session.step();
        session.wait();

        if reconstructing {
            session.read_output_by_index(0, &mut planar);
            // [3, H, W] → [H * W, 3].
            let pixels = config.image_size * config.image_size;
            for p in 0..pixels {
                for c in 0..3 {
                    colors_out[p * 3 + c] = planar[c * pixels + p];
                }
            }
        } else {
            session.read_output_by_index(0, &mut colors_out);
        }

        if refit_pending {
            // Refitting needs the full feature matrix, so this frame pays
            // for the larger readback. It happens once at startup and then
            // only on request.
            session.read_output_by_index(1, &mut features_out);
            basis = Basis::fit(
                &features_out,
                tokens,
                config.hidden_size,
                config.num_prefix_tokens(),
            );
            apply_basis(&mut session, &basis);
            refit_pending = false;
            log::info!("refitted colour basis from live features");
            // The colours just read were produced by the old basis; redo
            // the projection on the CPU so this frame is not displayed with
            // a stale palette.
            colors_out = basis.project(&features_out, tokens);
        }

        let grid = FeatureGrid {
            colors: colors_out.clone(),
            grid: if reconstructing {
                config.image_size
            } else {
                config.grid()
            },
            prefix_tokens: if reconstructing {
                0
            } else {
                config.num_prefix_tokens()
            },
            latency_ms: start.elapsed().as_secs_f64() * 1000.0,
        };
        if let Ok(mut slot) = latest.lock() {
            *slot = Some(grid);
        }
        generation.fetch_add(1, Ordering::Release);
        ready.store(true, Ordering::Release);
    }

    log::info!("inference worker stopped");
}

fn apply_basis(session: &mut meganeura::Session, basis: &Basis) {
    session.set_parameter("pca.weight", &basis.weight_matrix());
    session.set_parameter("pca.bias", &basis.bias_vector());
}