knn_scale/
knn_scale.rs

1//! Strong-scaling probe for one fixed problem.
2//!
3//! The box grows with `n` so the spacing stays 3.125, the cubic hot-path
4//! density. Set `RAYON_NUM_THREADS` before the process starts. Env:
5//! `KNN_SCALE_N` (default 262144), `KNN_SCALE_K` (default 4),
6//! `KNN_SCALE_REPS` (default 6), `KNN_SCALE_SHAPE` (`cubic` or `tilt`),
7//! `KNN_SCALE_HINT` (bin edge, default 3.0).
8//! `KNN_SCALE_SHUFFLE` (a seed) permutes the atoms, so their order no longer follows
9//! their position, as in a snapshot after a run.
10
11use std::time::Instant;
12
13use linkcell::{knearest_into, Cell};
14
15fn env_usize(key: &str, default: usize) -> usize {
16    std::env::var(key)
17        .ok()
18        .and_then(|s| s.parse().ok())
19        .unwrap_or(default)
20}
21
22fn nside_of(n: usize) -> usize {
23    ((n as f64).cbrt().round() as usize).max(1)
24}
25
26fn cubic(nside: usize) -> (Vec<[f64; 3]>, Cell) {
27    let boxl = nside as f64 * 3.125;
28    let cell = Cell::ortho(boxl, boxl, boxl).expect("box");
29    let mut xyz = Vec::with_capacity(nside * nside * nside);
30    let a = 3.125;
31    for iz in 0..nside {
32        for iy in 0..nside {
33            for ix in 0..nside {
34                xyz.push([ix as f64 * a, iy as f64 * a, iz as f64 * a]);
35            }
36        }
37    }
38    (xyz, cell)
39}
40
41/// Restricted triclinic box, tilt-reduced, scaled so the a edge matches
42/// the cubic spacing times `nside`.
43fn tilt(nside: usize) -> (Vec<[f64; 3]>, Cell) {
44    let lx = nside as f64 * 3.125;
45    let cell = Cell::from_vectors(
46        [lx, 0.0, 0.0],
47        [lx * 0.2, lx * 0.9, 0.0],
48        [lx * 0.05, lx * -0.08, lx * 0.95],
49        [0.0, 0.0, 0.0],
50    )
51    .expect("tilt");
52    let mut xyz = Vec::with_capacity(nside * nside * nside);
53    let nf = nside as f64;
54    for iz in 0..nside {
55        for iy in 0..nside {
56            for ix in 0..nside {
57                let s = [
58                    (ix as f64 + 0.5) / nf,
59                    (iy as f64 + 0.5) / nf,
60                    (iz as f64 + 0.5) / nf,
61                ];
62                xyz.push(cell.cartesian(s));
63            }
64        }
65    }
66    (xyz, cell)
67}
68
69/// Fisher-Yates with an xorshift stream from `seed`.
70fn shuffle(xyz: &mut [[f64; 3]], seed: u64) {
71    let mut s = seed.max(1);
72    for i in (1..xyz.len()).rev() {
73        s ^= s << 13;
74        s ^= s >> 7;
75        s ^= s << 17;
76        xyz.swap(i, (s % (i as u64 + 1)) as usize);
77    }
78}
79
80fn main() {
81    let nside = nside_of(env_usize("KNN_SCALE_N", 262_144));
82    let k = env_usize("KNN_SCALE_K", 4);
83    let reps = env_usize("KNN_SCALE_REPS", 6);
84    let shape = std::env::var("KNN_SCALE_SHAPE").unwrap_or_else(|_| "cubic".to_string());
85    let hint: f64 = std::env::var("KNN_SCALE_HINT")
86        .ok()
87        .and_then(|s| s.parse().ok())
88        .unwrap_or(3.0);
89    let (mut xyz, cell) = match shape.as_str() {
90        "cubic" => cubic(nside),
91        "tilt" => tilt(nside),
92        other => panic!("unknown KNN_SCALE_SHAPE {other}"),
93    };
94    if let Some(seed) = std::env::var("KNN_SCALE_SHUFFLE")
95        .ok()
96        .and_then(|s| s.parse().ok())
97    {
98        shuffle(&mut xyz, seed);
99    }
100    let threads = std::env::var("RAYON_NUM_THREADS").unwrap_or_else(|_| "default".to_string());
101    let mut out = vec![-1i32; xyz.len() * k];
102    linkcell::pop_engage();
103    linkcell::pop_prepare();
104    knearest_into(&xyz, &cell, k, None, Some(hint), &mut out).expect("warmup");
105    linkcell::pop_reset();
106    let mut acc = 0i32;
107    let started = Instant::now();
108    for _ in 0..reps {
109        knearest_into(&xyz, &cell, k, None, Some(hint), &mut out).expect("knearest");
110        acc = acc.wrapping_add(out[0]);
111    }
112    let wall_ns = started.elapsed().as_nanos() as u64;
113    std::hint::black_box(acc);
114    let useful = linkcell::pop_snapshot();
115    let ms = wall_ns as f64 / 1e6;
116    println!(
117        "shape={shape} threads={threads} n={} k={k} reps={reps} ms={ms:.1} ms_per={:.2}",
118        xyz.len(),
119        ms / reps as f64
120    );
121    let (lb, ce, pe) = linkcell::pop_efficiencies(&useful, wall_ns);
122    let list = useful
123        .iter()
124        .map(|ns| ns.to_string())
125        .collect::<Vec<_>>()
126        .join(",");
127    println!(
128        "pop shape={shape} threads={threads} n={} reps={reps} wall_ns={wall_ns} useful_ns={list} lb={lb:.4} ce={ce:.4} pe={pe:.4}",
129        xyz.len()
130    );
131}
knn_scale/knn_scale.rs

knn_scale/
knn_scale.rs
