Struct Pool 

pub struct Pool { /* private fields */ }

Persistent thread pool: shared job slot, epoch dispatch, caller participation.

Implementations

impl Pool

pub fn new(n_workers: usize) -> Self

Examples found in repository

examples/pool_lat.rs (line 15)

fn main() {
    for nt in [2usize, 4, 6, 8, 10] {
        let pool = Pool::new(nt);
        let noop: &(dyn Fn(usize, usize) + Sync) = &|_w, _n| {};
        // Warm: first dispatch spawns/parks the workers.
        for _ in 0..100 {
            pool.run(noop);
        }
        let iters = 20_000;
        let t0 = Instant::now();
        for _ in 0..iters {
            pool.run(noop);
        }
        let el = t0.elapsed().as_secs_f64();
        let per = el / iters as f64 * 1e6;
        println!(
            "threads={nt:2}  {per:7.2} us/dispatch  -> {:6.2} ms/token at 200 matvecs",
            per * 200.0 / 1e3
        );
    }
}

examples/matvec_bw.rs (line 60)

fn main() {
    let mut args = std::env::args().skip(1);
    let path = args
        .next()
        .expect("usage: matvec_bw <model.cmf> [sweep|one <tensor>]");
    let mode = args.next().unwrap_or_else(|| "sweep".to_string());
    let model = Arc::new(cortiq_core::CmfModel::open(&path).expect("open model"));

    // Real LM activations carry a few heavy channels (>8·rms); measured
    // mean on this model is ~3.7. NOUT models that distribution.
    let nout: usize = std::env::var("NOUT")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(4);
    let mk_x = |cols: usize| -> Vec<f32> {
        let mut x: Vec<f32> = (0..cols).map(|i| ((i % 17) as f32 - 8.0) / 8.0).collect();
        for k in 0..nout {
            x[k * 37 % cols] = 40.0;
        }
        x
    };

    if mode == "one" {
        let name = args
            .next()
            .unwrap_or_else(|| "model.embed_tokens.weight".to_string());
        let entry = model.tensor(&name).expect("tensor not found");
        let (rows, cols) = (entry.shape[0], entry.shape[1]);
        let nbytes = entry.nbytes as f64;
        println!(
            "tensor {name}: {rows}x{cols} {:?} = {:.1} MB, NOUT={nout}",
            entry.dtype,
            nbytes / 1e6
        );
        let t = QTensor::from_model(&model, &name).expect("wrap");
        let x = mk_x(cols);
        let mut out = vec![0f32; rows];
        t.matvec(&x, &mut out, None);
        for nt in [1usize, 2, 4, 6, 8, 10] {
            let pool = if nt == 1 { None } else { Some(Pool::new(nt)) };
            let iters = 8;
            t.matvec(&x, &mut out, pool.as_ref());
            let t0 = Instant::now();
            for _ in 0..iters {
                t.matvec(&x, &mut out, pool.as_ref());
            }
            let el = t0.elapsed().as_secs_f64();
            println!(
                "threads={nt:2}  {:6.2} ms/matvec  {:6.1} GB/s (sink {:.3})",
                el / iters as f64 * 1e3,
                nbytes * iters as f64 / el / 1e9,
                out[0]
            );
        }
        return;
    }

    // sweep: every 2-D q8_2f tensor once = one decode's worth of weights.
    let names: Vec<String> = model
        .tensors
        .iter()
        .filter(|t| t.dtype == TensorDtype::Q8_2f && t.shape.len() == 2)
        .map(|t| t.name.clone())
        .collect();
    let total_bytes: f64 = model
        .tensors
        .iter()
        .filter(|t| t.dtype == TensorDtype::Q8_2f && t.shape.len() == 2)
        .map(|t| t.nbytes as f64)
        .sum();
    println!(
        "sweep: {} q8_2f tensors, {:.2} GB total (= weights streamed per decode token), NOUT={nout}",
        names.len(),
        total_bytes / 1e9
    );

    let tensors: Vec<(QTensor, Vec<f32>, Vec<f32>)> = names
        .iter()
        .map(|n| {
            let e = model.tensor(n).unwrap();
            let (rows, cols) = (e.shape[0], e.shape[1]);
            (
                QTensor::from_model(&model, n).expect("wrap"),
                mk_x(cols),
                vec![0f32; rows],
            )
        })
        .collect();
    let mut tensors = tensors;

    // Whole-model residency pass BEFORE any timing: the first touch of a
    // 4.2 GB mmap faults ~260k pages, which would otherwise be charged to
    // whichever thread count happens to run first. REVERSE=1 flips the
    // order as a check that no first-touch cost is left in the table.
    for _ in 0..2 {
        for (t, x, out) in tensors.iter_mut() {
            t.matvec(x, out, None);
        }
    }
    let mut counts = vec![1usize, 2, 4, 6, 8, 10];
    if std::env::var("REVERSE").is_ok() {
        counts.reverse();
    }
    for nt in counts {
        let pool = if nt == 1 { None } else { Some(Pool::new(nt)) };
        // one warm pass, then two measured
        for (t, x, out) in tensors.iter_mut() {
            t.matvec(x, out, pool.as_ref());
        }
        let iters = 2;
        let t0 = Instant::now();
        for _ in 0..iters {
            for (t, x, out) in tensors.iter_mut() {
                t.matvec(x, out, pool.as_ref());
            }
        }
        let el = t0.elapsed().as_secs_f64();
        let per_tok = el / iters as f64;
        println!(
            "threads={nt:2}  {:7.1} ms/sweep  {:6.1} GB/s  -> weight-path-only ceiling {:5.1} tok/s",
            per_tok * 1e3,
            total_bytes * iters as f64 / el / 1e9,
            1.0 / per_tok
        );
    }
}

