Enable EverestCrystal on the CUDA (Cupy) GPU context#202
Enable EverestCrystal on the CUDA (Cupy) GPU context#202InnovateInfinity87 wants to merge 10 commits into
Conversation
…alue, stack scratch arrays, qsort CPU-guard + GPU insertion sort, kill-particle error idiom. GPU-ready; CPU numerics unchanged.
…alloc->out-params, frees dropped. CPU numerics unchanged.
…fill-via-pointer, EverestCrystal_free removed, missing /*gpufun*/ tags added. CPU numerics unchanged.
… xobjects /*restrict*/ token across the EverestCrystal compile path, so the specializer drops it on CUDA. CPU numerics unchanged (bitwise).
…ext; M_PI guard + /*gpukern*/ set_material tag (clean Cupy compile); get_s.h +1 spare slot (calculate_overlap_array_interval 1-element overrun, GPU-trapped); test fixture create_crystal(NULL,...). CPU numerics bitwise-unchanged; 24/24 crystal+geometry tests pass.
… (per-thread frame exceeds the 1024B default) — fixes the GPU-track cudaErrorIllegalAddress. xcoll EverestCrystal element-level helper: new module xcoll/scattering_routines/everest/gpu_stack.py exposing set_crystal_stack_limit(context, nbytes=16384), exported from that package's __init__.py, and called once at the end of EverestCrystal.__init__ (in xcoll/beam_elements/everest.py) with self._context. Cupy-only, idempotent, never-lowers; no import side-effect.. CPU numerics unchanged; GPU track matches CPU golden.
…enchmark - examples/everest_crystal_cpu_vs_gpu.py: also generate the T13 recursion-stress parity figure (10x crystal + grazing pencil); CPU/GPU both 19796 survivors, KS 5e-5. - examples/everest_crystal_gpu_benchmark.py (new): warm, device-synced CPU-vs-GPU timing of the crystal line across bunch sizes -> speed-up + throughput figure. - docs/gpu_everest_crystal.md: new 'Recursion-stress agreement' and 'Performance' sections. RTX 2070: break-even ~2-3k particles, plateau ~5-6x (5.9x @500k), GPU ~9.6 Mp/s vs ~1.6 Mp/s/core.
… change) - materials/material.py: remove unused public get_materials_context / set_materials_context (never called, not exported); keep _materials_context + _material_in_context. - everest/gpu_stack.py: drop the redundant _requested_stack_limit high-water global (the live deviceGetLimit already encodes it -> idempotent/never-lower) and the defensive _is_cupy_context wrapper; use the xcoll-standard isinstance(context, xo.ContextCupy). - examples: untrack generated everest_crystal_cpu_vs_gpu_summary.json (run output). GPU tests still pass (distribution match + stack-limit), CPU numerics unchanged.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 385750d116
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| int8_t n_hit = 0; | ||
| double* s = (double*) malloc(XC_MAX_CROSS_PER_SEGMENT*n_segments*sizeof(double)); | ||
| // +1 spare slot (see contract in methods.h:43-44 / first crossing above). | ||
| double s[XC_MAX_CROSS_PER_SEGMENT*XC_MAX_SEGMENTS + 1]; |
There was a problem hiding this comment.
When callers pass more than four segments, this fixed 9-slot stack buffer is no longer large enough for the existing CPU-only polygon helpers: create_open_polygon is still allowed to build num_polys + 1 segments, and the geometry data generator exercises 8-point open polygons, so find_crossing_with_vlimit can produce 9 hits and then calculate_overlap_array_interval writes the documented spare element at arr[*length]. That turns those valid geometry calls into a stack overwrite; the old malloc(2*n_segments) sizing avoided this.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d793425. The fixed-size buffer was only correct for the bounded element callers (collimator passes 3 segments, crystal 4 = XC_MAX_SEGMENTS); the geometry-test open polygons exceed that, as flagged. Restored n_segments-based sizing via malloc under #ifdef XO_CONTEXT_CPU (plus the +1 spare slot from the methods.h contract, and the contract comment is updated), and kept the bounded stack buffer on the GPU path where only EverestCrystal (n_segments <= XC_MAX_SEGMENTS) tracks. tests/test_geometry.py now passes 16/16 including the open_polygon cases; the crystal CPU-vs-GPU acceptance and CPU bitwise parity (maxdiff=0) are unchanged.
…ENTS polygon overflow) The malloc->stack GPU port hard-capped the get_s.h crossing buffers at XC_MAX_CROSS_PER_SEGMENT*XC_MAX_SEGMENTS+1 slots. Element callers stay within that bound (collimator passes 3 segments, crystal 4 == XC_MAX_SEGMENTS), but the geometry test generator builds open polygons with more sides (e.g. an 8-point open polygon -> 9 segments), and find_crossing / the calculate_overlap_array_interval spare write then run past the fixed buffer. Restore n_segments-based sizing via malloc on CPU (as upstream did, plus the +1 spare slot), and keep the bounded stack buffer on GPU, where only EverestCrystal (n_segments <= XC_MAX_SEGMENTS) tracks. CPU numerics stay bitwise-unchanged and the GPU crystal CPU-vs-GPU acceptance is unchanged; tests/test_geometry.py passes 16/16 (incl. the open-polygon cases) and the everest_crystal GPU tests still pass.
1 participant
Summary
EverestCrystal(and the Everest scattering + geometry routines it pulls in) nowcompiles and tracks on the CUDA context (
xobjects.ContextCupy), in addition tothe CPU context. This builds on #790 (which made
xtrack'sRandomRutherfordGPU-ready, the last CPU-only piece of the Everest physics).
CPU numerics are bitwise unchanged; CPU-vs-GPU agreement is validated
distributionally (KS / KL), the way the Rutherford RNG PR is.
Scope: this enables
EverestCrystalon the GPU.EverestBlockandEverestCollimatorare untouched and remain CPU-only.Background
Historically every Everest test excluded the GPU contexts with the note
# Rutherford RNG not on GPU: the Everest scattering samples the Rutherforddistribution through
RandomRutherford, which was CPU-only. #790 lifted that forthe RNG itself; this PR makes the crystal element and its kernels GPU-compilable.
What changed
Grouped as in the commit series:
scattering_routines/geometry/): the per-segmentfunction-pointer vtable becomes an
int8enum +switch; segments are storedby value instead of on the heap; the crossing/overlap scratch buffers are
fixed-size stack arrays;
qsortis kept on CPU under#ifdef XO_CONTEXT_CPUwith a GPU insertion-sort fallback; the kill-particle error idiom replaces the
CPU-only
exit.scattering_routines/everest/): the twopair-returning
mallochelpers (channel_transport,scamcs) return throughout-params, and the per-track scratch structs (
EverestData,CrystalGeometry,EverestCollData) are filled in place on the stack instead ofmalloc/free.restrictqualifier becomes the xobjects/*restrict*/token across the crystal compile path (so the specializer dropsit on CUDA);
M_PIis guarded for NVRTC;EverestCrystal_set_materialistagged
/*gpukern*/.Materialis landed in the consuming element's contextbefore track time (
_material_in_context), soMaterialDatalives where thekernel reads it instead of staying pinned to the host buffer. No-op on CPU.
scattering_routines/everest/gpu_stack.py):the crystal kernel's per-track structs exceed CUDA's default 1024-byte
per-thread stack frame, so a GPU track traps with
cudaErrorIllegalAddress.A small Cupy-only helper raises
cudaLimitStackSize(to a 16 KiB default) onceper process; it is idempotent, never lowers an existing limit, and is a no-op
on CPU. It is called from
EverestCrystal.__init__with the element's context,so user code needs no extra step.
CPU numerics unchanged (bitwise)
Every change above is either GPU-only or behaviour-preserving on CPU. Verified by
tracking a fixed-seed canonical bunch through
Drift + EverestCrystal + DriftonContextCpuon this branch vs cleanmain: the finalx, px, y, py, zeta, delta, statearrays are bit-identical (maxdiff = 0).Validation
tests/test_everest_crystal_gpu.py:test_everest_crystal_cpu_gpu_distribution_match— tracks the same fixed-seedbunch on CPU and GPU and asserts per-coordinate
KS < 0.02pluschanneling/surviving-fraction agreement. Skips cleanly (
importorskip("cupy"))where no CUDA GPU is available.
test_crystal_stack_limit_is_raised_on_gpu— asserts the per-thread stack israised above the 1024-byte default to the documented ≥ 16 KiB, is a no-op on
CPU, and is idempotent / never-lowering.
strip crystal): max KS
5.0e-05(≈ the1/Nfloor, three orders below the0.02gate), KL0, channeling fraction identical (0.16413), and19 911 / 20 000survivors on both contexts. The deepest-recursion stresscase (10×-longer crystal, grazing near-pencil beam) also agrees.
docs/gpu_everest_crystal.mddocuments the acceptance method, the stack-limitrequirement, and the overlaid CPU/GPU plots.
CPU vs GPU on the canonical
Drift + EverestCrystal + Driftline — per-coordinatedistributions (top: PDF overlay with KS/KL; middle: empirical CDFs with the KS gap;
bottom: GPU/CPU bin-count ratio against the ±1σ counting-noise band):
Channeling efficiency (channeled / surviving) vs incoming angle on CPU and GPU, with
the GPU − CPU residual on the right — the curves overlay across the acceptance:
Deepest-recursion stress case (10×-longer crystal tracked by a grazing near-pencil
beam) — the hardest case for the per-thread CUDA stack — also agrees per coordinate:
Performance
This is the first time
EverestCrystalcan run on the GPU at all.examples/everest_crystal_gpu_benchmark.pytimes one turn through the canonicalline on
ContextCpu(single core) vsContextCupy(warm kernels, devicesynchronized around the timed region). On one NVIDIA RTX 2070 vs a single CPU
core:
Break-even is ~2–3 k particles; the speed-up plateaus around 5–6× from ~50 k
(GPU ≈ 9.6 Mp/s vs ≈ 1.6 Mp/s on one CPU core). The crystal kernel is branchy and
stochastic (channeling vs amorphous paths, RNG rejection, sub-segment recursion),
so threads diverge and the plateau sits below a non-divergent element; the CPU
baseline is a single core. Absolute numbers are hardware-specific.
Notes
RandomRutherfordGPU change (#790), now inxtrackmain.docs/andexamples/.