open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/load_model.rs

//! Real-model bring-up: load any GGUF + exercise the
//! substrate stack against it.
//!
//! ## Usage
//!
//! ```bash
//! cargo run --release --example load_model -- path/to/model.gguf
//! cargo run --release --example load_model -- model.gguf --policy=none
//! cargo run --release --example load_model -- model.gguf --policy=top-n:10
//! cargo run --release --example load_model -- model.gguf --tier-pin-top=8
//! ```
//!
//! Prints:
//!
//! - Load wall-clock under the chosen [`PrefetchPolicy`]
//! - Tensor count + total bytes
//! - Top tensors by byte span (the ones a partial-load
//!   would prefetch first)
//! - Per-tensor access-hint distribution (Random /
//!   Sequential / WillNeed)
//! - Optional: register with a [`TierStore`] + pin the
//!   top-K tensors Hot; report final tier-store metrics.
//!
//! This binary doesn't run inference — execution lives in
//! `jouleclaw-omni`. It does verify the FULL load path
//! against any real GGUF on disk, including the prefetch +
//! tier-store wiring. Use it to:
//!
//! 1. Catch loader bugs against real (vs synthetic) models.
//! 2. Compare cold-load wall-clock across PrefetchPolicy.
//! 3. Confirm the access-hint classifier picks reasonable
//!    hints for a given architecture (Llama, Gemma, Qwen,
//!    Mistral all share GGUF naming conventions).

use jouleclaw_loader_gguf::{
    cpu_prefetch::PrefetchLevel,
    mmap_advise::AdviceKind,
    read_gguf_file_with_policy,
    tier_store::{StoreTensorMeta, TierConfig, TierStore},
    classify_access_hint, GgufBuffer, GgufModel, PrefetchPolicy,
};
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!(
            "usage: load_model <path.gguf>|--synth [--policy=none|all|per-tensor-only|top-n:N] [--tier-pin-top=K]"
        );
        std::process::exit(1);
    }
    let path = if args.iter().any(|a| a == "--synth") {
        synth_to_tempfile()?
    } else {
        PathBuf::from(&args[1])
    };
    let policy = parse_policy(&args).unwrap_or(PrefetchPolicy::All);
    let pin_top: Option<usize> = parse_pin_top(&args);
    let _ = PrefetchLevel::T0; // suppress unused-import on non-arch targets
    let _ = AdviceKind::WillNeed;

    println!("model:  {}", path.display());
    println!("policy: {:?}", policy);
    if let Some(k) = pin_top {
        println!("tier-pin-top: {}", k);
    }

    let t0 = Instant::now();
    let model = read_gguf_file_with_policy(&path, policy)
        .map_err(|e| format!("parse: {:?}", e))?;
    let load_secs = t0.elapsed().as_secs_f64();
    println!();
    println!("=== load ===");
    println!("  load_wall_clock_secs: {:.3}", load_secs);
    println!("  gguf_version: {}", model.version);
    println!("  metadata_keys: {}", model.metadata.len());
    println!("  tensor_count: {}", model.tensors.len());
    print_buffer_info(&model);
    println!();

    let metas = compute_metas(&model);
    let total_bytes: u64 = metas.iter().map(|m| m.size_bytes).sum();
    println!("=== tensor data ===");
    println!("  total_bytes: {}", total_bytes);
    println!(
        "  total_human: {}",
        humanize_bytes(total_bytes)
    );
    println!();

    print_top_n_by_size(&metas, 10);
    println!();
    print_hint_distribution(&model);
    println!();
    print_dtype_distribution(&model);
    println!();
    print_moe_structure(&model);
    println!();

    if let Some(k) = pin_top {
        println!("=== tier_store pin_top ===");
        run_tier_store(&path, &model, &metas, k)?;
    }
    Ok(())
}

fn parse_policy(args: &[String]) -> Option<PrefetchPolicy> {
    for a in args {
        if let Some(v) = a.strip_prefix("--policy=") {
            return Some(match v {
                "none" => PrefetchPolicy::None,
                "all" => PrefetchPolicy::All,
                "per-tensor-only" => PrefetchPolicy::PerTensorOnly,
                s if s.starts_with("top-n:") => {
                    let n: usize =
                        s.trim_start_matches("top-n:").parse().ok()?;
                    PrefetchPolicy::TopN(n)
                }
                _ => return None,
            });
        }
    }
    None
}