pub fn with_spin(n_workers: usize, spin_budget: usize) -> Self

Explicit spin budget (tests pin it without touching the env).

pub fn effective_threads() -> usize

The thread count from_env would use RIGHT NOW: forced (C ABI)

CMF_THREADS > big-core topology > available_parallelism−1. ≤1 means the model runs serial (no pool). Introspection (execution_mode, status endpoints) must report THIS, not available_parallelism.

Examples found in repository

examples/mimo_audio_dump.rs (line 303)

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let cmd = args.get(1).map(String::as_str).unwrap_or("");
    let out = PathBuf::from(arg(&args, "--out").expect("--out"));
    match cmd {
        "decode" => {
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            let flat: Vec<f32> = w.channels.concat();
            save_f32(&out, &[w.channels.len(), w.frames()], &flat);
            println!(
                "rate {} channels {} frames {}",
                w.sample_rate,
                w.channels.len(),
                w.frames()
            );
        }
        "frontend" => {
            std::fs::create_dir_all(&out).unwrap();
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let t0 = Instant::now();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            save_f32(
                &out.join("dec.npy"),
                &[w.channels.len(), w.frames()],
                &w.channels.concat(),
            );
            let chans: Vec<Vec<f32>> = w
                .channels
                .iter()
                .map(|c| mimo_audio::resample_sinc(c, w.sample_rate, mimo_audio::SAMPLE_RATE))
                .collect();
            save_f32(
                &out.join("chan24k.npy"),
                &[chans.len(), chans[0].len()],
                &chans.concat(),
            );
            let mono = mimo_audio::wav_to_mono_24k(&w).unwrap();
            save_f32(&out.join("wave24k.npy"), &[mono.len()], &mono);
            let pool = cortiq_engine::pool::Pool::from_env();
            let (mel, m) = mimo_audio::log_mel(&mono, pool.as_deref()).unwrap();
            save_f32(&out.join("mel.npy"), &[m, mimo_audio::N_MELS], &mel);
            println!(
                "rate {} channels {} frames {} -> {} samples, {m} mel frames, K {} ({:.3}s)",
                w.sample_rate,
                w.channels.len(),
                w.frames(),
                mono.len(),
                mimo_audio::audio_token_count(m, 4),
                t0.elapsed().as_secs_f64()
            );
        }
        "tower" => {
            std::fs::create_dir_all(&out).unwrap();
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let t0 = Instant::now();
            let audio = if src.extension().is_some_and(|e| e == "cmf") {
                let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
                MimoAudio::from_model(&model).expect("load towers")
            } else {
                MimoAudio::from_hf_dir(&src).expect("load towers")
            };
            let t_load = t0.elapsed().as_secs_f64();
            let (mel, m) = if let Some(mp) = arg(&args, "--mel") {
                let (shape, v) = load_f32(Path::new(&mp));
                assert_eq!(shape[1], mimo_audio::N_MELS);
                (v, shape[0])
            } else {
                let wav = std::fs::read(arg(&args, "--wav").expect("--wav or --mel")).unwrap();
                audio.wav_to_mel(&wav).unwrap()
            };
            let t1 = Instant::now();
            let feats = audio.features(&mel, m).unwrap();
            let t_feats = t1.elapsed().as_secs_f64();
            let d = audio.tokenizer.cfg.d_model;
            let rows = feats.len() / d;
            save_f32(&out.join("feats.npy"), &[rows, d], &feats);
            let t2 = Instant::now();
            let exact = audio.tokenizer.quantize(&feats, rows, false, audio.pool());
            let t_rvq = t2.elapsed().as_secs_f64();
            let rounded = audio.tokenizer.quantize(&feats, rows, true, audio.pool());
            let levels = exact.len() / rows;
            save_i32(&out.join("codes_exact.npy"), &[rows, levels], &exact);
            save_i32(&out.join("codes_bf16books.npy"), &[rows, levels], &rounded);
            let own = mimo_audio::AudioCodes {
                frames: rows,
                levels,
                codes: if audio.bf16_codebooks {
                    rounded.clone()
                } else {
                    exact.clone()
                },
            };
            let t3 = Instant::now();
            let emb_own = audio.embed_codes(&own).unwrap();
            let t_enc = t3.elapsed().as_secs_f64();
            save_f32(
                &out.join("embeds_own.npy"),
                &[emb_own.n_tokens, emb_own.dim],
                &emb_own.rows,
            );
            let fixed = match arg(&args, "--codes") {
                Some(cp) => {
                    let (shape, v) = load_codes(Path::new(&cp));
                    mimo_audio::AudioCodes {
                        frames: shape[0],
                        levels: shape[1],
                        codes: v,
                    }
                }
                None => own.clone(),
            };
            let emb = audio.embed_codes(&fixed).unwrap();
            save_f32(&out.join("embeds.npy"), &[emb.n_tokens, emb.dim], &emb.rows);
            let k = mimo_audio::audio_token_count(m, audio.encoder.cfg.group);
            let meta = serde_json::json!({
                "src": src.display().to_string(),
                "mel_frames": m,
                "segments": mimo_audio::segment_lengths(m),
                "codes": rows,
                "placeholder_count_K": k,
                "embed_rows_own": emb_own.n_tokens,
                "embed_rows_fixed": emb.n_tokens,
                "load_s": t_load,
                "features_s": t_feats,
                "rvq_s": t_rvq,
                "encoder_s": t_enc,
                "bf16_codebooks_default": audio.bf16_codebooks,
                "threads": cortiq_engine::pool::Pool::effective_threads(),
            });
            std::fs::write(
                out.join("tower.json"),
                serde_json::to_string_pretty(&meta).unwrap(),
            )
            .unwrap();
            println!("{meta}");
            assert_eq!(emb_own.n_tokens, k, "placeholder count != encoder rows");
        }
        "calib" => {
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
            let audio = MimoAudio::from_model(&model).expect("load towers");
            let dir = PathBuf::from(arg(&args, "--wav-dir").expect("--wav-dir"));
            let mut wavs: Vec<PathBuf> = std::fs::read_dir(&dir)
                .unwrap()
                .filter_map(|e| e.ok().map(|e| e.path()))
                .filter(|p| p.extension().is_some_and(|e| e == "wav"))
                .collect();
            wavs.sort();
            let t0 = Instant::now();
            cortiq_engine::gptq_capture::begin(true);
            let mut frames = 0usize;
            for w in &wavs {
                let emb = audio.embed_wav(&std::fs::read(w).unwrap()).unwrap();
                frames += emb.n_tokens;
                eprintln!(
                    "  {} -> {} rows ({:.0}s)",
                    w.display(),
                    emb.n_tokens,
                    t0.elapsed().as_secs_f64()
                );
            }
            let hess = cortiq_engine::gptq_capture::end();
            save_hessians(&out, &hess);
            println!(
                "{} clips, {frames} LLM rows, {} linears -> {} ({:.0}s)",
                wavs.len(),
                hess.len(),
                out.display(),
                t0.elapsed().as_secs_f64()
            );
        }
        _ => {
            eprintln!(
                "usage: mimo_audio_dump (decode|frontend|tower|calib) --out ... (see the source header)"
            );
            std::process::exit(2);
        }
    }
}

