linkcell 0.3.11

Periodic linked-cell cutoff pairs and k-nearest neighbour search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

"""Autotune the cutoff walk's thresholds and paths with Kernel Tuner.

    cargo build --release --features tune
    RAYON_NUM_THREADS=8 python3 scripts/tune-pairs.py [--n 4096] [--jitter 0.3]

Kernel Tuner (https://github.com/KernelTuner/kernel_tuner) compiles one
small C function per configuration with its tunable parameters as
`#define`s. That function sets the knobs through `lc_tune_set` and times
`pairs_within` through `lc_tune_pairs`, which returns the milliseconds
per call, the row count, and a checksum of the rows. Every configuration
is checked against the default one: the same `(i, j, S)` rows, and the
same `dist2` sum to within 1e-9. The thread count is the pool's, set by
`RAYON_NUM_THREADS` before the library starts, so each count is its own
run. The problem is the probe's: a lattice of `n` points in an 18 Å cube
at a 4 Å cutoff, moved by `--jitter` times the lattice step at random.

Knobs: `cell_hint` (bin edge, 0 for the cutoff), `split_pairs` (expected
pairs where the walk splits across threads), `grid_atoms` (atoms from
which a split walk builds its bins on several threads), `fused` (one
thread writes a full list from the tile), `chunks` (bin ranges per
thread in a split search), and `ahead_from` (output bytes per thread
from which the row writers prefetch ahead of their stores). Only the
knobs that can act at the pool's thread count are searched.
"""

import argparse
import ctypes
import json
import os
import sys

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
LIBDIR = os.path.join(ROOT, "target", "release")
LIB = os.path.join(LIBDIR, "liblinkcell.so")

KERNEL = r"""
#include <stddef.h>
extern "C" int lc_tune_set(int, double);
extern "C" int lc_tune_pairs(const double *, size_t, const void *, double,
                             double, size_t, double *);

extern "C" float tune_pairs(const double *xyz, const double *simbox,
                            double *out, int n) {
    lc_tune_set(0, split_pairs);
    lc_tune_set(1, grid_atoms);
    lc_tune_set(2, fused);
    lc_tune_set(3, chunks);
    lc_tune_set(4, ahead_from);
    if (lc_tune_pairs(xyz, (size_t)n, simbox, @CUTOFF@, cell_hint, @REPS@, out) != 0)
        return 1.0e30f;
    return (float)out[0];
}
"""

DEFAULTS = {
    "split_pairs": 10_000,
    "grid_atoms": 1_024,
    "fused": 1,
    "chunks": 1,
    "ahead_from": 3 << 20,
    "cell_hint": 0.0,
}
KEYS = {"split_pairs": 0, "grid_atoms": 1, "fused": 2, "chunks": 3, "ahead_from": 4}


