2 * The same run as `npm run bench`, on upstream SHTNS' own CUDA transforms.
3 *
4 * ./shtbench_gpu --preset schnak-spots --lmax 63 --steps 2000
5 * ./shtbench_gpu --mode transform --lmax 63 --steps 2000
6 *
7 * This is the like-for-like comparison the WGSL transforms exist to be measured
8 * against: single precision, everything resident on the GPU, nothing read back
9 * inside the loop. The only difference between this and `npm run bench` is what
10 * runs the transforms — SHTNS' hand-written CUDA kernels and cuFFT/VkFFT here,
11 * generated WGSL and a WGSL FFT there — on the same device.
12 *
13 * It keeps the state on the GPU the same way the WebGPU side does: cu_* are the
14 * on-device entry points, asynchronous on SHTNS' compute stream, and a batch of
15 * steps is launched before anything is waited for. `--batch` matches
16 * `npm run bench --batch`, so both sides can be made to synchronize equally
17 * often.
18 *
19 * Two things are measured, selected by --mode:
20 *
21 * - solver: one IMEX Euler timestep of models/<key>.m — 2*nspecies
22 * transforms, one reaction kernel on the grid, one spectral
23 * update kernel.
24 * - transform: one spectral -> grid -> spectral round trip and nothing else.
25 *
26 * Two SHTNS details worth knowing when reading the numbers:
27 *
28 * - in fp32 mode SHTNS runs the *Legendre recurrence* in fp64 when
29 * lmax <= 128 and the GPU has usable fp64 (SHT_L_RESCALE_FLY_FLOAT in
30 * sht_private.h), which WebGPU cannot do at all. Set
31 * SHTNS_GPU_REC_PREC=1 to force the recurrence into fp32 and get the closer
32 * comparison; the run prints which it got.
33 * - the spatial layout defaults to theta-contiguous, which is SHTNS' native
34 * and fastest. --layout phi matches what the WGSL side uses. The spectral
35 * layout and normalization are identical either way, so a state comparison
36 * is valid in both.
37 */
38#include "spec.h"
40#include <cuda_runtime.h>
41#include <shtns.h>
42#include <shtns_cuda.h>
44static const char *USAGE =
45 "usage: ./shtbench_gpu [options]\n"
46 "\n"
47 " --mode solver|transform what to measure (default solver)\n"
48 " --preset <key> schnak-spots | schnak-coarse | schnak-fine | brussel |\n"
49 " allencahn (default schnak-spots)\n"
50 " --lmax <n> spherical harmonic truncation (default 63)\n"
51 " --steps <n> timed steps, or round trips (default 2000)\n"
52 " --warmup <n> untimed steps first (default 100)\n"
53 " --batch <n> steps launched per synchronization (default 16), as in\n"
54 " `npm run bench -- --batch`\n"
55 " --seed <n> seed of the initial noise / spectrum (default 1)\n"
56 " --fp64 use SHTNS' double-precision GPU transforms instead of\n"
57 " single. Not comparable to WebGPU, which has no fp64;\n"
58 " useful as an accuracy and cost reference\n"
59 " --layout theta|phi spatial layout: theta-contiguous is SHTNS' native and\n"
60 " fastest, phi-contiguous is what the WGSL side uses\n"
61 " (default theta)\n"
62 " --polar-eps <x> SHTNS polar-optimization threshold (default 0)\n"
63 " --device <n> CUDA device index (default 0)\n"
64 " --digest after timing, re-run exactly --steps steps from the seed\n"
65 " and print a digest of the final spectral state\n"
66 " --dump-state <f> like --digest, and write the state to <f> as JSON, for\n"
67 " scripts/compare-native.mjs to diff\n"
68 " --<param> <v> any parameter of the preset's model, e.g. --dt 0.05\n"
69 " --json machine-readable output\n"
70 " --help\n"
71 "\n"
72 "Run the same spec through the WGSL transforms with:\n"
73 " npm run bench -- --preset <key> --lmax <n> --steps <n> (solver)\n"
74 " npm run bench:sht -- --lmax <n> --steps <n> (transform)";
76#define CU_CHECK(call) \
77 do { \
78 cudaError_t err_ = (call); \
79 if (err_ != cudaSuccess) { \
80 fprintf(stderr, "shtbench_gpu: %s failed at %s:%d: %s\n", #call, __FILE__, __LINE__, \
81 cudaGetErrorString(err_)); \
82 exit(1); \
83 } \
84 } while (0)
86/* ------------------------------------------------------------------- kernels
87 *
88 * The counterparts of the generated WGSL kernels: one thread per output
89 * element, reading the same inputs and doing the same arithmetic (see
90 * shtb_react / shtb_imex in spec.h, transcribed from the .m).
91 */
93template <typename real>
94__global__ void k_react(shtb_step_const<real> c, const real *u, const real *v, real *r1, real *r2,
95 long npts) {
96 const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
97 if (i >= npts) return;
98 real a = 0, b = 0;
99 shtb_react<real>(c, u[i], v ? v[i] : (real)0, &a, &b);
100 r1[i] = a;
101 if (r2) r2[i] = b;
102}
104/* Latitude lines are `stride` apart in the theta-contiguous layout, so the flat
105 * kernel above would step over padding. This one is used when there is any. */
106template <typename real>
107__global__ void k_react_strided(shtb_step_const<real> c, const real *u, const real *v, real *r1,
108 real *r2, long linelen, long stride) {
109 const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
110 if (i >= linelen) return;
111 const long o = (long)blockIdx.y * stride + i;
112 real a = 0, b = 0;
113 shtb_react<real>(c, u[o], v ? v[o] : (real)0, &a, &b);
114 r1[o] = a;
115 if (r2) r2[o] = b;
116}
118/* One thread per real, over the 2 x nlm spectral layout — `lam` carries l(l+1)
119 * duplicated across the real and imaginary halves, exactly as on the WGSL side. */
120template <typename real>
121__global__ void k_imex(shtb_step_const<real> c, int k, real *U, const real *R, const real *lam,
122 long n2) {
123 const long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
124 if (i >= n2) return;
125 U[i] = shtb_imex<real>(c, k, U[i], R[i], lam[i]);
126}
128/* Scatter the seeded perturbation, which is generated in [ilat*nphi + iphi]
129 * order, into whichever layout SHTNS is using. Also used to fill the uniform
130 * background. */
131template <typename real>
132__global__ void k_fill(real *f, real value, const float *noise, long nlat, long nphi,
133 long stride_lat, long stride_phi) {
134 const long ilat = (long)blockIdx.y;
135 const long iphi = (long)blockIdx.x * blockDim.x + threadIdx.x;
136 if (iphi >= nphi || ilat >= nlat) return;
137 const real n = noise ? (real)noise[ilat * nphi + iphi] : (real)0;
138 f[ilat * stride_lat + iphi * stride_phi] = value + n;
139}
141/* --------------------------------------------------------------------- setup */
143/* SHTNS' own device-buffer sizes, from init_cuda_buffer_fft() in sht_gpu.cu.
144 * Its kernels write a little past nlm ("one more data per m") and the Fourier
145 * stage needs room for the R2C form, so allocating exactly nlm or nlat*nphi is
146 * not enough. WARPSZE is 32 on every CUDA GPU SHTNS supports. */
147static long spec_alloc_reals(long nlm, int mmax) {
148 const long nlm2 = nlm + (mmax + 1);
149 return ((2 * nlm2 + 31) / 32) * 32;
150}
151static long spat_alloc_reals(long nlat_padded, long nphi, int mmax) {
152 const long extra = (nphi / 2 == mmax) ? 1 : 0;
153 return ((nlat_padded * (nphi + extra) + 31) / 32) * 32;
154}
156struct Layout {
157 long nlat, nphi, nlat_padded;
158 long stride_lat, stride_phi;
159 long linelen, nlines, stride; /* contiguous runs, for the reaction kernel */
160 int padded;
161};
163static Layout layout_of(shtns_cfg sht, int which) {
164 Layout l;
165 l.nlat = sht->nlat;
166 l.nphi = sht->nphi;
167 l.nlat_padded = sht->nlat_padded;
168 if (which == SHTB_LAYOUT_PHI) {
169 l.stride_lat = l.nphi;
170 l.stride_phi = 1;
171 l.nlines = l.nlat;
172 l.linelen = l.nphi;
173 l.stride = l.nphi;
174 l.padded = 0;
175 } else {
176 l.stride_lat = 1;
177 l.stride_phi = l.nlat_padded;
178 l.nlines = l.nphi;
179 l.linelen = l.nlat;
180 l.stride = l.nlat_padded;
181 l.padded = l.nlat_padded != l.nlat;
182 }
183 return l;
184}
186/* --------------------------------------------------------------- the run body
187 *
188 * Templated on the transform precision so fp32 and fp64 are the same code. The
189 * fp32 instantiation is the one that matters; fp64 is there as a reference.
190 */
191template <typename real>
192struct Run {
193 shtns_cfg sht;
194 const shtb_spec *spec;
195 /* The stream SHTNS was told to compute on, so our kernels are ordered against
196 * its transforms and one synchronization waits for the whole step. */
197 cudaStream_t stream;
198 Layout lay;
199 long nlm, n2, npts;
200 int nsp;
201 shtb_step_const<real> c;
202 double base[2];
204 real *dU[2], *dR[2], *dspat[2], *drspat[2], *dlam;
205 real *dTq[2]; /* transform mode ping-pong */
206 int tcur;
207 float *dnoise;
208 float *hnoise;
209 float *hstate;
211 void sync() { CU_CHECK(cudaStreamSynchronize(stream)); }
213 void alloc() {
214 const long spec_n = spec_alloc_reals(nlm, sht->mmax);
215 long spat_n = spat_alloc_reals(lay.nlat_padded, lay.nphi, sht->mmax);
216 if ((long)sht->nspat > spat_n) spat_n = (long)sht->nspat;
217 for (int k = 0; k < nsp; k++) {
218 CU_CHECK(cudaMalloc(&dU[k], sizeof(real) * (size_t)spec_n));
219 CU_CHECK(cudaMalloc(&dR[k], sizeof(real) * (size_t)spec_n));
220 CU_CHECK(cudaMalloc(&dspat[k], sizeof(real) * (size_t)spat_n));
221 CU_CHECK(cudaMalloc(&drspat[k], sizeof(real) * (size_t)spat_n));
222 CU_CHECK(cudaMemset(dU[k], 0, sizeof(real) * (size_t)spec_n));
223 CU_CHECK(cudaMemset(dR[k], 0, sizeof(real) * (size_t)spec_n));
224 CU_CHECK(cudaMemset(dspat[k], 0, sizeof(real) * (size_t)spat_n));
225 CU_CHECK(cudaMemset(drspat[k], 0, sizeof(real) * (size_t)spat_n));
226 }
227 dTq[0] = dU[0];
228 dTq[1] = dR[0];
229 tcur = 0;
230 CU_CHECK(cudaMalloc(&dlam, sizeof(real) * (size_t)n2));
231 CU_CHECK(cudaMalloc(&dnoise, sizeof(float) * (size_t)npts));
232 hnoise = (float *)malloc(sizeof(float) * (size_t)npts);
233 hstate = (float *)malloc(sizeof(float) * (size_t)(n2));
235 real *lam = (real *)malloc(sizeof(real) * (size_t)n2);
236 for (long lm = 0; lm < nlm; lm++) {
237 const int l = sht->li[lm];
238 lam[2 * lm] = lam[2 * lm + 1] = (real)(l * (l + 1));
239 }
240 CU_CHECK(cudaMemcpy(dlam, lam, sizeof(real) * (size_t)n2, cudaMemcpyHostToDevice));
241 free(lam);
242 }
244 void free_all() {
245 for (int k = 0; k < nsp; k++) {
246 cudaFree(dU[k]);
247 cudaFree(dR[k]);
248 cudaFree(dspat[k]);
249 cudaFree(drspat[k]);
250 }
251 cudaFree(dlam);
252 cudaFree(dnoise);
253 free(hnoise);
254 free(hstate);
255 }
257 void synth(real *qlm, real *spat);
258 void analys(real *spat, real *qlm);
260 void seed_state() {
261 shtb_seeded_noise(npts, spec->model->seed_amp, (uint32_t)spec->seed, hnoise);
262 CU_CHECK(cudaMemcpy(dnoise, hnoise, sizeof(float) * (size_t)npts, cudaMemcpyHostToDevice));
263 const int tpb = 256;
264 dim3 grid((unsigned)((lay.nphi + tpb - 1) / tpb), (unsigned)lay.nlat);
265 k_fill<real><<<grid, tpb, 0, stream>>>(dspat[0], (real)base[0], dnoise, lay.nlat, lay.nphi,
266 lay.stride_lat, lay.stride_phi);
267 analys(dspat[0], dU[0]);
268 if (nsp > 1) {
269 k_fill<real><<<grid, tpb, 0, stream>>>(dspat[1], (real)base[1], NULL, lay.nlat, lay.nphi,
270 lay.stride_lat, lay.stride_phi);
271 analys(dspat[1], dU[1]);
272 }
273 sync();
274 }
276 /* One IMEX Euler timestep. Nothing is synchronized: the calls queue on SHTNS'
277 * compute stream, which is what makes a batch of steps one submission's worth
278 * of work, as on the WebGPU side. */
279 void step() {
280 const int tpb = 256;
281 for (int k = 0; k < nsp; k++) synth(dU[k], dspat[k]);
282 if (lay.padded) {
283 dim3 grid((unsigned)((lay.linelen + tpb - 1) / tpb), (unsigned)lay.nlines);
284 k_react_strided<real><<<grid, tpb, 0, stream>>>(c, dspat[0], nsp > 1 ? dspat[1] : NULL,
285 drspat[0], nsp > 1 ? drspat[1] : NULL,
286 lay.linelen, lay.stride);
287 } else {
288 k_react<real><<<(unsigned)((npts + tpb - 1) / tpb), tpb, 0, stream>>>(
289 c, dspat[0], nsp > 1 ? dspat[1] : NULL, drspat[0], nsp > 1 ? drspat[1] : NULL, npts);
290 }
291 for (int k = 0; k < nsp; k++) analys(drspat[k], dR[k]);
292 for (int k = 0; k < nsp; k++)
293 k_imex<real><<<(unsigned)((n2 + tpb - 1) / tpb), tpb, 0, stream>>>(c, k, dU[k], dR[k], dlam,
294 n2);
295 }
297 void seed_spectrum() {
298 shtb_seeded_spectrum(spec->lmax, spec->lmax, (uint32_t)spec->seed, hstate);
299 tcur = 0;
300 if (sizeof(real) == sizeof(float)) {
301 CU_CHECK(cudaMemcpy(dTq[0], hstate, sizeof(float) * (size_t)n2, cudaMemcpyHostToDevice));
302 } else {
303 double *tmp = (double *)malloc(sizeof(double) * (size_t)n2);
304 for (long i = 0; i < n2; i++) tmp[i] = (double)hstate[i];
305 CU_CHECK(cudaMemcpy(dTq[0], tmp, sizeof(double) * (size_t)n2, cudaMemcpyHostToDevice));
306 free(tmp);
307 }
308 sync();
309 }
311 void round_trip() {
312 synth(dTq[tcur], dspat[0]);
313 analys(dspat[0], dTq[tcur ^ 1]);
314 tcur ^= 1;
315 }
317 /* the final spectral state of species 0 (or of the round trip), as float */
318 const float *read_state() {
319 sync();
320 const real *src = spec->mode == SHTB_MODE_TRANSFORM ? dTq[tcur] : dU[0];
321 if (sizeof(real) == sizeof(float)) {
322 CU_CHECK(cudaMemcpy(hstate, src, sizeof(float) * (size_t)n2, cudaMemcpyDeviceToHost));
323 } else {
324 double *tmp = (double *)malloc(sizeof(double) * (size_t)n2);
325 CU_CHECK(cudaMemcpy(tmp, src, sizeof(double) * (size_t)n2, cudaMemcpyDeviceToHost));
326 for (long i = 0; i < n2; i++) hstate[i] = (float)tmp[i];
327 free(tmp);
328 }
329 return hstate;
330 }
332 /* species 0 on the grid, for the range check */
333 void read_field(double *mn, double *mx, int *finite) {
334 for (int k = 0; k < nsp; k++) synth(dU[k], dspat[k]);
335 sync();
336 const long n = lay.nlines * lay.stride;
337 real *h = (real *)malloc(sizeof(real) * (size_t)n);
338 CU_CHECK(cudaMemcpy(h, dspat[0], sizeof(real) * (size_t)n, cudaMemcpyDeviceToHost));
339 *mn = INFINITY;
340 *mx = -INFINITY;
341 *finite = 1;
342 for (long l = 0; l < lay.nlines; l++)
343 for (long i = 0; i < lay.linelen; i++) {
344 const double x = (double)h[l * lay.stride + i];
345 if (x < *mn) *mn = x;
346 if (x > *mx) *mx = x;
347 if (!isfinite(x)) *finite = 0;
348 }
349 free(h);
350 }
351};
353template <>
354void Run<float>::synth(float *qlm, float *spat) {
355 cu_SH_to_spat_float(sht, (cplx_f *)qlm, spat, sht->lmax);
356}
357template <>
358void Run<float>::analys(float *spat, float *qlm) {
359 cu_spat_to_SH_float(sht, spat, (cplx_f *)qlm, sht->lmax);
360}
361template <>
362void Run<double>::synth(double *qlm, double *spat) {
363 cu_SH_to_spat(sht, (cplx *)qlm, spat, sht->lmax);
364}
365template <>
366void Run<double>::analys(double *spat, double *qlm) {
367 cu_spat_to_SH(sht, spat, (cplx *)qlm, sht->lmax);
368}
370/* ----------------------------------------------------------------------- main */
372template <typename real>
373static int run(shtns_cfg sht, cudaStream_t stream, const shtb_spec &spec, const char *adapter,
374 const char *runtime, const char *cfg_info) {
375 Run<real> r;
376 r.sht = sht;
377 r.spec = &spec;
378 r.stream = stream;
379 r.lay = layout_of(sht, spec.layout);
380 r.nlm = (long)sht->nlm;
381 r.n2 = 2 * r.nlm;
382 r.npts = r.lay.nlat * r.lay.nphi;
383 r.nsp = spec.model->nspecies;
384 r.c = shtb_make_step_const<real>(spec.model, &spec.params);
385 shtb_background(spec.model, &spec.params, r.base);
386 r.alloc();
388 const int transform_mode = spec.mode == SHTB_MODE_TRANSFORM;
389 const char *precision = sizeof(real) == 4 ? "fp32" : "fp64";
391 if (!spec.json) {
392 printf("shtbench_gpu — upstream SHTNS on CUDA, %s only\n\n",
393 transform_mode ? "transforms" : "solver");
394 printf(" mode %s\n", transform_mode
395 ? "transform (one synth + one analys per step)"
396 : "solver (one IMEX Euler timestep per step)");
397 if (!transform_mode) {
398 printf(" preset %s (models/%s.m: %d species)\n", spec.preset->label, spec.model->key,
399 r.nsp);
400 printf(" params ");
401 for (int i = 0; i < spec.model->nparams; i++)
402 printf("%s=%g ", spec.model->params[i].key,
403 *shtb_field_c(&spec.params, spec.model->params[i].off));
404 printf("\n");
405 }
406 printf(" grid lmax %d · %ldx%ld · nlm %ld\n", spec.lmax, r.lay.nlat, r.lay.nphi, r.nlm);
407 printf(" layout %s%s\n",
408 spec.layout == SHTB_LAYOUT_PHI ? "phi-contiguous" : "theta-contiguous (native)",
409 r.lay.padded ? ", padded" : "");
410 printf(" backend %s\n %s, %s\n %s\n", adapter, precision, runtime,
411 cfg_info ? cfg_info : "(no GPU config info)");
412 printf(" run %d warmup + %d timed steps in batches of %d, seed %d\n\n", spec.warmup,
413 spec.steps, spec.batch, spec.seed);
414 }
416 if (transform_mode)
417 r.seed_spectrum();
418 else
419 r.seed_state();
421 for (int i = 0; i < spec.warmup; i++) transform_mode ? r.round_trip() : r.step();
422 r.sync();
424 /* --- throughput: a batch launched together, waited for once ------------- */
425 const int batches = (spec.steps + spec.batch - 1) / spec.batch;
426 double launch_ms = 0;
427 int done = 0;
428 const double t0 = shtb_now_ms();
429 for (int b = 0; b < batches; b++) {
430 const int n = spec.steps - done < spec.batch ? spec.steps - done : spec.batch;
431 const double e0 = shtb_now_ms();
432 for (int i = 0; i < n; i++) transform_mode ? r.round_trip() : r.step();
433 launch_ms += shtb_now_ms() - e0;
434 r.sync();
435 done += n;
436 }
437 const double throughput_ms = (shtb_now_ms() - t0) / done;
439 /* --- latency: one step per synchronization, for the distribution -------- */
440 const int lat_steps = spec.steps < 200 ? spec.steps : 200;
441 double *samples = (double *)malloc(sizeof(double) * (size_t)lat_steps);
442 for (int i = 0; i < lat_steps; i++) {
443 const double a = shtb_now_ms();
444 transform_mode ? r.round_trip() : r.step();
445 r.sync();
446 samples[i] = shtb_now_ms() - a;
447 }
449 shtb_report rep;
450 memset(&rep, 0, sizeof(rep));
451 char libbuf[192], adapterbuf[160], runtimebuf[160];
452 rep.library = shtb_json_safe(libbuf, sizeof(libbuf), shtns_get_build_info());
453 rep.runtime = shtb_json_safe(runtimebuf, sizeof(runtimebuf), runtime);
454 rep.adapter = shtb_json_safe(adapterbuf, sizeof(adapterbuf), adapter);
455 rep.precision = precision;
456 rep.fourier = "cufft/vkfft";
457 rep.nlm = r.nlm;
458 rep.ops_per_step = transform_mode ? 2 : 2 * r.nsp + 1 + r.nsp;
459 rep.ms_per_step = throughput_ms;
460 rep.encode_ms_per_step = launch_ms / done;
461 rep.latency = shtb_stats(samples, lat_steps);
462 rep.have_latency = 1;
463 rep.steps_run = spec.warmup + spec.steps + lat_steps;
464 rep.model_t = transform_mode ? 0 : rep.steps_run * spec.params.dt;
466 if (transform_mode) {
467 const float *s = r.read_state();
468 rep.field_min = INFINITY;
469 rep.field_max = -INFINITY;
470 rep.finite = shtb_all_finite(s, r.n2);
471 for (long i = 0; i < r.n2; i++) {
472 if (s[i] < rep.field_min) rep.field_min = s[i];
473 if (s[i] > rep.field_max) rep.field_max = s[i];
474 }
475 } else {
476 r.read_field(&rep.field_min, &rep.field_max, &rep.finite);
477 }
479 /* A reproducible state to compare against a WebGPU run: exactly --steps steps
480 * from the seed, separate from the timed runs above. */
481 if (spec.digest) {
482 if (transform_mode) {
483 r.seed_spectrum();
484 rep.input_digest = shtb_digest_of(r.hstate, r.n2);
485 rep.have_input_digest = 1;
486 for (int i = 0; i < spec.steps; i++) r.round_trip();
487 } else {
488 r.seed_state();
489 for (int i = 0; i < spec.steps; i++) r.step();
490 }
491 rep.digest = shtb_digest_of(r.read_state(), r.n2);
492 rep.have_digest = 1;
493 }
495 if (spec.json) {
496 shtb_print_json(&spec, &rep);
497 } else {
498 printf(" %.3f ms/step %.1f steps/s", rep.ms_per_step, 1000.0 / rep.ms_per_step);
499 if (!transform_mode) printf(" %.2f model time/s", spec.params.dt * 1000.0 / rep.ms_per_step);
500 printf(" (batches of %d)\n", spec.batch);
501 printf(" of which CPU kernel launching: %.3f ms/step (%.0f%% — the rest is the GPU)\n",
502 rep.encode_ms_per_step, 100.0 * rep.encode_ms_per_step / rep.ms_per_step);
503 printf(" one step per sync: %.3f ms mean · median %.3f · p05 %.3f · p95 %.3f · min %.3f\n",
504 rep.latency.mean_ms, rep.latency.median_ms, rep.latency.p05_ms, rep.latency.p95_ms,
505 rep.latency.min_ms);
506 if (transform_mode)
507 printf(" i.e. %.3f ms per single transform\n", rep.ms_per_step / 2);
508 else
509 printf(" after %d steps: t = %.2f, field ∈ [%.4f, %.4f] (contrast %.4f)%s\n",
510 rep.steps_run, rep.model_t, rep.field_min, rep.field_max,
511 rep.field_max - rep.field_min, rep.finite ? "" : " — NOT FINITE");
512 if (rep.have_digest) {
513 printf("\n state after %d steps from seed %d:\n", spec.steps, spec.seed);
514 printf(" n=%ld min=%.9g max=%.9g mean=%.9g rms=%.9g\n", rep.digest.n, rep.digest.min,
515 rep.digest.max, rep.digest.mean, rep.digest.rms);
516 }
517 printf("\n Compare with `npm run bench --json` on this machine: same GPU, same\n"
518 " precision, same grid — the difference is the transform implementation.\n"
519 " scripts/compare-native.mjs runs both and lines the numbers up.\n");
520 }
522 if (spec.dump_state && rep.have_digest) {
523 if (shtb_dump_state(spec.dump_state, &spec, &rep, r.hstate, r.n2) != 0) {
524 fprintf(stderr, "shtbench_gpu: cannot write %s\n", spec.dump_state);
525 return 1;
526 }
527 if (!spec.json) printf("\n wrote %s\n", spec.dump_state);
528 }
530 free(samples);
531 r.free_all();
532 return rep.finite ? 0 : 1;
533}
535int main(int argc, char **argv) {
536 shtb_spec spec;
537 double polar_eps = 0.0;
538 int device = 0;
540 /* --polar-eps and --device are ours, not part of the shared spec. */
541 int argc2 = 0;
542 char **argv2 = (char **)malloc(sizeof(char *) * (size_t)argc);
543 for (int i = 0; i < argc; i++) {
544 if (strcmp(argv[i], "--polar-eps") == 0 && i + 1 < argc) {
545 polar_eps = atof(argv[++i]);
546 continue;
547 }
548 if (strncmp(argv[i], "--polar-eps=", 12) == 0) {
549 polar_eps = atof(argv[i] + 12);
550 continue;
551 }
552 if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) {
553 device = atoi(argv[++i]);
554 continue;
555 }
556 if (strncmp(argv[i], "--device=", 9) == 0) {
557 device = atoi(argv[i] + 9);
558 continue;
559 }
560 argv2[argc2++] = argv[i];
561 }
562 int rc = shtb_parse_spec(argc2, argv2, &spec, USAGE);
563 free(argv2);
564 if (rc) return rc == 1 ? 0 : rc;
566 CU_CHECK(cudaSetDevice(device));
567 cudaDeviceProp prop;
568 CU_CHECK(cudaGetDeviceProperties(&prop, device));
569 char adapter[128];
570 snprintf(adapter, sizeof(adapter), "%s (sm_%d%d, %d SMs)", prop.name, prop.major, prop.minor,
571 prop.multiProcessorCount);
572 int rtv = 0, drv = 0;
573 cudaRuntimeGetVersion(&rtv);
574 cudaDriverGetVersion(&drv);
575 char runtime[128];
576 snprintf(runtime, sizeof(runtime), "CUDA runtime %d.%d, driver %d.%d", rtv / 1000,
577 (rtv % 1000) / 10, drv / 1000, (drv % 1000) / 10);
579 shtns_verbose(0);
580 shtns_cfg sht = shtns_create(spec.lmax, spec.lmax, 1, sht_orthonormal);
581 if (!sht) {
582 fprintf(stderr, "shtbench_gpu: shtns_create failed\n");
583 return 1;
584 }
585 /* Our reaction and update kernels have to be ordered against SHTNS'
586 * transforms, so both go on one stream we own. cushtns_set_streams must come
587 * before shtns_set_grid, which is where the GPU (and its FFT plan) is set up. */
588 cudaStream_t stream = 0;
589 CU_CHECK(cudaStreamCreate(&stream));
590 cushtns_set_streams(sht, stream, 0);
592 int flags = sht_gauss | SHT_ALLOW_GPU | SHT_SCALAR_ONLY;
593 flags |= spec.layout == SHTB_LAYOUT_PHI ? SHT_PHI_CONTIGUOUS : SHT_THETA_CONTIGUOUS;
594 if (spec.fp32) flags |= SHT_FP32;
595 if (shtns_set_grid(sht, (enum shtns_type)flags, polar_eps, spec.nlat, spec.nphi) <= 0) {
596 fprintf(stderr, "shtbench_gpu: shtns_set_grid failed for lmax %d on a %dx%d grid\n", spec.lmax,
597 spec.nlat, spec.nphi);
598 return 1;
599 }
600 if ((int)sht->nlat != spec.nlat || (int)sht->nphi != spec.nphi) {
601 fprintf(stderr, "shtbench_gpu: SHTNS chose a %ux%u grid, not the %dx%d asked for\n", sht->nlat,
602 sht->nphi, spec.nlat, spec.nphi);
603 return 1;
604 }
606 /* cushtns_get_cfg_info() returns NULL when the GPU was never initialized, which
607 * is how SHT_ALLOW_GPU failing shows up — and with SHT_FP32 the CPU fallback
608 * would read fp64 out of fp32 buffers, so stop rather than produce a number
609 * for the wrong thing. */
610 const char *cfg_info = cushtns_get_cfg_info(sht);
611 if (!cfg_info) {
612 fprintf(stderr,
613 "shtbench_gpu: SHTNS did not initialize the GPU for this grid (lmax %d, %dx%d).\n"
614 " Was it built with --enable-cuda, and does nlat %% 4 == 0 hold for the\n"
615 " theta-contiguous layout? Try --layout phi.\n",
616 spec.lmax, spec.nlat, spec.nphi);
617 return 1;
618 }
620 const int status = spec.fp32 ? run<float>(sht, stream, spec, adapter, runtime, cfg_info)
621 : run<double>(sht, stream, spec, adapter, runtime, cfg_info);
622 shtns_destroy(sht);
623 cudaStreamDestroy(stream);
624 return status;
625}