pub fn from_env() -> Option<Arc<Self>>

Pool sized from CMF_THREADS (see module docs). None = serial. Without the env, heterogeneous ARM defaults to its BIG cores.

Examples found in repository

examples/mimo_audio_dump.rs (line 214)

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let cmd = args.get(1).map(String::as_str).unwrap_or("");
    let out = PathBuf::from(arg(&args, "--out").expect("--out"));
    match cmd {
        "decode" => {
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            let flat: Vec<f32> = w.channels.concat();
            save_f32(&out, &[w.channels.len(), w.frames()], &flat);
            println!(
                "rate {} channels {} frames {}",
                w.sample_rate,
                w.channels.len(),
                w.frames()
            );
        }
        "frontend" => {
            std::fs::create_dir_all(&out).unwrap();
            let wav = std::fs::read(arg(&args, "--wav").expect("--wav")).unwrap();
            let t0 = Instant::now();
            let w = mimo_audio::decode_wav(&wav).unwrap();
            save_f32(
                &out.join("dec.npy"),
                &[w.channels.len(), w.frames()],
                &w.channels.concat(),
            );
            let chans: Vec<Vec<f32>> = w
                .channels
                .iter()
                .map(|c| mimo_audio::resample_sinc(c, w.sample_rate, mimo_audio::SAMPLE_RATE))
                .collect();
            save_f32(
                &out.join("chan24k.npy"),
                &[chans.len(), chans[0].len()],
                &chans.concat(),
            );
            let mono = mimo_audio::wav_to_mono_24k(&w).unwrap();
            save_f32(&out.join("wave24k.npy"), &[mono.len()], &mono);
            let pool = cortiq_engine::pool::Pool::from_env();
            let (mel, m) = mimo_audio::log_mel(&mono, pool.as_deref()).unwrap();
            save_f32(&out.join("mel.npy"), &[m, mimo_audio::N_MELS], &mel);
            println!(
                "rate {} channels {} frames {} -> {} samples, {m} mel frames, K {} ({:.3}s)",
                w.sample_rate,
                w.channels.len(),
                w.frames(),
                mono.len(),
                mimo_audio::audio_token_count(m, 4),
                t0.elapsed().as_secs_f64()
            );
        }
        "tower" => {
            std::fs::create_dir_all(&out).unwrap();
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let t0 = Instant::now();
            let audio = if src.extension().is_some_and(|e| e == "cmf") {
                let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
                MimoAudio::from_model(&model).expect("load towers")
            } else {
                MimoAudio::from_hf_dir(&src).expect("load towers")
            };
            let t_load = t0.elapsed().as_secs_f64();
            let (mel, m) = if let Some(mp) = arg(&args, "--mel") {
                let (shape, v) = load_f32(Path::new(&mp));
                assert_eq!(shape[1], mimo_audio::N_MELS);
                (v, shape[0])
            } else {
                let wav = std::fs::read(arg(&args, "--wav").expect("--wav or --mel")).unwrap();
                audio.wav_to_mel(&wav).unwrap()
            };
            let t1 = Instant::now();
            let feats = audio.features(&mel, m).unwrap();
            let t_feats = t1.elapsed().as_secs_f64();
            let d = audio.tokenizer.cfg.d_model;
            let rows = feats.len() / d;
            save_f32(&out.join("feats.npy"), &[rows, d], &feats);
            let t2 = Instant::now();
            let exact = audio.tokenizer.quantize(&feats, rows, false, audio.pool());
            let t_rvq = t2.elapsed().as_secs_f64();
            let rounded = audio.tokenizer.quantize(&feats, rows, true, audio.pool());
            let levels = exact.len() / rows;
            save_i32(&out.join("codes_exact.npy"), &[rows, levels], &exact);
            save_i32(&out.join("codes_bf16books.npy"), &[rows, levels], &rounded);
            let own = mimo_audio::AudioCodes {
                frames: rows,
                levels,
                codes: if audio.bf16_codebooks {
                    rounded.clone()
                } else {
                    exact.clone()
                },
            };
            let t3 = Instant::now();
            let emb_own = audio.embed_codes(&own).unwrap();
            let t_enc = t3.elapsed().as_secs_f64();
            save_f32(
                &out.join("embeds_own.npy"),
                &[emb_own.n_tokens, emb_own.dim],
                &emb_own.rows,
            );
            let fixed = match arg(&args, "--codes") {
                Some(cp) => {
                    let (shape, v) = load_codes(Path::new(&cp));
                    mimo_audio::AudioCodes {
                        frames: shape[0],
                        levels: shape[1],
                        codes: v,
                    }
                }
                None => own.clone(),
            };
            let emb = audio.embed_codes(&fixed).unwrap();
            save_f32(&out.join("embeds.npy"), &[emb.n_tokens, emb.dim], &emb.rows);
            let k = mimo_audio::audio_token_count(m, audio.encoder.cfg.group);
            let meta = serde_json::json!({
                "src": src.display().to_string(),
                "mel_frames": m,
                "segments": mimo_audio::segment_lengths(m),
                "codes": rows,
                "placeholder_count_K": k,
                "embed_rows_own": emb_own.n_tokens,
                "embed_rows_fixed": emb.n_tokens,
                "load_s": t_load,
                "features_s": t_feats,
                "rvq_s": t_rvq,
                "encoder_s": t_enc,
                "bf16_codebooks_default": audio.bf16_codebooks,
                "threads": cortiq_engine::pool::Pool::effective_threads(),
            });
            std::fs::write(
                out.join("tower.json"),
                serde_json::to_string_pretty(&meta).unwrap(),
            )
            .unwrap();
            println!("{meta}");
            assert_eq!(emb_own.n_tokens, k, "placeholder count != encoder rows");
        }
        "calib" => {
            let src = PathBuf::from(arg(&args, "--src").expect("--src"));
            let model = Arc::new(cortiq_core::CmfModel::open(&src).expect("open cmf"));
            let audio = MimoAudio::from_model(&model).expect("load towers");
            let dir = PathBuf::from(arg(&args, "--wav-dir").expect("--wav-dir"));
            let mut wavs: Vec<PathBuf> = std::fs::read_dir(&dir)
                .unwrap()
                .filter_map(|e| e.ok().map(|e| e.path()))
                .filter(|p| p.extension().is_some_and(|e| e == "wav"))
                .collect();
            wavs.sort();
            let t0 = Instant::now();
            cortiq_engine::gptq_capture::begin(true);
            let mut frames = 0usize;
            for w in &wavs {
                let emb = audio.embed_wav(&std::fs::read(w).unwrap()).unwrap();
                frames += emb.n_tokens;
                eprintln!(
                    "  {} -> {} rows ({:.0}s)",
                    w.display(),
                    emb.n_tokens,
                    t0.elapsed().as_secs_f64()
                );
            }
            let hess = cortiq_engine::gptq_capture::end();
            save_hessians(&out, &hess);
            println!(
                "{} clips, {frames} LLM rows, {} linears -> {} ({:.0}s)",
                wavs.len(),
                hess.len(),
                out.display(),
                t0.elapsed().as_secs_f64()
            );
        }
        _ => {
            eprintln!(
                "usage: mimo_audio_dump (decode|frontend|tower|calib) --out ... (see the source header)"
            );
            std::process::exit(2);
        }
    }
}

