# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find your NUMA effect: first touch.

Allocates an array and scans it with all cores under three setups: (a) one
thread initializes it, then all threads scan; (b) it is initialized in
parallel with the same split as the scan; (c) its pages are interleaved
across NUMA nodes (set_mempolicy, the call behind numactl --interleave=all).
Prints the machine's NUMA nodes and the scan bandwidth of each setup.
Linux only (it reads /sys and calls set_mempolicy); needs cc.
"""
import argparse, json, os, subprocess, tempfile

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--gb", type=float, default=1.0, help="Array size (GB).")
p.add_argument("--threads", type=int, default=os.cpu_count(), help="Threads.")
p.add_argument("--repeats", type=int, default=5, help="Scans per setup (the best is reported).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

C = r"""
#define _GNU_SOURCE
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/syscall.h>
#include <time.h>
#include <unistd.h>
static long *arr; static size_t n; static int T;
static double now(void){ struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec / 1e9; }
static void *init(void *a){ long id = (long)a; size_t lo = n * id / T, hi = n * (id + 1) / T; for (size_t i = lo; i < hi; i++) arr[i] = i; return NULL; }
static volatile long sums[256];
static void *scan(void *a){ long id = (long)a; size_t lo = n * id / T, hi = n * (id + 1) / T; long s = 0; for (size_t i = lo; i < hi; i++) s += arr[i]; sums[id * 8 % 256] = s; return NULL; }
static void run(void *(*f)(void *), int threads){ pthread_t th[256]; int t0 = T; if (threads == 1) T = 1;
  for (long i = 0; i < T; i++) pthread_create(&th[i], NULL, f, (void *)i); for (int i = 0; i < T; i++) pthread_join(th[i], NULL); T = t0; }
int main(int argc, char **argv){
  double gb = atof(argv[1]); T = atoi(argv[2]); int reps = atoi(argv[3]); char mode = argv[4][0];
  n = (size_t)(gb * 1e9 / sizeof(long));
  if (mode == 'c') { unsigned long mask = 0; FILE *f = fopen("/sys/devices/system/node/online", "r"); int a = 0, b = 0;
    if (f && fscanf(f, "%d-%d", &a, &b) >= 1) { if (b < a) b = a; for (int i = a; i <= b; i++) mask |= 1UL << i; } else mask = 1;
    if (f) fclose(f);
    if (syscall(SYS_set_mempolicy, 3 /* MPOL_INTERLEAVE */, &mask, sizeof(mask) * 8)) perror("set_mempolicy"); }
  arr = malloc(n * sizeof(long));
  run(init, mode == 'a' ? 1 : T);                                  /* first touch decides where each page lives */
  double best = 1e9; for (int r = 0; r < reps; r++) { double t = now(); run(scan, T); double e = now() - t; if (e < best) best = e; }
  long total = 0; for (int i = 0; i < 256; i++) total += sums[i];
  fprintf(stderr, "%ld\n", total);                                 /* read the sums, so the scans cannot be removed */
  printf("%f\n", n * sizeof(long) / best / 1e9);
}
"""
d = tempfile.mkdtemp(); src = os.path.join(d, "numa.c"); exe = os.path.join(d, "numa"); open(src, "w").write(C)
subprocess.run(["cc", "-O2", "-pthread", src, "-o", exe], check=True)
nodes = open("/sys/devices/system/node/online").read().strip() if os.path.exists("/sys/devices/system/node/online") else "unknown"
node_count = len(os.listdir("/sys/devices/system/node")) if os.path.isdir("/sys/devices/system/node") else 0
node_count = sum(1 for x in os.listdir("/sys/devices/system/node") if x.startswith("node") and x[4:].isdigit()) if node_count else 1
rows = []
for mode, label in (("a", "(a) one thread initializes, all scan"), ("b", "(b) parallel first touch, same split"), ("c", "(c) interleaved across nodes")):
    gbs = float(subprocess.run([exe, str(args.gb), str(args.threads), str(args.repeats), mode], capture_output=True, text=True, check=True).stdout)
    rows.append({"setup": label, "gb_per_s": gbs})
for r in rows: r["vs_b_pct"] = (rows[1]["gb_per_s"] / r["gb_per_s"] - 1) * 100
summary = {"numa_nodes": node_count, "nodes_online": nodes, "threads": args.threads, "spread_pct": (max(r["gb_per_s"] for r in rows) / min(r["gb_per_s"] for r in rows) - 1) * 100}
print(f"NUMA nodes: {node_count} (online: {nodes}), {args.threads} threads, {args.gb:g} GB array")
for r in rows: print(f"{r['setup']:<40} {r['gb_per_s']:6.2f} GB/s  (scan time {r['vs_b_pct']:+.1f}% vs parallel first touch)")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"setups": rows}}, f)