/ concept-collection / turing-sphere
Sign in
concept-collection / turing-sphere
413 lines · 16.2 KBBlameHistoryRaw
1/*
2 * The same run as `npm run bench`, on upstream SHTNS on the CPU (fp64).
3 *
4 * ./shtbench --preset schnak-spots --lmax 63 --steps 2000
5 * ./shtbench --mode transform --lmax 63 --steps 2000
6 *
7 * Two things are measured, selected by --mode:
8 *
9 * - solver: the IMEX Euler timestep of models/<key>.m — 4 transforms, a
10 * reaction on the grid, and the spectral update — which is what
11 * the app's `solver` number and `npm run bench` measure.
12 * - transform: one spectral -> grid -> spectral round trip and nothing else,
13 * which is the library-against-library number. The solver does
14 * one of these per species per step.
15 *
16 * This is fp64 throughout, because SHTNS' single precision exists only on the
17 * GPU. So it is not a like-for-like comparison against the fp32 WGSL
18 * transforms — it is the accuracy reference (how far has fp32 drifted?) and the
19 * "what does a well-optimized CPU do" reference. shtbench_gpu is the
20 * like-for-like one.
21 *
22 * The reaction is a transcription of the .m rather than the .m itself; see
23 * spec.h. It is checked, not trusted: `compare-native.mjs --check` diffs the
24 * final state against a WebGPU run of the actual .m.
25 */
26#include "spec.h"
28#include <shtns.h>
30static const char *USAGE =
31 "usage: ./shtbench [options]\n"
32 "\n"
33 " --mode solver|transform what to measure (default solver)\n"
34 " --preset <key> schnak-spots | schnak-coarse | schnak-fine | brussel |\n"
35 " allencahn (default schnak-spots)\n"
36 " --lmax <n> spherical harmonic truncation (default 63)\n"
37 " --steps <n> timed steps, or round trips (default 2000)\n"
38 " --warmup <n> untimed steps first (default 100)\n"
39 " --seed <n> seed of the initial noise / spectrum (default 1)\n"
40 " --threads <n> OpenMP threads (default: the library's choice)\n"
41 " --batch <n> accepted for symmetry with shtbench_gpu and ignored:\n"
42 " every CPU transform call is synchronous, so there is\n"
43 " nothing to batch\n"
44 " --layout theta|phi spatial layout: theta-contiguous is SHTNS' native and\n"
45 " fastest, phi-contiguous is what the WGSL side uses\n"
46 " (default theta)\n"
47 " --polar-eps <x> SHTNS polar-optimization threshold. 0 disables it, which\n"
48 " is what the WGSL transforms do; SHTNS' own default is\n"
49 " 1e-10 (default 0)\n"
50 " --digest after timing, re-run exactly --steps steps from the seed\n"
51 " and print a digest of the final spectral state\n"
52 " --dump-state <f> like --digest, and write the state to <f> as JSON, for\n"
53 " scripts/compare-native.mjs to diff\n"
54 " --<param> <v> any parameter of the preset's model, e.g. --dt 0.05\n"
55 " --json machine-readable output\n"
56 " --help\n"
57 "\n"
58 "Run the same spec through the WGSL transforms with:\n"
59 " npm run bench -- --preset <key> --lmax <n> --steps <n> (solver)\n"
60 " npm run bench:sht -- --lmax <n> --steps <n> (transform)";
62/*
63 * The layout arithmetic. SHTNS stores a spatial field either theta-contiguous
64 * (its native layout) or phi-contiguous (what the WGSL side uses), possibly
65 * with padding between lines. The reaction is pointwise so it does not care,
66 * but the seeded perturbation is indexed by (ilat, iphi) and does.
67 *
68 * Either way the field is `nlines` contiguous runs of `linelen` doubles,
69 * `stride` apart.
70 */
71struct Grid {
72 long nlat, nphi, nlat_padded;
73 int layout;
74 long nlines, linelen, stride;
75};
77static Grid grid_of(shtns_cfg sht, int layout) {
78 Grid g;
79 g.nlat = sht->nlat;
80 g.nphi = sht->nphi;
81 g.nlat_padded = sht->nlat_padded;
82 g.layout = layout;
83 if (layout == SHTB_LAYOUT_PHI) {
84 g.nlines = g.nlat;
85 g.linelen = g.nphi;
86 g.stride = g.nphi;
87 } else {
88 g.nlines = g.nphi;
89 g.linelen = g.nlat;
90 g.stride = g.nlat_padded;
91 }
92 return g;
95static inline long spat_index(const Grid &g, long ilat, long iphi) {
96 return g.layout == SHTB_LAYOUT_PHI ? ilat * g.nphi + iphi : iphi * g.nlat_padded + ilat;
99static void fill_uniform(const Grid &g, double *f, double value) {
100 for (long l = 0; l < g.nlines; l++)
101 for (long i = 0; i < g.linelen; i++) f[l * g.stride + i] = value;
104/* value + the seeded perturbation, which arrives in [ilat*nphi + iphi] order */
105static void fill_perturbed(const Grid &g, double *f, double value, const float *noise) {
106 for (long ilat = 0; ilat < g.nlat; ilat++)
107 for (long iphi = 0; iphi < g.nphi; iphi++)
108 f[spat_index(g, ilat, iphi)] = value + (double)noise[ilat * g.nphi + iphi];
111static void react_field(const Grid &g, const shtb_step_const<double> &c, const double *u,
112 const double *v, double *r1, double *r2) {
113#ifdef _OPENMP
114#pragma omp parallel for schedule(static)
115#endif
116 for (long l = 0; l < g.nlines; l++) {
117 const long o = l * g.stride;
118 for (long i = 0; i < g.linelen; i++) {
119 double a = 0, b = 0;
120 shtb_react<double>(c, u[o + i], v ? v[o + i] : 0.0, &a, &b);
121 r1[o + i] = a;
122 if (r2) r2[o + i] = b;
123 }
124 }
127static void imex_update(const shtb_step_const<double> &c, int nspecies, long nlm, double **U,
128 const double *const *R, const float *lam) {
129 const long n2 = 2 * nlm;
130 for (int k = 0; k < nspecies; k++) {
131#ifdef _OPENMP
132#pragma omp parallel for schedule(static)
133#endif
134 for (long i = 0; i < n2; i++)
135 U[k][i] = shtb_imex<double>(c, k, U[k][i], R[k][i], (double)lam[i]);
136 }
139static void field_range(const double *f, long nlines, long linelen, long stride, double *mn,
140 double *mx, int *finite) {
141 *mn = INFINITY;
142 *mx = -INFINITY;
143 *finite = 1;
144 for (long l = 0; l < nlines; l++)
145 for (long i = 0; i < linelen; i++) {
146 double x = f[l * stride + i];
147 if (x < *mn) *mn = x;
148 if (x > *mx) *mx = x;
149 if (!isfinite(x)) *finite = 0;
150 }
153int main(int argc, char **argv) {
154 shtb_spec spec;
155 double polar_eps = 0.0;
157 /* --polar-eps is ours, not part of the shared spec; take it out first. */
158 int argc2 = 0;
159 char **argv2 = (char **)malloc(sizeof(char *) * (size_t)argc);
160 for (int i = 0; i < argc; i++) {
161 if (strcmp(argv[i], "--polar-eps") == 0 && i + 1 < argc) {
162 polar_eps = atof(argv[++i]);
163 continue;
164 }
165 if (strncmp(argv[i], "--polar-eps=", 12) == 0) {
166 polar_eps = atof(argv[i] + 12);
167 continue;
168 }
169 argv2[argc2++] = argv[i];
170 }
171 int rc = shtb_parse_spec(argc2, argv2, &spec, USAGE);
172 free(argv2);
173 if (rc) return rc == 1 ? 0 : rc;
174 /* Every CPU transform call is synchronous, so there is nothing to batch. */
175 spec.batch = 1;
177 const int quiet = spec.json;
178 const int transform_mode = spec.mode == SHTB_MODE_TRANSFORM;
179 shtns_verbose(0);
180 const int threads = shtns_use_threads(spec.threads);
182 shtns_cfg sht = shtns_create(spec.lmax, spec.lmax, 1, sht_orthonormal);
183 if (!sht) {
184 fprintf(stderr, "shtbench: shtns_create failed\n");
185 return 1;
186 }
187 const int layout_flag =
188 spec.layout == SHTB_LAYOUT_PHI ? SHT_PHI_CONTIGUOUS : SHT_THETA_CONTIGUOUS;
189 if (shtns_set_grid(sht, (enum shtns_type)(sht_gauss | layout_flag | SHT_SCALAR_ONLY), polar_eps,
190 spec.nlat, spec.nphi) <= 0) {
191 fprintf(stderr, "shtbench: shtns_set_grid failed for lmax %d on a %dx%d grid\n", spec.lmax,
192 spec.nlat, spec.nphi);
193 return 1;
194 }
195 if ((int)sht->nlat != spec.nlat || (int)sht->nphi != spec.nphi) {
196 fprintf(stderr, "shtbench: SHTNS chose a %ux%u grid, not the %dx%d asked for\n", sht->nlat,
197 sht->nphi, spec.nlat, spec.nphi);
198 return 1;
199 }
201 const Grid g = grid_of(sht, spec.layout);
202 const long nlm = (long)sht->nlm;
203 const long nspat = (long)sht->nspat;
204 const int nsp = spec.model->nspecies;
206 /* Laplace-Beltrami eigenvalues, 2 x nlm with the value duplicated across the
207 * real and imaginary halves — the layout eigenvalues() builds in
208 * src/mgpu/model.ts. Held in float so both sides divide by the same number. */
209 float *lam = (float *)malloc(sizeof(float) * (size_t)(2 * nlm));
210 for (long lm = 0; lm < nlm; lm++) {
211 const int l = sht->li[lm];
212 lam[2 * lm] = lam[2 * lm + 1] = (float)(l * (l + 1));
213 }
215 double *spat[2] = {NULL, NULL};
216 double *rspat[2] = {NULL, NULL};
217 cplx *Q[2] = {NULL, NULL};
218 cplx *R[2] = {NULL, NULL};
219 for (int k = 0; k < nsp; k++) {
220 spat[k] = (double *)shtns_malloc(sizeof(double) * (size_t)nspat);
221 rspat[k] = (double *)shtns_malloc(sizeof(double) * (size_t)nspat);
222 Q[k] = (cplx *)shtns_malloc(sizeof(cplx) * (size_t)nlm);
223 R[k] = (cplx *)shtns_malloc(sizeof(cplx) * (size_t)nlm);
224 memset(spat[k], 0, sizeof(double) * (size_t)nspat);
225 memset(rspat[k], 0, sizeof(double) * (size_t)nspat);
226 /* through the double view: a cplx array is [re, im] pairs, and memset on
227 * std::complex itself is a non-trivial-type warning */
228 memset((double *)Q[k], 0, sizeof(double) * (size_t)(2 * nlm));
229 memset((double *)R[k], 0, sizeof(double) * (size_t)(2 * nlm));
230 }
231 float *noise = (float *)malloc(sizeof(float) * (size_t)(g.nlat * g.nphi));
232 float *state32 = (float *)malloc(sizeof(float) * (size_t)(2 * nlm));
234 const shtb_step_const<double> c = shtb_make_step_const<double>(spec.model, &spec.params);
235 double base[2];
236 shtb_background(spec.model, &spec.params, base);
238 /* --- solver: init and one timestep, from models/<key>.m ------------------ */
239 auto seed_state = [&]() {
240 shtb_seeded_noise(g.nlat * g.nphi, spec.model->seed_amp, (uint32_t)spec.seed, noise);
241 fill_perturbed(g, spat[0], base[0], noise);
242 spat_to_SH(sht, spat[0], Q[0]);
243 if (nsp > 1) {
244 fill_uniform(g, spat[1], base[1]);
245 spat_to_SH(sht, spat[1], Q[1]);
246 }
247 };
248 auto step = [&]() {
249 for (int k = 0; k < nsp; k++) SH_to_spat(sht, Q[k], spat[k]);
250 react_field(g, c, spat[0], nsp > 1 ? spat[1] : NULL, rspat[0], nsp > 1 ? rspat[1] : NULL);
251 for (int k = 0; k < nsp; k++) spat_to_SH(sht, rspat[k], R[k]);
252 double *Ud[2] = {(double *)Q[0], (double *)Q[1]};
253 const double *Rd[2] = {(const double *)R[0], (const double *)R[1]};
254 imex_update(c, nsp, nlm, Ud, Rd, lam);
255 };
257 /* --- transform: one synth + one analys, ping-ponging the two buffers ----- */
258 cplx *tq[2] = {Q[0], R[0]};
259 int tcur = 0;
260 auto seed_spectrum = [&]() {
261 shtb_seeded_spectrum(spec.lmax, spec.lmax, (uint32_t)spec.seed, state32);
262 tcur = 0;
263 double *q = (double *)tq[0];
264 for (long i = 0; i < 2 * nlm; i++) q[i] = (double)state32[i];
265 };
266 auto round_trip = [&]() {
267 SH_to_spat(sht, tq[tcur], spat[0]);
268 spat_to_SH(sht, spat[0], tq[tcur ^ 1]);
269 tcur ^= 1;
270 };
272 if (transform_mode)
273 seed_spectrum();
274 else
275 seed_state();
277 if (!quiet) {
278 printf("shtbench — upstream SHTNS on the CPU, %s only\n\n",
279 transform_mode ? "transforms" : "solver");
280 printf(" mode %s\n", transform_mode
281 ? "transform (one synth + one analys per step)"
282 : "solver (one IMEX Euler timestep per step)");
283 if (!transform_mode) {
284 printf(" preset %s (models/%s.m: %d species)\n", spec.preset->label, spec.model->key,
285 nsp);
286 printf(" params ");
287 for (int i = 0; i < spec.model->nparams; i++)
288 printf("%s=%g ", spec.model->params[i].key,
289 *shtb_field_c(&spec.params, spec.model->params[i].off));
290 printf("\n");
291 }
292 printf(" grid lmax %d · %ldx%ld · nlm %ld\n", spec.lmax, g.nlat, g.nphi, nlm);
293 printf(" layout %s%s\n",
294 spec.layout == SHTB_LAYOUT_PHI ? "phi-contiguous" : "theta-contiguous (native)",
295 (long)sht->nlat_padded != g.nlat ? ", padded" : "");
296 printf(" backend %s\n fp64, %d thread%s, polar opt %g\n",
297 shtns_get_build_info(), threads, threads == 1 ? "" : "s", polar_eps);
298 printf(" run %d warmup + %d timed steps, seed %d\n\n", spec.warmup, spec.steps,
299 spec.seed);
300 }
302 for (int i = 0; i < spec.warmup; i++) {
303 if (transform_mode)
304 round_trip();
305 else
306 step();
307 }
309 double *samples = (double *)malloc(sizeof(double) * (size_t)spec.steps);
310 const double t0 = shtb_now_ms();
311 for (int i = 0; i < spec.steps; i++) {
312 const double a = shtb_now_ms();
313 if (transform_mode)
314 round_trip();
315 else
316 step();
317 samples[i] = shtb_now_ms() - a;
318 }
319 const double total = shtb_now_ms() - t0;
321 shtb_report rep;
322 memset(&rep, 0, sizeof(rep));
323 char libbuf[192], adapterbuf[64];
324 rep.library = shtb_json_safe(libbuf, sizeof(libbuf), shtns_get_build_info());
325 rep.runtime = "cpu";
326 snprintf(adapterbuf, sizeof(adapterbuf), "CPU, %d thread%s", threads, threads == 1 ? "" : "s");
327 rep.adapter = adapterbuf;
328 rep.precision = "fp64";
329 rep.fourier = "fftw";
330 rep.nlm = nlm;
331 rep.ops_per_step = transform_mode ? 2 : 2 * nsp + 2;
332 rep.ms_per_step = total / spec.steps;
333 rep.encode_ms_per_step = 0; /* nothing is deferred: every call is synchronous */
334 rep.latency = shtb_stats(samples, spec.steps);
335 rep.have_latency = 1;
336 /* the range reported below is the state as it stands now: warmup included */
337 rep.steps_run = spec.warmup + spec.steps;
338 rep.model_t = transform_mode ? 0 : rep.steps_run * spec.params.dt;
340 /* Did the run stay finite and develop contrast? The app shows the same range
341 * for the first species under its stats line. */
342 if (transform_mode) {
343 field_range((const double *)tq[tcur], 1, 2 * nlm, 0, &rep.field_min, &rep.field_max,
344 &rep.finite);
345 } else {
346 for (int k = 0; k < nsp; k++) SH_to_spat(sht, Q[k], spat[k]);
347 field_range(spat[0], g.nlines, g.linelen, g.stride, &rep.field_min, &rep.field_max,
348 &rep.finite);
349 }
351 /* A reproducible state to compare against a WebGPU run: exactly --steps steps
352 * from the seed, separate from the timed run above (which has warmup in it). */
353 if (spec.digest) {
354 if (transform_mode) {
355 seed_spectrum();
356 rep.input_digest = shtb_digest_of(state32, 2 * nlm);
357 rep.have_input_digest = 1;
358 for (int i = 0; i < spec.steps; i++) round_trip();
359 } else {
360 seed_state();
361 for (int i = 0; i < spec.steps; i++) step();
362 }
363 const double *q = (const double *)(transform_mode ? tq[tcur] : Q[0]);
364 for (long i = 0; i < 2 * nlm; i++) state32[i] = (float)q[i];
365 rep.digest = shtb_digest_of(state32, 2 * nlm);
366 rep.have_digest = 1;
367 }
369 if (spec.json) {
370 shtb_print_json(&spec, &rep);
371 } else {
372 printf(" %.3f ms/step %.1f steps/s", rep.ms_per_step, 1000.0 / rep.ms_per_step);
373 if (!transform_mode) printf(" %.2f model time/s", spec.params.dt * 1000.0 / rep.ms_per_step);
374 printf("\n");
375 printf(" per step: %.3f ms mean · median %.3f · p05 %.3f · p95 %.3f · min %.3f\n",
376 rep.latency.mean_ms, rep.latency.median_ms, rep.latency.p05_ms, rep.latency.p95_ms,
377 rep.latency.min_ms);
378 if (transform_mode)
379 printf(" i.e. %.3f ms per single transform\n", rep.ms_per_step / 2);
380 else
381 printf(" after %d steps: t = %.2f, field ∈ [%.4f, %.4f] (contrast %.4f)%s\n",
382 rep.steps_run, rep.model_t, rep.field_min, rep.field_max,
383 rep.field_max - rep.field_min, rep.finite ? "" : " — NOT FINITE");
384 if (rep.have_digest) {
385 printf("\n state after %d steps from seed %d:\n", spec.steps, spec.seed);
386 printf(" n=%ld min=%.9g max=%.9g mean=%.9g rms=%.9g\n", rep.digest.n, rep.digest.min,
387 rep.digest.max, rep.digest.mean, rep.digest.rms);
388 }
389 printf("\n This is fp64. The like-for-like fp32 comparison against the WGSL\n"
390 " transforms is ./shtbench_gpu; this run is the accuracy reference.\n");
391 }
393 if (spec.dump_state && rep.have_digest) {
394 if (shtb_dump_state(spec.dump_state, &spec, &rep, state32, 2 * nlm) != 0) {
395 fprintf(stderr, "shtbench: cannot write %s\n", spec.dump_state);
396 return 1;
397 }
398 if (!spec.json) printf("\n wrote %s\n", spec.dump_state);
399 }
401 free(samples);
402 free(noise);
403 free(state32);
404 free(lam);
405 for (int k = 0; k < nsp; k++) {
406 shtns_free(spat[k]);
407 shtns_free(rspat[k]);
408 shtns_free(Q[k]);
409 shtns_free(R[k]);
410 }
411 shtns_destroy(sht);
412 return rep.finite ? 0 : 1;
moveopenescclose