knn_hot/
knn_hot.rs

1//! Tight k-nearest loop for perf / cachegrind / roofline.
2//!
3//! Env: `KNN_HOT_N` (default 4096), `KNN_HOT_K` (default 4),
4//! `KNN_HOT_REPS` (default 40), `KNN_HOT_SHAPE`:
5//! `cubic` (default), `slab` (eOn / readcon orthorhombic film),
6//! `column` (seams water column), `hex` (rgsaddle hexagonal cell),
7//! `tilt` (seams LAMMPS dump with an xy tilt).
8
9use std::time::Instant;
10
11use linkcell::{knearest_into, Cell};
12
13fn env_usize(key: &str, default: usize) -> usize {
14    std::env::var(key)
15        .ok()
16        .and_then(|s| s.parse().ok())
17        .unwrap_or(default)
18}
19
20fn lattice(n: usize, boxl: f64) -> Vec<[f64; 3]> {
21    let nside = ((n as f64).cbrt().ceil() as usize).max(1);
22    let a = boxl / nside as f64;
23    let mut xyz = Vec::with_capacity(n);
24    for i in 0..n {
25        let ix = i % nside;
26        let iy = (i / nside) % nside;
27        let iz = i / (nside * nside);
28        xyz.push([ix as f64 * a, iy as f64 * a, iz as f64 * a]);
29    }
30    xyz
31}
32
33fn frac_grid(cell: &Cell, n: usize) -> Vec<[f64; 3]> {
34    let nside = ((n as f64).cbrt().ceil() as usize).max(1);
35    let mut xyz = Vec::with_capacity(n);
36    for i in 0..n {
37        let ix = i % nside;
38        let iy = (i / nside) % nside;
39        let iz = i / (nside * nside);
40        let s = [
41            (ix as f64 + 0.5) / nside as f64,
42            (iy as f64 + 0.5) / nside as f64,
43            (iz as f64 + 0.5) / nside as f64,
44        ];
45        xyz.push(cell.cartesian(s));
46    }
47    xyz
48}
49
50fn nside_of(n: usize) -> usize {
51    ((n as f64).cbrt().round() as usize).max(1)
52}
53
54/// Orthorhombic film at the cubic spacing, with an equal vacuum along z.
55/// This is an eOn slab or a readcon CON, scaled up.
56fn slab(n: usize) -> (Vec<[f64; 3]>, Cell) {
57    let nside = nside_of(n);
58    let film = nside as f64 * 3.125;
59    let cell = Cell::ortho(film, film, film * 2.0).expect("slab");
60    let mut xyz = Vec::with_capacity(n);
61    for i in 0..n {
62        let ix = i % nside;
63        let iy = (i / nside) % nside;
64        let iz = i / (nside * nside);
65        xyz.push([
66            (ix as f64 + 0.5) / nside as f64 * film,
67            (iy as f64 + 0.5) / nside as f64 * film,
68            (iz as f64 + 0.5) / nside as f64 * film,
69        ]);
70    }
71    (xyz, cell)
72}
73
74/// Tall orthorhombic column at the cubic spacing. Points occupy a band
75/// and the rest of z is vacuum, as in a seams nanotube trajectory.
76fn column(n: usize) -> (Vec<[f64; 3]>, Cell) {
77    let nside = nside_of(n);
78    let side = nside as f64 * 3.125;
79    let cell = Cell::ortho(side, side, side * 4.0).expect("column");
80    let mut xyz = Vec::with_capacity(n);
81    for i in 0..n {
82        let ix = i % nside;
83        let iy = (i / nside) % nside;
84        let iz = i / (nside * nside);
85        xyz.push([
86            (ix as f64 + 0.5) / nside as f64 * side,
87            (iy as f64 + 0.5) / nside as f64 * side,
88            (iz as f64 + 0.5) / nside as f64 * side + side * 1.5,
89        ]);
90    }
91    (xyz, cell)
92}
93
94fn system(shape: &str, n: usize) -> (Vec<[f64; 3]>, Cell) {
95    match shape {
96        "cubic" => {
97            let boxl = 50.0;
98            (
99                lattice(n, boxl),
100                Cell::ortho(boxl, boxl, boxl).expect("box"),
101            )
102        }
103        "slab" => slab(n),
104        "column" => column(n),
105        "hex" => {
106            // rgsaddle hexagonal cell, scaled to the cubic spacing.
107            let cell = Cell::from_vectors(
108                [50.0, 0.0, 0.0],
109                [25.0, 43.30127018922193, 0.0],
110                [0.0, 0.0, 55.55555555555556],
111                [0.0, 0.0, 0.0],
112            )
113            .expect("hex");
114            (frac_grid(&cell, n), cell)
115        }
116        "tilt" => {
117            // genice sH dump (BOX BOUNDS xy xz yz), scaled to the cubic spacing.
118            let s = 2.85;
119            let cell = Cell::from_lammps(
120                -6.2106057 * s,
121                24.8424228 * s,
122                0.0,
123                10.7570846 * s,
124                0.0,
125                20.131291 * s,
126                -6.2106057 * s,
127                0.0,
128                0.0,
129            )
130            .expect("tilt");
131            (frac_grid(&cell, n), cell)
132        }
133        other => panic!("unknown KNN_HOT_SHAPE {other}"),
134    }
135}
136
137fn main() {
138    let n = env_usize("KNN_HOT_N", 4096);
139    let k = env_usize("KNN_HOT_K", 4);
140    let reps = env_usize("KNN_HOT_REPS", 40);
141    let shape = std::env::var("KNN_HOT_SHAPE").unwrap_or_else(|_| "cubic".to_string());
142    let (xyz, cell) = system(&shape, n);
143    let mut out = vec![-1i32; xyz.len() * k];
144    let mut acc = 0i32;
145    let started = Instant::now();
146    for _ in 0..reps {
147        knearest_into(&xyz, &cell, k, None, Some(3.0), &mut out).expect("knearest");
148        acc = acc.wrapping_add(out[0]);
149    }
150    let ms = started.elapsed().as_secs_f64() * 1000.0;
151    std::hint::black_box(acc);
152    eprintln!(
153        "shape={shape} n={} k={} reps={} ms={ms:.1}",
154        xyz.len(),
155        k,
156        reps
157    );
158}
knn_hot/knn_hot.rs

knn_hot/
knn_hot.rs
