Cargo.toml

[package]
name = "mem"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
rayon = "1.5"
indicatif = {version = "0.16", features = ["rayon"]}
ndarray = { version = "0.15", features = ["rayon", "serde", "blas"] }
blas-src = { version = "*", default-features = false, features = ["openblas"] }
openblas-src = { version = "0.6.1", default-features = false, features = ["cblas"] }
ndarray-npy = "0.8"
ndarray-rand = "0.14"
memmap2 = "0.5"
bincode = "1.3"
serde_bytes = "0.11.6"
rkyv = { version = "0.7", features = ["validation"] }
abomonation = "0.7.3"

[[bin]]
name = "main"
path = "src/main.rs"

[[bin]]
name = "viewer"
path = "src/viewer.rs"

viewer.rs

extern crate blas_src;

use ndarray_npy::{write_npy, ViewNpyExt};

use memmap2::Mmap;
use std::{fs::OpenOptions, time::Instant};

use ndarray::{Array, Array1, ArrayView2};
use ndarray_rand::{rand_distr::Uniform, RandomExt};

struct Embeddings {
    mmap: Mmap,
}

impl Embeddings {
    fn view(&self) -> ndarray::ArrayBase<ndarray::ViewRepr<&f32>, ndarray::Dim<[usize; 2]>> {
        ArrayView2::<f32>::view_npy(&self.mmap).expect("viewed mmap")
    }
}

fn main() -> Result<(), &'static str> {
    let d = 256;
    let n = 10_000;
    let x = Array1::random(d, Uniform::<f32>::new(0., 1.));

    let path = "/tmp/embeds.npy";
    write_embeddings(n, d, path);
    let embeddings = get_embeddings(path);

    println!("view()");
    let start = Instant::now();
    let v = embeddings.view();
    println!("took: {:?}", start.elapsed());
    println!("in npy view row(i)");
    let start = Instant::now();
    for i in 0..n {
        let row = v.row(i);
        let _res = x.dot(&row);
    }
    println!("took: {:?}", start.elapsed());

    Ok(())
}

fn get_embeddings(path: &str) -> Embeddings {
    let file = OpenOptions::new()
        .read(true)
        .open(path)
        .expect("opened file");
    let mmap = unsafe { Mmap::map(&file).expect("memmap'd embeds.npy") };
    let embeddings = Embeddings { mmap };
    embeddings
}

fn write_embeddings(n: usize, d: usize, path: &str) {
    println!("generating ({},{}) embeddings", n, d);
    let start = Instant::now();
    let mut ys = Vec::new();
    for _i in 0..n {
        let y = Array1::random(d, Uniform::<f32>::new(0., 1.));
        ys.push(y);
    }
    println!("took: {:?}", start.elapsed());

    println!("writing npy view");
    let start = Instant::now();
    let nd = Array::from_shape_vec((n, d), ys.into_iter().flatten().collect())
        .expect("reshape vec of ndarrays into ndarray");
    write_npy(path, &nd).expect("write npy");
    println!("took: {:?}", start.elapsed());
}