GPU: cross-vendor CubeCL renderer + fused CE transport kernel - #13
Merged
Merged
Conversation
gpu: cross-vendor CubeCL geometry preview renderer (f64) Adds a GPU ray-cast preview of the recursive CSG geometry written once in CubeCL #[cube] kernels, compiling to CUDA / HIP-ROCm / Vulkan / Metal / WebGPU. Runs in f64 (confirmed on the wgpu/Vulkan backend when the adapter exposes SHADER_F64) so it matches the CPU/CUDA geometry walk; the broken CPU ray-cast normals (lattice-plane collapse + too-coarse occupancy gradient) are bypassed entirely. - cubecl 0.10 added as a general (non-gated) workspace dependency, wgpu runtime always on. First step toward replacing the hand-written gpu/cuda/*.cu kernels + cudarc with one cross-vendor source. - geometry::flat: lifted the Geometry->SoA flattening (build_host_tables / pack_surface / flatten_region) out of the cuda-gated gpu_recursive module so the CUDA upload path and the new CubeCL renderer share one device-layout definition. No behaviour change; 8 recursive-geometry tests + cuda build still green. - gpu_render: ports gr_surf_eval / gr_find_cell / gr_trace_step / gr_surf_normal / raycast_preview to f64 #[cube] kernels. Single-exit rewrite (CubeCL forbids early return), depth-4 SoA local-array stack, on-demand surface eval (no per-thread scratch), 2-blob upload to stay within storage-binding limits. Rect lattices + nested universes walked fully; hex deferred. - preview_scene --cubecl-out <png>: headless cross-vendor GPU render, honouring --cam-azim/--cam-elev/--zoom/--resolution. - Tracked hero images under docs/images (Godiva + PWR assembly GPU renders, 2D scene cross-sections). Tests: 441/441 lib (new wgpu sphere-render smoke test included); cuda build compiles. @
The binary opaque/skip transparency made water vanish entirely, so every ray stacked all the fuel pins behind it into a flat striped wall. Replace it with true alpha compositing in the ray-march loop: moderating/coolant fluids (water, D2O, solution, borated coolant) carry a per-cm absorption (Beer-Lambert) plus a dimmed tint, accumulated front-to-back as the ray crosses each segment; air/void stay perfectly clear; solids composite under the accumulated haze and terminate the ray. Because a ray-marcher visits surfaces strictly front-to-back, this is exact order-independent transparency with no depth buffer, depth peeling, or alpha-to-coverage -- those are rasterizer workarounds for out-of-order primitives, which a per-ray march does not have. - pack_scene gains an absorb (alpha-per-cm) channel, packed into the f64 blob; preview_scene derives it from material names. - opaque_mask now treats water / H2O / D2O / moderator / coolant / solution / borated as transparent (reactor-geometry preview convention) so fuel structure reads through the fluid. - Re-rendered hero images: LCT-008 now shows a 3D rod forest; the PWR 17x17 assembly shows pin columns + guide-tube gaps with depth. Tests 441/441; cuda build green.
Reframes the project from "SVD compression study" to the continuous-energy MC engine it became, leading with the GPU scene-preview hero images and benchmark graphs. - Origin story: how answering one SVD question grew a real MC code. - Hero renders (CubeCL cross-vendor GPU): LCT-008 rod forest, PWR 17x17 assembly, Godiva sphere; plus 2-D cross-sections. - Benchmark section with graphs: four-way k_inf vs OpenMC, ICSBEP k_calc vs handbook, SVD singular spectrum, throughput (Godiva + PWR), and the memory comparison. - Honest SVD verdict, now including the memory case the single-temperature bar chart hides: one rank-k basis reconstructs ANY temperature (near-flat in N_T vs a table's linear growth) and maps far better into GPU/cache memory than irregular pointwise tables -- the project's original cache-resident thesis. Backed by paper/sections/memory_vs_precision.tex (off-library regime). - Documents both GPU backends (CUDA transport + CubeCL renderer) and flags the transport-kernel CubeCL port as in progress. Benchmark plots copied from the gitignored outputs/ into tracked docs/images/.
First transport kernel ported to CubeCL: constant cross-sections per material, one batch of histories, fission sites banked atomically. Proves the cross-vendor transport plumbing in f64 — geometry walk (find_cell / trace_step), PCG-XSH-RR RNG in u64, isotropic scatter, and u32 atomics (portable fission-bank cursor + counters; i64/f64 atomics need backend extensions wgpu doesn't guarantee). The host pieces are validated in isolation on NVIDIA/Vulkan: RNG, atomics, find_cell, helper nesting, and 4 KB of private arrays all run correctly. The full kernel compiles to valid WGSL but faults at dispatch (STATUS_ACCESS_VIOLATION) — a known upstream SPIR-V bug with large thread-private Array state (tracel-ai/cubecl#1336). Shrinking the state isn't workable for recursive geometry, so transport stays on the CUDA backend until the fix lands. Kernel + host harness kept in-tree and ready to re-enable; the batch test is #[ignore]d. Bug confirmed upstream with our NVIDIA repro (docs/cubecl_1336_comment.md). Tests 441 passed / 1 ignored; cuda build green.
The const-XS transport #[cube] kernel works through CubeCL's CUDA runtime — same source that faults on wgpu/Vulkan. Enables cubecl/cuda under our `cuda` feature and adds cuda_const_xs_batch_runs: 4096 histories, absorbed+leaked == 4096 (each ends once), ~7.3k fission sites all inside the sphere. So the kernel is functional on NVIDIA now; only the cross-vendor Vulkan/Metal path waits on the #1336 SPIR-V private-Array fix. Updated the module status accordingly; the Vulkan test stays #[ignore]d, the CUDA test runs. Also fixes test-only imports in gpu_recursive.rs (Vec3/CellFill) that broke `cargo test --features cuda` after the geometry::flat extraction. Tests: 441 default / 456 cuda passed, 1 ignored (Vulkan transport).
A/B against the legacy .cu const-XS kernel on the 2x2 reflective lattice exposed a real bug: the CubeCL kernel re-ran find_cell after *every* event, including reflections. On a reflective surface the particle sits exactly on the boundary, so the re-find is numerically unstable and reported depth==0 -> spurious leak. Result was 190731 / 200000 false leaks and k=0.074 vs the correct 1.6. Fix: only re-resolve the cell after a transmission crossing (matches the .cu reference); reflection / collision / scatter keep the current stack. After the fix both backends agree: legacy .cu k=1.59995 leak=0 CubeCL k=1.59889 leak=0 |Δk| = 106 pcm (1σ_MC ≈ 283 pcm) PASS New bin `cubecl_vs_cu_keff` (--features cuda) runs the three-way comparison (CPU ref / legacy .cu / CubeCL-CUDA) on identical scene, source, and seeds. Tests 456 cuda passed, 1 ignored.
Foundation for the CE transport port (target: HEU-COMP-INTER, multi- nuclide fast/intermediate). gpu_ce_cubecl ports the genuinely new CE piece: per-nuclide pointwise sigma on the device, one binary search on the shared energy grid + log-log interpolation per reaction, exactly mirroring SvdXsProvider::lookup. Carries 5 channels (elastic, fission, capture, inelastic, n2n) — the ones that matter for a fast/intermediate spectrum. A/B (tests/cubecl_ce_xs.rs, --features cuda) reconstructs total micro sigma(E) over all 19 HEU-COMP-INTER-003 nuclides at a geometric energy sweep, CPU provider vs CubeCL CUDA runtime: 1007 (nuclide, E) points checked, worst rel err = 8.07e-13 i.e. bit-exact to f64 round-off. The grid search + interp are correct; the transport loop on top reuses the const-XS structure already A/B-validated against the legacy .cu. Also fixes a pre-existing compile error in metal_stats_diag (the `openmc=<path>` arg was referenced but never declared/parsed; predates this branch, surfaced when building cuda bins). Tests: default lib green; CE XS A/B passes on CUDA.
Extends gpu_ce_cubecl toward an event-based CE transport kernel: pack_ce_scene packs the per-nuclide grids+sigma PLUS a material table (per material: list of (nuclide_idx, atom_density)) and the per-nuclide awr / nu_bar arrays, so the device can sum macroscopic Sigma_t(E) over a material's nuclides. pack_ce (the XS-lookup A/B) now delegates with an empty material table -- that test is unchanged. Event-based design (mirrors gpu/cuda/transport_event_based.cu): keeps per-particle state in global SoA Arrays with small per-kernel private state, which should avoid the thread-private-Array fault cubecl#1336 -- potentially unblocking Vulkan as well as CUDA. Kernel stages (trace + reaction events + k tally) land next. Foundation compiles; CE XS A/B still passes.
Second CE physics piece ported by wiring the CPU data directly: extract_angular pulls each nuclide's NuclideKernels.elastic_angle (energies + per-energy TabularMuDist mu/cdf/pdf/histogram) into flat device blobs; pack_angular lays them out; the #[cube] sample_mu_at + sample_mu_bin mirror AngularDistribution::sample_mu (energy bracket + stochastic bin pick) and TabularMuDist::sample_with_xi (histogram linear CDF or lin-lin quadratic CDF inversion) exactly. A/B (tests/cubecl_ce_angular.rs, --features cuda): for every HEU-COMP-INTER-003 nuclide with elastic angular data, 200k samples at 7 fast/intermediate energies, CPU vs CubeCL-CUDA. Mean mu and variance agree within 5 sigma MC noise: 133 (nuclide, E) means checked, worst |Delta mean mu| = 5.9e-3 This is the forward-peaked-elastic structure that historically biased ieu-met-fast-001 by hundreds of pcm when sampled wrong. CE physics ported so far: sigma lookup (bit-exact) + elastic angular (within MC noise). Remaining: fission chi sampler + free-gas target, then the fused trace_and_sample event kernel + k_eff A/B.
Third CE physics piece, same wiring pattern: extract_fission pulls each nuclide's NuclideKernels.fission_energy_dist into flat device blobs, tagged by law -- tabular ContinuousTabular (per-incident-E e_out/pdf/cdf slices) or closed-form Watt / Maxwell / Evaporation (a/b or theta lin-lin tables). The #[cube] fis_sample_energy mirrors EnergyDistribution::sample exactly: tabular path does the stochastic energy-bin pick + quadratic CDF inversion (TabularEnergyDist::sample_with_xi) + OpenMC scaled kinematic remap to the interpolated [E1,EK]; closed-form path runs the Cranberg/Coveyou rejection samplers (WattLaw::sample / sample_maxwell / sample_evaporation). A/B (tests/cubecl_ce_fission.rs, --features cuda): 200k samples at 6 incident energies for every fissile HEU-COMP-INTER-003 nuclide, CPU vs CubeCL-CUDA. Mean outgoing energy + variance agree within 5 sigma MC noise: 24 (nuclide, E) means checked, worst rel Delta-mean = 3.3e-3 CE physics ported + validated: sigma lookup (bit-exact), elastic angular (MC noise), fission chi (MC noise). Remaining: free-gas target velocity (self-contained), then the fused trace_and_sample event kernel + k_eff A/B vs .cu.
Assembles the validated CE primitives into one fused single-generation transport kernel: unified scene pack (geometry SoA + per-nuclide CE sigma + elastic angular + fission chi + material table, one i32/f64 blob with a base-offset header), and ce_transport_kernel — one source neutron per thread, looping: find_cell -> sum macroscopic Sigma_t over the cell material's nuclides via the CE lookup -> sample distance vs surface -> pick nuclide proportional to Sigma_t,nuc -> analog reaction pick (elastic/inelastic/n2n/fission/capture) -> elastic via free-gas target velocity + two-body CM with the tabulated angular sampler -> fission banks nu-bar neutrons at chi-sampled energies -> capture/leak kill. Host ce_generation drives it and returns the fission bank + counters (single-generation k = fissions / source). Reuses the three A/B-validated samplers (CE sigma lookup, elastic angular, fission chi) via rebased variants that read each sub-blob at its base offset in the unified pack. CUDA-only until cubecl#1336. Compiles clean (default + the kernel). k_eff A/B vs the CPU/.cu lands next.
First whole-transport run of the fused CubeCL CE kernel on a real fissile sphere (U-235 + U-238 from ENDF, r=8.7 cm, vacuum). Output is physically correct: coll=55535 leak=9939 cap=466 fissions=24006 k(1 generation)=1.200 mean fission E=2.045 MeV k~1.2 is right for a bare HEU sphere's single-generation multiplication; 2.0 MeV is the textbook Watt/chi peak; histories balance (fission + leak + capture account for the population). Genuine continuous-energy MC -- real cross-sections, energy-dependent lookup, free-gas elastic, tabulated angular scatter, fission chi -- in one CubeCL kernel on the GPU (CUDA runtime; Vulkan still blocked by cubecl#1336). tests/cubecl_ce_transport.rs gates collisions>0, fissions banked, k in (0.3,3), mean fission E in (0.3,5) MeV. Also fixes a pre-existing stale CudaRunner initializer in cuda_runs.rs (missing nxn_mode field added in 4a4f1eb) that broke cargo test --features cuda.
…x guards The find_cell/trace_step/surf_eval/surf_dist/cell_contains/cross_or_die walk was hand-copied ~1:1 across gpu_transport_cubecl.rs, gpu_ce_cubecl.rs, and (partially) gpu_render.rs. That duplication had already produced a real drift (surf_dist's PLANE_GENERAL tolerance was 1e-300 in two copies and 1e-30 in the third) and hid a real bug: find_cell treated CellFill::Void as an instant leak instead of a cell to free-stream through (both kernels' "void: free-stream to next surface" handling downstream was already correct and never reached). Pull the walk into one shared module, gpu_cubecl_geom.rs, imported by all three kernel files; fix the Void leak in the single copy. Tag constants (SURF_*/REGION_*/FILL_*/BC_*) now come from geometry::flat's existing single source of truth instead of being redeclared three times. Also add geometry::flat::check_gpu_supported / check_gpu_rotation_supported, mirroring the loud-error guard GpuRecursiveContext::build already has for per-cell Mat3 rotation (the CubeCL walk is axis-aligned only) plus a hex-lattice check for the transport paths (the CubeCL walk has no hex descent, unlike the legacy CUDA path) — wired into pack_transport / pack_ce_full / pack_scene so an unsupported scene now errors instead of silently mis-transporting or mis-rendering. 454/454 lib tests green with --features cuda (RTX A1000); the 4 cubecl_ce_* GPU integration tests and cubecl_vs_cu_keff (legacy .cu vs CubeCL A/B, |Δk| = 106 pcm, within MC noise) all still pass.
…lation Two numerical-safety gaps in the fused CE transport kernel's fission and cross-section sampling, found during review: - fis_sample_energy / fis_sample_energy_base (Watt, Maxwell, Evaporation laws): on rejection-sampling exhaustion (128 tries) the fallback returned a literal sentinel of 2 eV instead of the last attempted draw clamped into [1e-5, max_e]. The CPU reference (WattLaw::sample, MaxwellLaw::sample_maxwell/sample_evaporation) already does the clamp-last-draw safety valve; the GPU copy now matches it. Only affects the rare case where a (mis-)configured (a, b, u) tuple collapses the acceptance band — normal-case sampling is unaffected (cubecl_ce_fission / cubecl_ce_transport reproduce identical statistics before and after). - rx_interp / ce_macro_rx: the log-log interpolation fraction wasn't clamped to [0, 1], so an energy above/below a nuclide's tabulated grid extrapolated the power law instead of holding the boundary cross-section. xs_provider.rs's CPU path already clamps (`log_frac.clamp(0.0, 1.0)`); the GPU copies now do the same. 454/454 lib tests green with --features cuda; cubecl_ce_xs, cubecl_ce_fission, and cubecl_ce_transport all still pass with unchanged statistics (RTX A1000).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a cross-vendor GPU scene renderer and a fused continuous-energy (CE) neutron transport kernel, both written once in CubeCL and compiled to CUDA (validated on an RTX A1000 here; also targets ROCm/Vulkan/Metal/WebGPU) as an alternative to the hand-maintained
gpu/cuda/*.cukernels.gpu_render.rs— interactive/headless 3-D ray-cast preview of the engine's own recursive CSG geometry, with front-to-back alpha compositing for translucent fluids.gpu_transport_cubecl.rs— const-XS transport kernel proving the CubeCL plumbing (geometry walk + PCG RNG + collision sampling + atomic fission banking).gpu_ce_cubecl.rs— CE cross-section lookup, elastic/fission-χ angular sampling, and the fused single-generation CE transport kernel, each validated with an A/B statistical/bit-exact test against the CPU reference.geometry/flat.rs— geometry SoA flattening lifted out of thecuda-gatedgpu_recursivemodule so it's the single source of truth shared by the CUDA and CubeCL backends.The last two commits are from a review pass on this branch before merging:
gpu: dedupe CubeCL geometry walk; fix void-leak + missing rotation/hex guards— the geometry walk (find_cell/trace_step/surf_eval/surf_dist/cell_contains/cross_or_die) had been hand-copied across all three CubeCL files. That duplication had already caused a real numeric drift (a tolerance was1e-300in two copies and1e-30in the third) and hid a real bug:find_celltreatedCellFill::Voidas an instant leak instead of a cell to free-stream through. Pulled into one sharedgpu_cubecl_geom.rsmodule and fixed in the one place; also added the loud-error guard for per-cell rotation and hex lattices that the legacy CUDA path already has (GpuRecursiveContext::build), now viageometry::flat::check_gpu_supported.gpu: fix Watt/Maxwell/Evaporation exhaustion fallback + CE XS extrapolation— two numerical-safety gaps in the CE kernel found in the same pass: the fission-energy rejection sampler fell back to a literal 2 eV sentinel on exhaustion instead of clamping the last draw (mirroring the CPU'sWattLaw::samplesafety valve), and the CE cross-section log-log interpolation didn't clamp the interpolation fraction to[0, 1], so energies outside a nuclide's tabulated grid would extrapolate instead of holding the boundary value.Test plan
All run on this dev box's RTX A1000 (sm_86), CUDA 13.2 driver:
cargo build --release(default features) — cleancargo build --release --features cuda— cleancargo test --release --lib— 439/439cargo test --release --lib --features cuda— 454/454cargo test --release --features cuda --test cubecl_ce_xs— CE XS A/B, worst rel err 8.07e-13cargo test --release --features cuda --test cubecl_ce_angular— angular A/B, worst |Δmean μ| 5.93e-3cargo test --release --features cuda --test cubecl_ce_fission— fission χ A/B, worst rel Δmean 3.271e-3cargo test --release --features cuda --test cubecl_ce_transport— fused CE kernel runs end-to-end, sane k/fission-bank statscargo run --release --features cuda --bin cubecl_vs_cu_keff— legacy.cuvs CubeCL const-XS k_eff, |Δk| = 106 pcm (within 1σ_MC ≈ 283 pcm)