pub fn n_workers(&self) -> usize

Spawned worker threads (the caller joins each job on top).

pub fn bind_numa(&self, regions: &[&[u8]])

Keep the pool on the NUMA node that holds regions (the model’s weight bytes). Linux with two or more nodes only; CMF_NUMA=0 turns it off, CMF_NUMA=node:<n> forces a node.

WHY: decode streams every weight once per token, and on a two-socket host the page cache holds a file on whichever node read it. Unpinned, the scheduler spreads the workers over both sockets and half the matvec rows cross the socket link. Measured on a 2×EPYC 7763 pod with the model’s pages all on node 0 (31 CPUs of cgroup quota): a STREAM-style read over a node-0 buffer gives 42 GB/s from 31 unpinned threads and 74 GB/s from 31 threads kept on node 0. The mask is the node’s physical cores (first SMT sibling) when there are enough of them for the pool, else the whole node; never narrower than the pool, so nothing oversubscribes. Threads are bound to a SET of cores, not to one core each: the scheduler still balances inside the node. The calling thread adopts the same mask on its next dispatch.

pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync))

Run f(row_start, row_end) over 0..rows, self-balancing.

One dispatch, but workers pull row-ranges from a shared cursor instead of each taking a fixed 1/n slice. On a heterogeneous CPU (Apple Silicon: 4 P-cores + 6 E-cores here) a static split makes every matvec end at the SLOWEST core’s pace while the fast ones idle at the barrier; pulling by grain lets a P-core take several chunks for each one an E-core takes, so skew collapses to a single grain. Row ranges stay disjoint and each row’s dot is computed exactly as in the serial path → bit-identical output.