def lattice(n, boxl, jitter, seed):
    m = int(np.ceil(np.cbrt(n)))
    while m**3 < n:
        m += 1
    idx = np.arange(m**3)[:n]
    pts = np.stack([idx % m, (idx // m) % m, idx // (m * m)], axis=1).astype(np.float64)
    rng = np.random.default_rng(seed)
    pts += 0.5 + jitter * (rng.random(pts.shape) - 0.5)
    return np.ascontiguousarray(pts * boxl / m)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=4096)
    ap.add_argument("--box", type=float, default=18.0)
    ap.add_argument("--cutoff", type=float, default=4.0)
    ap.add_argument("--jitter", type=float, default=0.0)
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--reps", type=int, default=50, help="timed calls per Kernel Tuner run")
    ap.add_argument("--iterations", type=int, default=5, help="Kernel Tuner runs per configuration")
    ap.add_argument("--strategy", default="brute_force")
    ap.add_argument("--cache", default=None, help="Kernel Tuner cache file (JSON)")
    args = ap.parse_args()
    # Kernel Tuner writes its sources and objects in the working directory.
    os.chdir(os.environ.get("TMPDIR", "/tmp"))

    threads = int(os.environ.get("RAYON_NUM_THREADS", os.cpu_count()))
    # The library, and the worker pool inside it, stay loaded while Kernel
    # Tuner loads and unloads one wrapper per configuration.
    lib = ctypes.CDLL(LIB, mode=ctypes.RTLD_GLOBAL)
    lib.lc_tune_pairs.restype = ctypes.c_int
    lib.lc_tune_set.argtypes = [ctypes.c_int, ctypes.c_double]

    from kernel_tuner import tune_kernel

    n = args.n
    xyz = lattice(n, args.box, args.jitter, args.seed)
    simbox = np.array([args.box, 0, 0, 0, args.box, 0, 0, 0, args.box, 0, 0, 0], dtype=np.float64)
    out = np.zeros(5, dtype=np.float64)

    def call(knobs, hint, reps):
        for key, slot in KEYS.items():
            lib.lc_tune_set(slot, float(knobs[key]))
        res = np.zeros(5, dtype=np.float64)
        ok = lib.lc_tune_pairs(
            xyz.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_size_t(n),
            simbox.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_double(args.cutoff),
            ctypes.c_double(hint),
            ctypes.c_size_t(reps),
            res.ctypes.data_as(ctypes.c_void_p),
        )
        if ok != 0:
            sys.exit("lc_tune_pairs failed")
        return res

    ref = call(DEFAULTS, 0.0, args.reps)
    print(f"default: {ref[0]:.4f} ms per call, {int(ref[1])} rows, {threads} threads", flush=True)

    edges = [0.0] + [round(f * args.cutoff, 4) for f in (0.5, 0.5625, 0.6, 0.75, 0.9, 1.125, 1.25, 1.5)]
    tune_params = {"cell_hint": edges, "ahead_from": [0, 1 << 20, 2 << 20, 3 << 20, 6 << 20, 1 << 40]}
    if threads > 1:
        tune_params["split_pairs"] = [2_500, 5_000, 10_000, 20_000, 40_000]
        tune_params["grid_atoms"] = [512, 1_024, 2_048, 4_096, 8_192, 1 << 30]
        tune_params["chunks"] = [1, 2, 3, 4]
        tune_params["fused"] = [DEFAULTS["fused"]]
    else:
        tune_params["fused"] = [1, 0]
        tune_params["split_pairs"] = [DEFAULTS["split_pairs"]]
        tune_params["grid_atoms"] = [DEFAULTS["grid_atoms"]]
        tune_params["chunks"] = [DEFAULTS["chunks"]]

    def verify(answer, result, atol=None):
        want, got = answer[2], result[2]
        same_rows = np.array_equal(want[1:4], got[1:4])
        same_d2 = abs(want[4] - got[4]) <= 1e-9 * max(1.0, abs(want[4]))
        return bool(same_rows and same_d2)

    source = KERNEL.replace("@CUTOFF@", repr(args.cutoff)).replace("@REPS@", str(args.reps))
    results, _ = tune_kernel(
        "tune_pairs",
        source,
        n,
        [xyz, simbox, out, np.int32(n)],
        tune_params,
        lang="C",
        compiler_options=["-O2", f"-L{LIBDIR}", "-llinkcell", f"-Wl,-rpath,{LIBDIR}"],
        answer=[None, None, ref.copy(), None],
        verify=verify,
        iterations=args.iterations,
        strategy=args.strategy,
        cache=args.cache,
        quiet=True,
    )
    ok = [r for r in results if isinstance(r.get("time"), float)]
    ok.sort(key=lambda r: r["time"])
    print("best configurations (ms per call):")
    for r in ok[:8]:
        print("  " + json.dumps({k: r[k] for k in ["time"] + list(tune_params)}))
    best = ok[0]
    knobs = {k: best[k] for k in KEYS}
    check = call(knobs, best["cell_hint"], args.reps * 4)
    again = call(DEFAULTS, 0.0, args.reps * 4)
    print(
        f"recheck: best {check[0]:.4f} ms, default {again[0]:.4f} ms "
        f"({again[0] / check[0]:.2f}x), {len(results)} configurations"
    )


if __name__ == "__main__":
    main()