fn parse_pin_top(args: &[String]) -> Option<usize> {
    for a in args {
        if let Some(v) = a.strip_prefix("--tier-pin-top=") {
            return v.parse().ok();
        }
    }
    None
}

fn print_buffer_info(model: &GgufModel) {
    match &model.buf {
        GgufBuffer::Owned(v) => {
            println!("  buffer_variant: Owned (size = {} bytes)", v.len());
        }
        GgufBuffer::Mapped { mmap, data_offset } => {
            println!(
                "  buffer_variant: Mapped (mmap_len = {}, data_offset = {})",
                mmap.len(),
                data_offset
            );
        }
        GgufBuffer::Backing { backing, data_offset } => {
            println!(
                "  buffer_variant: Backing (len = {}, data_offset = {})",
                backing.len(),
                data_offset
            );
        }
    }
}

fn compute_metas(model: &GgufModel) -> Vec<StoreTensorMeta> {
    let data_end_in_data: u64 = match &model.buf {
        GgufBuffer::Mapped { mmap, data_offset } => {
            (mmap.len() - data_offset) as u64
        }
        GgufBuffer::Backing { backing, data_offset } => {
            (backing.len() - data_offset) as u64
        }
        GgufBuffer::Owned(v) => v.len() as u64,
    };
    let mut sorted: Vec<usize> = (0..model.tensors.len()).collect();
    sorted.sort_by_key(|&i| model.tensors[i].offset);
    let mut metas = Vec::with_capacity(model.tensors.len());
    for (rank, &idx) in sorted.iter().enumerate() {
        let t = &model.tensors[idx];
        let next_off = sorted
            .get(rank + 1)
            .map(|&j| model.tensors[j].offset)
            .unwrap_or(data_end_in_data);
        metas.push(StoreTensorMeta {
            name: t.name.clone(),
            offset_in_data: t.offset,
            size_bytes: next_off.saturating_sub(t.offset),
        });
    }
    metas
}

fn print_top_n_by_size(metas: &[StoreTensorMeta], n: usize) {
    println!("=== top tensors by byte span ===");
    let mut sorted: Vec<&StoreTensorMeta> = metas.iter().collect();
    sorted.sort_by_key(|a| std::cmp::Reverse(a.size_bytes));
    for m in sorted.iter().take(n) {
        println!(
            "  {:<60} {:>12} bytes ({})",
            m.name,
            m.size_bytes,
            humanize_bytes(m.size_bytes)
        );
    }
}

fn print_hint_distribution(model: &GgufModel) {
    println!("=== access-hint distribution ===");
    let mut counts: BTreeMap<String, u32> = BTreeMap::new();
    for t in &model.tensors {
        let hint = classify_access_hint(t);
        let key = format!("{:?}", hint);
        *counts.entry(key).or_insert(0) += 1;
    }
    for (k, v) in &counts {
        println!("  {:<16} {:>6}", k, v);
    }
}

fn run_tier_store(
    path: &std::path::Path,
    model: &GgufModel,
    metas: &[StoreTensorMeta],
    pin_top: usize,
) -> Result<(), Box<dyn std::error::Error>> {
    let data_offset = match &model.buf {
        GgufBuffer::Mapped { data_offset, .. } => *data_offset as u64,
        _ => {
            println!("  (skipped — tier_store needs a Mapped buffer)");
            return Ok(());
        }
    };
    let store = TierStore::new(TierConfig::default());
    store.register_model("m1", path, data_offset)?;
    store.register_tensors("m1", metas);
    // Pin the top-K by byte span Hot.
    let mut sorted: Vec<&StoreTensorMeta> = metas.iter().collect();
    sorted.sort_by_key(|a| std::cmp::Reverse(a.size_bytes));
    let t0 = Instant::now();
    for (i, m) in sorted.iter().take(pin_top).enumerate() {
        store.pin_hot("m1", m, i as u64 + 1)?;
    }
    let pin_secs = t0.elapsed().as_secs_f64();
    let s = store.metrics.snapshot();
    println!("  pinned_top_k: {}", pin_top);
    println!("  pin_wall_clock_secs: {:.6}", pin_secs);
    println!("  hot_bytes:  {} ({})", store.hot_bytes(), humanize_bytes(store.hot_bytes()));
    println!("  warm_bytes: {} ({})", store.warm_bytes(), humanize_bytes(store.warm_bytes()));
    println!("  cold_bytes: {} ({})", store.cold_bytes(), humanize_bytes(store.cold_bytes()));
    println!("  metrics: {:?}", s);
    Ok(())
}

fn print_dtype_distribution(model: &GgufModel) {
    println!("=== quant format distribution ===");
    let mut counts: BTreeMap<String, (u32, u64)> = BTreeMap::new();
    for t in &model.tensors {
        let key = format!("{:?}", t.dtype);
        let entry = counts.entry(key).or_insert((0, 0));
        entry.0 += 1;
        // approximate byte span via block layout if quantized
        let n_elems = t.shape.iter().map(|&d| d as usize).product::<usize>();
        let bytes = if let Some(sz) = t.dtype.element_size() {
            (n_elems * sz) as u64
        } else if let Some((bb, eb)) = t.dtype.block_layout() {
            ((n_elems / eb) * bb) as u64
        } else {
            0
        };
        entry.1 += bytes;
    }
    for (k, (count, bytes)) in &counts {
        println!(
            "  {:<10} count={:>4}  size={}",
            k,
            count,
            humanize_bytes(*bytes)
        );
    }
}

fn print_moe_structure(model: &GgufModel) {
    // Two distinct MoE GGUF layout conventions:
    //
    //   A. Per-expert tensor (wave 32 routing — DeepSeek-V2,
    //      older Mixtral exports):
    //        blk.{N}.{kind}_exps.{E}.weight
    //
    //   B. Stacked-expert tensor (wave 33 routing — Mellum 2,
    //      Qwen3-MoE, modern Mixtral exports):
    //        blk.{N}.{kind}_exps.weight       (3-D shape [E, ...])
    //
    // The stacked form holds ALL experts in one tensor; per-
    // token activation = byte-range slice into that tensor.

    let re_per = regex::Regex::new(
        r"^blk\.(?P<blk>\d+)\.(?P<kind>[a-z_]+)_exps\.(?P<exp>\d+)\.weight$",
    )
    .unwrap();
    let re_stacked = regex::Regex::new(
        r"^blk\.(?P<blk>\d+)\.(?P<kind>[a-z_]+)_exps\.weight$",
    )
    .unwrap();

    let mut per_expert: BTreeMap<usize, BTreeMap<String, u32>> = BTreeMap::new();
    let mut stacked: Vec<(usize, String, Vec<u64>)> = Vec::new();
    for t in &model.tensors {
        if let Some(caps) = re_per.captures(&t.name) {
            let blk: usize = caps["blk"].parse().unwrap();
            let kind = caps["kind"].to_string();
            *per_expert.entry(blk).or_default().entry(kind).or_insert(0) += 1;
        } else if let Some(caps) = re_stacked.captures(&t.name) {
            let blk: usize = caps["blk"].parse().unwrap();
            let kind = caps["kind"].to_string();
            stacked.push((blk, kind, t.shape.clone()));
        }
    }

    if !per_expert.is_empty() {
        let total_blocks = per_expert.len();
        let first_block = *per_expert.keys().next().unwrap();
        let kinds = per_expert.get(&first_block).unwrap();
        let experts = kinds.values().next().copied().unwrap_or(0);
        println!("=== MoE structure (per-expert tensor layout) ===");
        println!("  blocks_with_experts: {}", total_blocks);
        println!("  experts_per_block:   {}", experts);
        println!(
            "  expert_kinds:        {:?}",
            kinds.keys().collect::<Vec<_>>()
        );
        println!("  → wave-32 routing applies (one tensor per expert)");
    } else if !stacked.is_empty() {
        // Group by block.
        let mut by_block: BTreeMap<usize, Vec<(String, Vec<u64>)>> = BTreeMap::new();
        for (blk, kind, shape) in &stacked {
            by_block.entry(*blk).or_default().push((kind.clone(), shape.clone()));
        }
        let total_blocks = by_block.len();
        let first_block = *by_block.keys().next().unwrap();
        let first_kinds = by_block.get(&first_block).unwrap();
        // Build the wave-33 stacked-expert index.
        let idx = jouleclaw_loader_gguf::moe::build_stacked_expert_index(model, 0);
        // GGUF shape order is REVERSED from PyTorch: a
        // [num_experts, ffn_inner, hidden] PyTorch tensor
        // is stored in GGUF as [hidden, ffn_inner, num_experts].
        // So the expert axis is the LAST element of the shape.
        let num_experts = first_kinds
            .first()
            .and_then(|(_, s)| s.last().copied())
            .unwrap_or(0);
        println!("=== MoE structure (stacked-expert tensor layout) ===");
        println!("  blocks_with_experts: {}", total_blocks);
        println!("  num_experts (from trailing axis): {}", num_experts);
        println!(
            "  expert_kinds:        {:?}",
            first_kinds.iter().map(|(k, _)| k).collect::<Vec<_>>()
        );
        println!("  first tensor shape:  {:?}", first_kinds[0].1);
        println!(
            "  → wave-33 stacked-expert routing applies (one tensor per (block, kind), \
             expert axis = trailing dim in GGUF order)"
        );
        // Surface the wave-33 index's view + per-expert stride.
        println!("  wave-33 index entries: {}", idx.len());
        if let Some(first) = idx.first() {
            println!(
                "  first entry: block={} kind={} num_experts={} stride_bytes={} (= {})",
                first.block,
                first.kind,
                first.num_experts,
                first.expert_stride(),
                humanize_bytes(first.expert_stride()),
            );
            let active_bytes_per_token: u64 = idx
                .iter()
                // Per token, top-2 active per (block, kind) × all blocks × all kinds.
                .map(|e| 2 * e.expert_stride())
                .sum();
            let total_expert_bytes: u64 =
                idx.iter().map(|e| e.tensor_byte_len).sum();
            println!(
                "  per-token active expert bytes (top-2): {} (= {})",
                active_bytes_per_token,
                humanize_bytes(active_bytes_per_token)
            );
            println!(
                "  total stacked-expert bytes: {} (= {})",
                total_expert_bytes,
                humanize_bytes(total_expert_bytes)
            );
            let savings = 1.0
                - (active_bytes_per_token as f64 / total_expert_bytes as f64);
            println!(
                "  → routing saves {:.1}% of expert bandwidth per token",
                savings * 100.0
            );
        }
    } else {
        println!("=== MoE structure ===");
        println!("  (none — no expert-named tensors found)");
    }
}

fn synth_to_tempfile() -> Result<PathBuf, Box<dyn std::error::Error>> {
    use jouleclaw_loader_gguf::synthetic::{synthesize_llama_gguf, SyntheticConfig};
    let cfg = SyntheticConfig {
        vocab_size: 256,
        embedding_length: 128,
        block_count: 4,
        feed_forward_length: 256,
        head_count: 4,
        head_count_kv: 4,
        ..SyntheticConfig::default()
    };
    let bytes = synthesize_llama_gguf(&cfg);
    let dir = std::env::temp_dir();
    let pid = std::process::id();
    let path = dir.join(format!("jouleclaw-load-model-synth-{}.gguf", pid));
    std::fs::write(&path, &bytes)?;
    eprintln!("[synth] wrote {} bytes → {}", bytes.len(), path.display());
    Ok(path)
}

fn humanize_bytes(n: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;
    if n >= GB {
        format!("{:.2} GiB", n as f64 / GB as f64)
    } else if n >= MB {
        format!("{:.2} MiB", n as f64 / MB as f64)
    } else if n >= KB {
        format!("{:.2} KiB", n as f64 / KB as f64)
    } else {
        format!("{} B", n)
    }
}