pub fn run_many(&self, parts: &[(usize, &(dyn Fn(usize, usize) + Sync))])

Multi-matrix job: one dispatch serves SEVERAL row spaces (roadmap §3 P0 — «одна внешняя публикация job на слой»). Parts are laid out back-to-back in a virtual row space and pulled by grain from one shared cursor, so QKV or gate+up cost a single barrier instead of one each. Each part’s f(start, end) sees its OWN row indices — per-row math and outputs are bit-identical to separate run_rows calls.

pub fn run(&self, f: &(dyn Fn(usize, usize) + Sync))

Run f(worker_idx, n_participants) on every worker AND the calling thread (worker_idx = n_workers() for the caller); returns when all participants have finished.

Examples found in repository

examples/pool_lat.rs (line 19)

fn main() {
    for nt in [2usize, 4, 6, 8, 10] {
        let pool = Pool::new(nt);
        let noop: &(dyn Fn(usize, usize) + Sync) = &|_w, _n| {};
        // Warm: first dispatch spawns/parks the workers.
        for _ in 0..100 {
            pool.run(noop);
        }
        let iters = 20_000;
        let t0 = Instant::now();
        for _ in 0..iters {
            pool.run(noop);
        }
        let el = t0.elapsed().as_secs_f64();
        let per = el / iters as f64 * 1e6;
        println!(
            "threads={nt:2}  {per:7.2} us/dispatch  -> {:6.2} ms/token at 200 matvecs",
            per * 200.0 / 1e3
        );
    }
}

Trait Implementations

impl Drop for Pool

fn drop(&mut self)

Executes the destructor for this type.

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned.