From b51e11f276d602b3c137e159bda0f3638b3d0353 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:13:48 +0200 Subject: [PATCH 1/3] Hand-written CUDA vs code generation: port bstar_parallel_3form both ways (decision for #690) - bstar_parallel_3form_cuda.cu: hand-written CUDA version, 1:1 with pyccel, on the existing device helpers; new device helper eval_0form_spline_mpi in evaluation_kernels_3d.cuh. - 14 parity cases (every mapping, alpha = 1/0/mixed, wrap-around in mod, hole) in cuda_parity_cases.py; passes the CPU emulation. - pic/tests/codegen_spike: prototype generator translating the pyccel source (kernel and all helpers) into numba.cuda; agrees with pyccel in the numba.cuda simulator and under njit. Test-only, skipped without numba. - CUDA_STRATEGY.md: decision section (hand-written stays), comparison and what would change it. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 91 +++- .../bsplines/evaluation_kernels_3d.cuh | 18 + .../bstar_parallel_3form_cuda.cu | 75 +++ .../pic/tests/codegen_spike/__init__.py | 1 + .../pic/tests/codegen_spike/compare.py | 61 +++ .../pic/tests/codegen_spike/numba_codegen.py | 458 ++++++++++++++++++ .../tests/codegen_spike/test_numba_codegen.py | 29 ++ src/struphy/pic/tests/cuda_parity_cases.py | 32 ++ 8 files changed, 762 insertions(+), 3 deletions(-) create mode 100644 src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu create mode 100644 src/struphy/pic/tests/codegen_spike/__init__.py create mode 100644 src/struphy/pic/tests/codegen_spike/compare.py create mode 100644 src/struphy/pic/tests/codegen_spike/numba_codegen.py create mode 100644 src/struphy/pic/tests/codegen_spike/test_numba_codegen.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..2a621ce88 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,7 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -315,7 +315,7 @@ algorithm is `discrete_gradient_1st_order`; its chain comes first. | 16–17 | `push_gc_bxEstar_discrete_gradient_2nd_order`, `push_gc_Bstar_discrete_gradient_2nd_order` | P | | 18–19 | `push_gc_bxEstar_explicit_multistage`, `push_gc_Bstar_explicit_multistage` | P | -These are the largest kernels; evaluate code generation (see [Open questions](#open-questions)) before porting 13–19 by hand. +These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. **Step 4 – Hybrid MHD–kinetic 6D (current and pressure coupling)** @@ -366,13 +366,98 @@ matrix data). - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. -- **Single-source alternatives.** Hand-written CUDA stays the default. Generating whole kernels from the Python source (`cupyx.jit`, numba-cuda, or a pyccel CUDA backend) is worth a look before the guiding-center kernels (the largest ones) are ported. Those tools take flat arguments, which `fields` also provides. +- **Single-source alternatives.** Decided for now (#690): hand-written CUDA stays the default; see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690) for the comparison and for what would change the decision. - **Polar splines and MHD equilibria on the GPU** (left after PR 19): spline mappings run on the device, but `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels), and `EQDSKequilibrium` cannot be created on CuPy (its SciPy splines get device arrays). A `Tokamak` on CuPy builds its default equilibrium on the host. Needed once a model with a polar domain or an EQDSK equilibrium runs on the GPU. +## Hand-written CUDA vs. code generation (decision for #690) + +**Decision: the guiding-center and hybrid kernels (steps 3–5 of the [porting order](#porting-order)) are ported by +hand, as `_cuda.cu` next to the pyccel kernel, like every kernel so far.** None of the generators is adopted. +This is decided again if a condition under [What would change the decision](#what-would-change-the-decision) +comes true. + +To decide, one guiding-center kernel, `bstar_parallel_3form` (step 3, #7: drift-kinetic B*∥ as a 3-form; mapping +Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported in two ways: + +- **By hand:** `pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel, with + 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping (all analytic and four spline mappings), + with α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)`, a negative `eta_n`, a hole + and two output columns. This port stays in the code. +- **Generated:** a prototype generator, `pic/tests/codegen_spike/numba_codegen.py`, that translates the pyccel + source of the kernel *and of every helper it calls* into one numba.cuda module (`compare.py` runs it, + `test_numba_codegen.py` tests it; skipped without numba). It only exists for this comparison and can be + deleted. There is no `kernels` in the module names, so `struphy compile` does not pick it up. + +### Results + +| | Hand-written CUDA | Generated (numba.cuda) | +|---|---|---| +| Code for this kernel | 75 lines `.cu` (36 without comments), plus one new 5-line helper `eval_0form_spline_mpi` in `bsplines/evaluation_kernels_3d.cuh`; everything else is reused (`df` for 14 mappings, `det`, `get_spans`, `eval_spline_mpi_kernel`) | 0 lines by hand, but a 458-line generator. The generated module has 490 lines and contains its own copy of the whole helper chain (24 device functions); every generated kernel repeats it | +| Agreement with pyccel (14 cases) | CPU emulation of the real `.cu` source: all agree, max abs. difference 5.7·10⁻¹⁴. Removing the `mod` from the CUDA kernel makes the emulation test fail (115 of 3225 values) | numba.cuda **simulator**: all agree, max 6.3·10⁻¹³; the same translation compiled with `numba.njit`: the same values. The simulator interprets the Python one thread at a time and compiles nothing | +| Compiles for a GPU? | Not checked here (no GPU); the same NVRTC path as the 15 kernels before it | Not checked (numba.cuda needs NVVM from the CUDA toolkit, which macOS does not have). Not checked either: named tuples of local arrays in device code, and the 38 flat kernel parameters | +| Argument classes | Takes `MarkerArgs`, `DomainArgs`, `DerhamArgs` by value, like every kernel; cunumpy's `CudaKernel` checks the call against the signature | numba does not accept the C structs. The generator turns each argument object into one kernel parameter per struct field (38 parameters) and rebuilds the objects as named tuples inside the kernel. Calling it needs a second launch path next to cunumpy's (flattening, numba's array conversion, its own compile cache) | +| 1:1 with pyccel | Names, arguments and order as in pyccel, kept by review | Exactly 1:1 for the parts the generator understands; whatever it does not understand fails with `NotImplementedError` (see the last row) | +| Parity tests, emulation | The existing tests run unchanged: `PARITY_CASES` entry, GPU parity with `assert_kernels_agree`, CPU emulation, which compiles the `.cu` itself | The parity cases are reused, but the simulator checks the Python, not the code the GPU will run. cunumpy's emulation and `assert_kernels_agree` do not apply to numba kernels | +| Debugging | A C++ compiler and printf in emulation; `-lineinfo`, cunumpy's debug mode, compute-sanitizer and Nsight on the GPU; the source you read is the source you run | pdb in the simulator, where Python float arithmetic hides bugs specific to the GPU. On the GPU, errors point into generated code. The generator's own mistakes are a new kind of bug (see below) | +| Expected performance (no GPU numbers possible here) | Per-thread arrays with fixed size (`MAX_SPLINE_DEGREE`), `__device__ inline` helpers; atomics where needed (`cunumpy_atomic_add`) | About the same local memory (the same `MAX_LOCAL` bounds; `b_der_splines_slim` has the same 9×9 `values`), plus loops that fill `zeros` arrays. numba applies Python's wrap-around for negative indices to every array access (one extra compare/select each), passes shape, strides and data pointer per array, and has a higher launch overhead than `RawKernel`. The generator cannot tell which writes need atomics: a translated accumulation kernel would race | +| Dependencies | None beyond CuPy | numba + llvmlite + numba-cuda (≈175 MB installed), and numba only supports a range of NumPy versions (numba 0.68 was needed for NumPy 2.4). The numba.cuda that comes with numba is deprecated in favour of NVIDIA's numba-cuda | +| The other 59 pusher and accumulation kernels | Each is a transcription onto existing helpers; new helpers as needed (e.g. 1-/2-form evaluations exist, `eval_0form_spline_mpi` was added here) | Tried on all 60 entry kernels (`.claude/scripts/codegen_probe.py`, not committed): 31 translate, 29 do not (2D slice assignments such as `c[:, :] = ...` in the linear-algebra helpers, `numpy.abs` in the SPH helpers). "Translates" does not mean "is correct": none of them ran, and the accumulation kernels would race | + +What the generator had to change in the pyccel code, and the risks that come with it: + +- `empty`/`zeros` with a size known only at run time (`zeros(pn)`) become per-thread arrays of `MAX_LOCAL` + entries. This silently changes `len(x)`/`x.shape` of these arrays, e.g. `diff[:] = diff * pn` in + `b_der_splines_slim` now loops over 9 entries instead of `pn` (harmless there, because the extra entries are + zero). +- Array expressions on slices (`eta[:] = alpha[:3] * eta_k + ...`) become loops over the elements, because numba.cuda + cannot allocate the temporaries. Expressions with 2D slices are not supported yet. +- `numpy.mod` has to be translated exactly as pyccel compiles it (Fortran `MODULO`, `a - floor(a / b) * b`). +- The scratch arrays of `DerhamArguments` (`bn1`, ..., `bd3`), which every marker shares in pyccel, have to become + per-thread arrays; the hand-written `SplineScratch` does the same. +- The marker loop becomes one thread per marker; `continue` becomes `return`. + +**The other generators.** `cupyx.jit` (CuPy's experimental Python-to-CUDA transpiler) would need the same +translation of the pyccel source: its reference lists elementary scalar operations, thread indices, shared memory +and atomics, but no per-thread arrays and no struct arguments, and "CuPy's array operations ... are not supported". +It adds no dependency, but it was not run here (no CuPy). **pyccel 2.2.1 has no CUDA backend** (`--language` is +Fortran, C or Python). The CUDA work happens in the separate fork [pyccel-cuda](https://github.com/pyccel/pyccel-cuda) +(`@kernel`/`@device`, last commit January 2025); its open issues still include thread indexing, passing arrays as +arguments and device arrays, and nothing there handles class arguments like `MarkerArguments`. + +### Reasoning + +1. **The shared part is already written.** What a generator would save is mostly the helper chain (mappings, + metric, splines, linear algebra, boundary conditions), and PRs 10–19 have already ported it by hand and tested + it, every mapping included. What is left per kernel is its body: here 36 lines that read like the pyccel code. +2. **A generator is a second compiler to maintain.** The prototype has 458 lines for a subset of what the kernels + use, and its mistakes (wrong array lengths, missing atomics) are silent. pyccel already is the compiler for the + CPU; a second one written by ourselves only pays off for many more kernels than ~45. +3. **It does not fit the infrastructure.** Everything since PR 8 (structs, signature checks in `CudaKernel`, + compile cache, debug mode, `assert_kernels_agree`, CPU emulation of the real source, no host transfers) is + built for CUDA C++ source. numba kernels would need a second launch and test path in cunumpy and struphy. +4. **It is harder to test.** The CPU emulation compiles the `.cu` that will run on the GPU; the numba simulator + only interprets Python. The hand-written port is gated better without a GPU than the generated one. +5. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change + to a pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. + +### What would change the decision + +- **pyccel gets a CUDA backend** in a release (not the fork) that accepts our argument classes and arrays. Then the + `.cu` files could be generated by pyccel in place, behind the same `Kernel`, structs and tests; the parity cases + would check them unchanged. This is the most attractive single-source route; check pyccel's releases before + step 4. +- **Porting by hand becomes the bottleneck**, e.g. new models add many kernels, or the pyccel kernels change often. + Then the next step is a generator that writes **`.cu` text** (C++) instead of numba: the AST transformations of + the prototype (helper calls, local arrays, slice loops, marker loop to threads, `mod`) carry over, and the output + goes through the existing structs, `CudaKernel`, emulation and parity tests. Accumulation kernels would still need + their atomics marked by hand. +- **GPU measurements** that show a generated kernel to be as fast as the hand-written one would remove the + performance argument; they do not touch points 2–4. + ## PR 10 implementation notes Device headers now provide Cuboid Jacobians, matrix inversion/vector products, diff --git a/src/struphy/bsplines/evaluation_kernels_3d.cuh b/src/struphy/bsplines/evaluation_kernels_3d.cuh index ccb10b6f7..3c4c353ad 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.cuh +++ b/src/struphy/bsplines/evaluation_kernels_3d.cuh @@ -71,6 +71,24 @@ __device__ inline double eval_spline_mpi_kernel(int p1, int p2, int p3, const do return spline_value; } +/** + * Evaluate a 0-form spline, as in evaluation_kernels_3d.eval_0form_spline_mpi. + * + * @param span1 Knot span index along the first axis (from get_spans). + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param args_derham Spline degrees and start indices. + * @param scratch N-spline values from get_spans; pyccel reads them from args_derham.bn1, bn2, bn3. + * @param form_coeffs Coefficients of the 0-form (N N N). + * @return out, the value of the 0-form at the point of get_spans. + */ +__device__ inline double eval_0form_spline_mpi(int span1, int span2, int span3, const DerhamArgs& args_derham, + const SplineScratch& scratch, Array3D form_coeffs) { + double out = eval_spline_mpi_kernel(args_derham.pn[0], args_derham.pn[1], args_derham.pn[2], scratch.bn1, + scratch.bn2, scratch.bn3, span1, span2, span3, form_coeffs, args_derham.starts); + return out; +} + /** * Evaluate the three components of a 1-form spline, as in evaluation_kernels_3d.eval_1form_spline_mpi. * diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu new file mode 100644 index 000000000..ec7389284 --- /dev/null +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -0,0 +1,75 @@ +#include "struphy/bsplines/evaluation_kernels_3d.cuh" +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" + +/** + * Evaluate B*_parallel as a 3-form, as in bstar_parallel_3form_kernels.bstar_parallel_3form. + * + * B^{*3}_parallel(Z_p) = sqrt(g) (B . b + epsilon v_parallel,p ((curl b_0) . b_0)(eta_p)) at the weighted average + * Z_{p,i} = alpha_i Z^{n+1,k}_{p,i} + (1 - alpha_i) Z^n_{p,i} (i = 1, ..., 4); one thread per marker row. + * + * @param alpha Weights of the average (at least four entries): alpha[0:3] for eta, alpha[3] for v_parallel. + * @param output_indices Marker column of the result in output_indices[0]; -1 skips the output. + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param epsilon Scaling parameter epsilon in front of v_parallel. + * @param B_dot_b_coeffs Coefficients of the 0-form B . b. + * @param curl_unit_b_dot_b0 Coefficients of the 0-form (curl b_0) . b_0. + * + * Holes (markers[ip, 0] == -1) are skipped, as in pyccel (which does not use valid_mks here). + */ +extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long long* output_indices, + MarkerArgs args_markers, DomainArgs args_domain, + DerhamArgs args_derham, double epsilon, + Array3D B_dot_b_coeffs, Array3D curl_unit_b_dot_b0) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // allocate stack arrays + double eta_k[3], eta_n[3], eta[3], dfm[9]; + + // get marker arguments + int first_init_idx = args_markers.first_init_idx; + int first_shift_idx = args_markers.first_shift_idx; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1.0) return; + + for (int j = 0; j < 3; ++j) { + eta_k[j] = args_markers.markers(ip, j) + args_markers.markers(ip, first_shift_idx + j); + eta_n[j] = args_markers.markers(ip, first_init_idx + j); + + eta[j] = alpha[j] * eta_k[j] + (1.0 - alpha[j]) * eta_n[j]; + // numpy.mod(eta, 1.0) (Fortran MODULO): the result has the sign of the divisor + eta[j] = eta[j] - floor(eta[j]); + } + + double v_k = args_markers.markers(ip, 3); + double v_n = args_markers.markers(ip, first_init_idx + 3); + double v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n; + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta[0], eta[1], eta[2], args_domain, dfm); + + double det_df = struphy_cuda::det(dfm); + + // spline evaluation; CUDA-only scratch holds the spline values pyccel keeps in args_derham + struphy_cuda::SplineScratch scratch; + struphy_cuda::get_spans(eta[0], eta[1], eta[2], args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // compute B*_parallel + double B_dot_b = struphy_cuda::eval_0form_spline_mpi(span1, span2, span3, args_derham, scratch, B_dot_b_coeffs); + + double b_star_parallel = + struphy_cuda::eval_0form_spline_mpi(span1, span2, span3, args_derham, scratch, curl_unit_b_dot_b0); + + b_star_parallel *= epsilon * v; + b_star_parallel += B_dot_b; + b_star_parallel *= det_df; + + if (output_indices[0] >= 0) args_markers.markers(ip, output_indices[0]) = b_star_parallel; +} diff --git a/src/struphy/pic/tests/codegen_spike/__init__.py b/src/struphy/pic/tests/codegen_spike/__init__.py new file mode 100644 index 000000000..d1b189c74 --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/__init__.py @@ -0,0 +1 @@ +"""Prototype for issue #690 (hand-written CUDA vs. code generation); not used by struphy, see ``numba_codegen``.""" diff --git a/src/struphy/pic/tests/codegen_spike/compare.py b/src/struphy/pic/tests/codegen_spike/compare.py new file mode 100644 index 000000000..ae72c1a6f --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/compare.py @@ -0,0 +1,61 @@ +"""Run the numba kernel generated from ``bstar_parallel_3form`` on its parity cases and compare it with pyccel. + +Usage (numba is not a struphy dependency):: + + NUMBA_ENABLE_CUDASIM=1 python -m struphy.pic.tests.codegen_spike.compare cuda # numba.cuda simulator + python -m struphy.pic.tests.codegen_spike.compare cpu # same translation, numba.njit + +The simulator interprets the generated Python thread by thread (it does not compile anything); ``cpu`` compiles the +same translation in nopython mode. Prints one line per case and a summary; exits with 1 if a case disagrees. +""" + +import sys +import time + +import cunumpy as xp +import numpy as np + +from struphy.pic.tests.codegen_spike.numba_codegen import flatten, load + +MODULE = "struphy.pic.pushing.kernels.bstar_parallel_3form.bstar_parallel_3form_kernels" +NAME = "bstar_parallel_3form" +BLOCK = 128 + + +def compare(target): + from struphy.pic.pushing.kernels.bstar_parallel_3form import bstar_parallel_3form + from struphy.pic.tests.cuda_parity_cases import PARITY_CASES + from struphy.pic.tests.test_cuda_emulation import arrays + + kernel, path = load(MODULE, NAME, target) + print(f"generated {path} ({len(path.read_text().splitlines())} lines)") + spec = PARITY_CASES[NAME] + failures, bitwise = 0, 0 + with xp.use_backend("numpy"): + for case in spec.cases: + reference_args, generated_args = spec.build(case), spec.build(case) + bstar_parallel_3form(*reference_args) + flat = flatten(generated_args) + start = time.perf_counter() + if target == "cuda": + n_markers = generated_args[2].n_markers + kernel[(n_markers + BLOCK - 1) // BLOCK, BLOCK](*flat) + else: + kernel(*flat) + seconds = time.perf_counter() - start + reference, generated = arrays(reference_args), arrays(generated_args) + diff = max( + float(np.max(np.abs(r.astype(float) - g.astype(float)), initial=0.0)) + for r, g in zip(reference, generated) + ) + same = all(np.array_equal(r, g, equal_nan=True) for r, g in zip(reference, generated)) + ok = all(np.allclose(g, r, rtol=spec.rtol, atol=spec.atol) for r, g in zip(reference, generated)) + failures += not ok + bitwise += same + print(f"domain {case[0]:2d}: max |diff| = {diff:.2e}, bitwise equal: {same}, {seconds:.2f} s") + print(f"{target}: {len(spec.cases) - failures}/{len(spec.cases)} cases agree, {bitwise} bitwise") + return failures + + +if __name__ == "__main__": + sys.exit(1 if compare(sys.argv[1] if len(sys.argv) > 1 else "cpu") else 0) diff --git a/src/struphy/pic/tests/codegen_spike/numba_codegen.py b/src/struphy/pic/tests/codegen_spike/numba_codegen.py new file mode 100644 index 000000000..062df877d --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/numba_codegen.py @@ -0,0 +1,458 @@ +"""Prototype for issue #690: generate a ``numba.cuda`` kernel from the Python source of a pyccel kernel. + +Not used by struphy. It exists to compare code generation with the hand-written ``_cuda.cu`` files on one +real kernel (``bstar_parallel_3form``); the result of the comparison is in ``CUDA_STRATEGY.md`` (decision for +#690). The module name contains no ``kernels``, so ``struphy compile`` ignores it. + +The generator works on the source files (the imported kernel modules are usually the compiled pyccel extensions): + +- it parses the entry kernel's module, follows every call into struphy modules (``import a.b as x`` and + ``from a.b import f``) and translates each reached function once, as a numba device function named + ``__``; +- it strips the pyccel decorators and type annotations, maps the NumPy scalar functions to ``math``, and + ``numpy.mod`` to the Fortran ``MODULO`` pyccel generates; +- ``empty``/``zeros`` become ``cuda.local.array`` with a constant size (a run-time size such as ``zeros(pn)`` gets + ``MAX_LOCAL`` entries, like the hand-written ``MAX_SPLINE_DEGREE`` arrays); +- array expressions on slices (``eta[:] = alpha[:3] * eta_k + ...``, ``bn[:] = 0.0``) become element loops, since + numba.cuda cannot allocate the temporaries; +- the entry kernel's marker loop ``for ip in range(n): ...`` becomes ``ip = cuda.grid(1)`` and ``continue`` becomes + ``return``; +- argument objects (``MarkerArguments``, ``DomainArguments``, ``DerhamArguments``) are passed as the fields of + their CUDA structs (``CudaMarkerArguments.struct.fields``, one kernel parameter each) and rebuilt as named tuples + at the top of the kernel; ``DerhamArguments`` gets its scratch arrays ``bn1``, ..., ``bd3`` as per-thread local + arrays there, so the helpers can stay one to one with pyccel. + +With ``target="cpu"`` the same translation is compiled with ``numba.njit`` (NumPy allocations, the marker loop kept), +which checks that the generated code types in nopython mode; numba.cuda itself needs a GPU (or NVVM) to compile, and +its simulator (``NUMBA_ENABLE_CUDASIM=1``) only interprets the Python. +""" + +import ast +import importlib.util +import re +import sys +import tempfile +from pathlib import Path + +from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments + +MAX_LOCAL = 9 # MAX_SPLINE_DEGREE + 1 of bsplines_kernels.cuh + +STRUCT_CLASSES = { + "MarkerArguments": CudaMarkerArguments, + "DerhamArguments": CudaDerhamArguments, + "DomainArguments": CudaDomainArguments, +} +# pyccel scratch arrays of the argument classes: per-thread local arrays in the generated kernel +SCRATCH = {"DerhamArguments": ("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")} + +NUMPY_MATH = { + "cos": "cos", + "sin": "sin", + "tan": "tan", + "arctan": "atan", + "arcsin": "asin", + "arccos": "acos", + "arctan2": "atan2", + "sqrt": "sqrt", + "exp": "exp", + "log": "log", + "floor": "floor", +} +NUMPY_ALLOC = ("empty", "zeros") + + +def _source_file(module): + """The ``.py`` file of `module` (the import system would find the compiled pyccel extension first).""" + spec = importlib.util.find_spec(module) + return Path(spec.origin).parent / (module.rsplit(".", 1)[-1] + ".py") + + +class _Module: + """Functions and imports of one source module.""" + + def __init__(self, name): + self.name = name + self.short = name.rsplit(".", 1)[-1] + tree = ast.parse(_source_file(name).read_text()) + self.functions = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)} + self.module_aliases = {} # alias -> module + self.imported = {} # name -> (module, name) + for node in tree.body: + if isinstance(node, ast.Import): + for a in node.names: + self.module_aliases[a.asname or a.name] = a.name + elif isinstance(node, ast.ImportFrom): + for a in node.names: + self.imported[a.asname or a.name] = (node.module, a.name) + + def resolve(self, func): + """``(module, name)`` of a called function node, or ``None``.""" + if isinstance(func, ast.Name): + if func.id in self.functions: + return self.name, func.id + return self.imported.get(func.id) + if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + module = self.module_aliases.get(func.value.id) + if module is not None: + return module, func.attr + return None + + +def _array_ndim(annotation): + """Number of dimensions of a pyccel array annotation such as ``"float[:,:]"``, or 0.""" + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + match = re.search(r"\[([:,\s]+)\]", annotation.value) + if match: + return match.group(1).count(":") + return 0 + + +def _struct_type(annotation): + if isinstance(annotation, ast.Constant) and annotation.value in STRUCT_CLASSES: + return annotation.value + if isinstance(annotation, ast.Name) and annotation.id in STRUCT_CLASSES: + return annotation.id + return None + + +def _expr(source): + return ast.parse(source, mode="eval").body + + +def _stmt(source): + return ast.parse(source).body + + +class _Elementwise(ast.NodeTransformer): + """Index a slice expression at element ``index``: ``a[lo:hi]`` -> ``a[lo + index]``, 1D array ``x`` -> ``x[index]``.""" + + def __init__(self, index, arrays): + self.index = index + self.arrays = arrays + + def _element(self, slc): + lower = slc.lower if slc.lower is not None else ast.Constant(0) + return ast.BinOp(lower, ast.Add(), ast.Name(self.index, ast.Load())) + + def visit_Subscript(self, node): + if isinstance(node.slice, ast.Slice): + return ast.Subscript(node.value, self._element(node.slice), node.ctx) + if isinstance(node.slice, ast.Tuple) and any(isinstance(e, ast.Slice) for e in node.slice.elts): + elts = [self._element(e) if isinstance(e, ast.Slice) else e for e in node.slice.elts] + return ast.Subscript(node.value, ast.Tuple(elts, ast.Load()), node.ctx) + return node + + def visit_Name(self, node): + if self.arrays.get(node.id) == 1: + return ast.Subscript(ast.Name(node.id, ast.Load()), ast.Name(self.index, ast.Load()), node.ctx) + if node.id in self.arrays and self.arrays[node.id] > 1: + raise NotImplementedError(f"array expression on the {self.arrays[node.id]}D array {node.id}") + return node + + +def _slice_length(target): + """Number of elements written by a sliced assignment target (one slice per target).""" + slices = target.slice.elts if isinstance(target.slice, ast.Tuple) else [target.slice] + positions = [i for i, s in enumerate(slices) if isinstance(s, ast.Slice)] + if len(positions) != 1: + raise NotImplementedError(f"assignment to {ast.unparse(target)}: exactly one slice is supported") + slc = slices[positions[0]] + if slc.step is not None: + raise NotImplementedError(f"assignment to {ast.unparse(target)}: slice steps are not supported") + if slc.upper is None: + upper = ast.Subscript(ast.Attribute(target.value, "shape", ast.Load()), ast.Constant(positions[0]), ast.Load()) + else: + upper = slc.upper + if slc.lower is None: + return upper + return ast.BinOp(upper, ast.Sub(), slc.lower) + + +def _has_slice(target): + if not isinstance(target, ast.Subscript): + return False + s = target.slice + return isinstance(s, ast.Slice) or (isinstance(s, ast.Tuple) and any(isinstance(e, ast.Slice) for e in s.elts)) + + +class _ContinueToReturn(ast.NodeTransformer): + """``continue`` of the marker loop -> ``return`` (nested loops keep theirs).""" + + def visit_Continue(self, node): + return ast.Return(None) + + def visit_For(self, node): + return node + + def visit_While(self, node): + return node + + +class Generator: + """Translate a pyccel entry kernel and every struphy function it calls into one numba module.""" + + def __init__(self, target="cuda"): + assert target in ("cuda", "cpu") + self.target = target + self.modules = {} + self.translated = {} # (module, name) -> mangled name + self.sources = [] # translated functions, callees first + self.counter = 0 + + def _module(self, name): + if name not in self.modules: + self.modules[name] = _Module(name) + return self.modules[name] + + def _is_struphy_function(self, ref): + if ref is None or not ref[0].startswith("struphy."): + return False + try: + return ref[1] in self._module(ref[0]).functions + except (ModuleNotFoundError, FileNotFoundError, AttributeError): + return False + + def _numpy_name(self, module, node): + """The NumPy name a call or name node refers to in `module`, or ``None``.""" + if isinstance(node, ast.Name): + ref = module.imported.get(node.id) + if ref is not None and ref[0] == "numpy": + return ref[1] + return None + + def _fresh(self, prefix): + self.counter += 1 + return f"{prefix}{self.counter}" + + # ------------------------------------------------------------------ allocations + + def _allocation(self, module, target, call, arrays): + kind = self._numpy_name(module, call.func) + shape = call.args[0] + ndim = len(shape.elts) if isinstance(shape, ast.Tuple) else 1 + dtype = "int64" + for kw in call.keywords: + if kw.arg == "dtype" and isinstance(kw.value, ast.Name) and kw.value.id == "float": + dtype = "float64" + if len(call.args) > 1 and isinstance(call.args[1], ast.Name) and call.args[1].id == "float": + dtype = "float64" + if not call.keywords and len(call.args) == 1: + dtype = "float64" + arrays[target] = ndim + if self.target == "cpu": + return _stmt(f"{target} = np.{kind}({ast.unparse(shape)}, dtype=np.{dtype})") + elts = shape.elts if isinstance(shape, ast.Tuple) else [shape] + sizes = [e.value if isinstance(e, ast.Constant) else MAX_LOCAL for e in elts] + size = sizes[0] if ndim == 1 else "(" + ", ".join(map(str, sizes)) + ")" + stmts = _stmt(f"{target} = cuda.local.array({size}, {dtype})") + if kind == "zeros": + loops = "".join(f"{' ' * d}for _z{d} in range({s}):\n" for d, s in enumerate(sizes)) + index = ", ".join(f"_z{d}" for d in range(ndim)) + stmts += _stmt(f"{loops}{' ' * ndim}{target}[{index}] = 0") + return stmts + + # ------------------------------------------------------------------ functions + + def translate(self, module_name, name, entry=False): + """Translate function `name` of `module_name`; returns its name in the generated module.""" + key = (module_name, name) + if key in self.translated: + return self.translated[key] + module = self._module(module_name) + mangled = f"{module.short}__{name}" + self.translated[key] = mangled + node = ast.parse(ast.unparse(module.functions[name])).body[0] # a fresh copy + + # parameters: array dimensions and argument objects from the pyccel annotations + arrays = {a.arg: _array_ndim(a.annotation) for a in node.args.args if _array_ndim(a.annotation)} + structs = {a.arg: _struct_type(a.annotation) for a in node.args.args if _struct_type(a.annotation)} + for a in node.args.args: + a.annotation = None + node.returns = None + node.decorator_list = [] + node.name = mangled + if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant): + node.body = node.body[1:] or [ast.Pass()] + + node.body = self._statements(module, node.body, arrays) + + if entry: + self._make_entry(node, structs) + decorator = "@cuda.jit" if self.target == "cuda" else "@njit" + else: + decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" + self.sources.append(f"{decorator}\n{ast.unparse(ast.fix_missing_locations(node))}\n") + return mangled + + def _statements(self, module, body, arrays): + out = [] + for stmt in body: + out += self._statement(module, stmt, arrays) + return out + + def _statement(self, module, stmt, arrays): + # local arrays + if ( + isinstance(stmt, ast.Assign) + and isinstance(stmt.value, ast.Call) + and self._numpy_name(module, stmt.value.func) in NUMPY_ALLOC + ): + (target,) = stmt.targets + return self._allocation(module, target.id, stmt.value, arrays) + # array expressions on slices -> element loops + if isinstance(stmt, (ast.Assign, ast.AugAssign)): + target = stmt.targets[0] if isinstance(stmt, ast.Assign) else stmt.target + if _has_slice(target): + index = self._fresh("_k") + length = _slice_length(target) + element = _Elementwise(index, arrays) + value = self._expressions(module, element.visit(stmt.value)) + new_target = element.visit(target) + if isinstance(stmt, ast.Assign): + inner = ast.Assign([new_target], value) + else: + inner = ast.AugAssign(new_target, stmt.op, value) + loop = ast.For( + ast.Name(index, ast.Store()), + ast.Call(ast.Name("range", ast.Load()), [length], []), + [inner], + [], + lineno=stmt.lineno, + ) + return [loop] + # compound statements: recurse into their bodies + for field in ("body", "orelse"): + if hasattr(stmt, field) and isinstance(getattr(stmt, field), list): + setattr(stmt, field, self._statements(module, getattr(stmt, field), arrays)) + for field in ("value", "test", "iter"): + if hasattr(stmt, field) and isinstance(getattr(stmt, field), ast.AST): + setattr(stmt, field, self._expressions(module, getattr(stmt, field))) + if isinstance(stmt, ast.Assign): + stmt.targets = [self._expressions(module, t) for t in stmt.targets] + if isinstance(stmt, ast.AugAssign): + stmt.target = self._expressions(module, stmt.target) + return [stmt] + + def _expressions(self, module, node): + generator = self + + class Calls(ast.NodeTransformer): + def visit_Call(self, call): + self.generic_visit(call) + ref = module.resolve(call.func) + numpy_name = generator._numpy_name(module, call.func) + if generator._is_struphy_function(ref): + call.func = ast.Name(generator.translate(*ref), ast.Load()) + elif numpy_name in NUMPY_MATH: + call.func = _expr(f"math.{NUMPY_MATH[numpy_name]}") + elif numpy_name == "mod": + call.func = ast.Name("_numpy_mod", ast.Load()) + elif numpy_name == "shape": + return ast.Attribute(call.args[0], "shape", ast.Load()) + elif numpy_name is not None: + raise NotImplementedError(f"numpy.{numpy_name} in {module.name}") + return call + + def visit_Name(self, name): + if generator._numpy_name(module, name) == "pi": + return _expr("math.pi") + return name + + return Calls().visit(node) + + def _make_entry(self, node, structs): + """Flatten the argument objects into struct fields, rebuild them, and map the marker loop to threads.""" + params, prologue = [], [] + for a in node.args.args: + kind = structs.get(a.arg) + if kind is None: + params.append(a.arg) + continue + fields = [f.name for f in STRUCT_CLASSES[kind].struct.fields] + params += [f"{a.arg}__{f}" for f in fields] + values = [f"{a.arg}__{f}" for f in fields] + if self.target == "cuda": + values += [f"cuda.local.array({MAX_LOCAL}, float64)" for _ in SCRATCH.get(kind, ())] + else: + values += [f"np.empty({MAX_LOCAL})" for _ in SCRATCH.get(kind, ())] + prologue += _stmt(f"{a.arg} = {kind}_T({', '.join(values)})") + node.args.args = [ast.arg(p) for p in params] + + body = prologue + node.body + if self.target == "cuda": + for i, stmt in enumerate(body): + if ( + isinstance(stmt, ast.For) + and isinstance(stmt.iter, ast.Call) + and getattr(stmt.iter.func, "id", None) == "range" + and len(stmt.iter.args) == 1 + ): + ip, n = stmt.target.id, ast.unparse(stmt.iter.args[0]) + loop_body = [_ContinueToReturn().visit(s) for s in stmt.body] + body = body[:i] + _stmt(f"{ip} = cuda.grid(1)\nif {ip} >= {n}:\n return") + loop_body + break + else: + raise NotImplementedError("no marker loop 'for ip in range(n)' in the entry kernel") + node.body = body + + def module_source(self): + decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" + header = [ + f"# Generated by struphy.pic.tests.codegen_spike.numba_codegen (target={self.target}); do not edit.", + "import math", + "from collections import namedtuple", + "", + "import numpy as np", + "from numba import cuda, float64, int64, njit", + "", + ] + for kind, cls in STRUCT_CLASSES.items(): + names = [f.name for f in cls.struct.fields] + list(SCRATCH.get(kind, ())) + header.append(f"{kind}_T = namedtuple({kind + '_T'!r}, {names!r})") + header += [ + "", + "", + decorator, + "def _numpy_mod(a, b):", + " # numpy.mod as pyccel compiles it (Fortran MODULO)", + " return a - math.floor(a / b) * b", + "", + "", + ] + return "\n".join(header) + "\n\n".join(self.sources) + + +def generate(module_name, name, target="cuda"): + """Source of the numba module for the pyccel kernel `name` in `module_name` (and every function it calls).""" + generator = Generator(target) + generator.translate(module_name, name, entry=True) + return generator.module_source() + + +def load(module_name, name, target="cuda", directory=None): + """Generate, write to `directory` (a temporary one by default) and import; returns the kernel function.""" + source = generate(module_name, name, target) + directory = Path(directory or tempfile.mkdtemp(prefix="struphy_numba_codegen_")) + stem = f"generated_{name}_{target}" + path = directory / f"{stem}.py" + path.write_text(source) + spec = importlib.util.spec_from_file_location(stem, path) + module = importlib.util.module_from_spec(spec) + sys.modules[stem] = module + spec.loader.exec_module(module) + return getattr(module, f"{module_name.rsplit('.', 1)[-1]}__{name}"), path + + +def flatten(args): + """Kernel arguments with every argument object replaced by the fields of its CUDA struct (host arrays).""" + from struphy.pic.tests.cuda_emulation import _struct_field + + flat = [] + for value in args: + kind = type(value).__name__ + if kind in STRUCT_CLASSES: + flat += [_struct_field(value, f.name) for f in STRUCT_CLASSES[kind].struct.fields] + else: + flat.append(value) + return flat diff --git a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py new file mode 100644 index 000000000..7648440b6 --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py @@ -0,0 +1,29 @@ +"""The numba kernel generated from ``bstar_parallel_3form`` agrees with pyccel (issue #690 prototype). + +Skipped unless numba (with numba-cuda for the simulator) is installed; numba is not a struphy dependency. +Runs in a subprocess because ``NUMBA_ENABLE_CUDASIM`` must be set before numba is imported. +""" + +import importlib.util +import subprocess +import sys + +import pytest + +from struphy.geometry.tests.test_domain import serial_child_env + +requires_numba = pytest.mark.skipif(importlib.util.find_spec("numba") is None, reason="numba is not installed") + + +@requires_numba +@pytest.mark.parametrize("target", ["cpu", "cuda"]) +def test_generated_kernel_agrees_with_pyccel(target): + """``cpu``: the translation compiled with numba.njit; ``cuda``: the numba.cuda kernel in the CUDA simulator.""" + extra = {"NUMBA_ENABLE_CUDASIM": "1"} if target == "cuda" else {} + result = subprocess.run( + [sys.executable, "-m", "struphy.pic.tests.codegen_spike.compare", target], + env=serial_child_env(**extra), + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..7ee91c0c0 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -71,6 +71,31 @@ def reflect_args(axis): return (xp.asarray(markers), Cuboid(r1=2.0, r2=3.0, r3=4.0).args_domain, xp.asarray(outside_inds), axis) +# guiding-centre evaluation kernels: every mapping, the evaluation point at the new position (alpha = 1), the old one +# (alpha = 0) and in between; the output column alternates. (index into geometry_domain(), alpha, output column) +BSTAR_PARALLEL_ALPHAS = ((1.0,) * 6, (0.0,) * 6, (0.3, 0.6, 0.9, 0.5, 0.0, 0.0)) +BSTAR_PARALLEL_3FORM_CASES = tuple( + (domain, BSTAR_PARALLEL_ALPHAS[domain % 3], 19 if domain % 2 == 0 else 24) for domain in range(N_GEOMETRY_DOMAINS) +) + + +def bstar_parallel_3form_args(case): + """Markers with shifts in [0, 1), so eta + shift wraps around, a negative eta_n in row 0 and a hole in row 2.""" + domain, alpha, output_column = case + args_markers, _ = marker_arguments((0, 0, 0)) + args_markers.markers[0, 8] = -0.75 # eta_n[0] = mod(-0.75, 1.0) = 0.25 (not 0, where some Jacobians are singular) + args_markers.markers[2, 0] = -1.0 # a hole: pyccel tests markers[ip, 0] == -1, not valid_mks + return ( + xp.asarray(alpha), + xp.asarray([output_column], dtype=np.int64), + args_markers, + geometry_domain(domain).args_domain, + derham_arguments(), + 0.37, + *spline_coefficients(n=2, seed=23), + ) + + # ---------------------------------------------------------------- accumulation @@ -265,6 +290,13 @@ def size_of(index, per_thread=1): BOUNDARY_CONDITIONS, push_weights_with_efield_lin_va_args, **PUSHER_TOLERANCES ), "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), + # one thread per marker row; the first array, alpha, does not tell the launch size + "bstar_parallel_3form": ParityCases( + BSTAR_PARALLEL_3FORM_CASES, + bstar_parallel_3form_args, + n_threads=lambda args: args[2].n_markers, + **GEOMETRY_TOLERANCES, + ), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), From 7d18d8cb01cec3e862f47719a57c97a5f8659128 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 11:29:47 +0200 Subject: [PATCH 2/3] Remove the numba code-generation prototype (#690 decided: hand-written CUDA only) The decision for #690 is made: kernels are not generated from the pyccel source. Remove pic/tests/codegen_spike and shorten the decision section in CUDA_STRATEGY.md. The hand-written port of bstar_parallel_3form and its parity cases stay. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 108 +---- .../pic/tests/codegen_spike/__init__.py | 1 - .../pic/tests/codegen_spike/compare.py | 61 --- .../pic/tests/codegen_spike/numba_codegen.py | 458 ------------------ .../tests/codegen_spike/test_numba_codegen.py | 29 -- 5 files changed, 26 insertions(+), 631 deletions(-) delete mode 100644 src/struphy/pic/tests/codegen_spike/__init__.py delete mode 100644 src/struphy/pic/tests/codegen_spike/compare.py delete mode 100644 src/struphy/pic/tests/codegen_spike/numba_codegen.py delete mode 100644 src/struphy/pic/tests/codegen_spike/test_numba_codegen.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 8988497d1..775e97382 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -27,7 +27,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). -- [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)). +- [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -316,7 +316,7 @@ algorithm is `discrete_gradient_1st_order`; its chain comes first. | 16–17 | `push_gc_bxEstar_discrete_gradient_2nd_order`, `push_gc_Bstar_discrete_gradient_2nd_order` | P | | 18–19 | `push_gc_bxEstar_explicit_multistage`, `push_gc_Bstar_explicit_multistage` | P | -These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. +These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. **Step 4 – Hybrid MHD–kinetic 6D (current and pressure coupling)** @@ -367,7 +367,7 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. -- **Single-source alternatives.** Decided for now (#690): hand-written CUDA stays the default; see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690) for the comparison and for what would change the decision. +- **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). - **Polar splines on the GPU** (left after PR 19): spline mappings run on the device, but `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels). Needed @@ -402,90 +402,34 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat children died with SIGILL and no output), with `faulthandler` and one OpenMP thread, and a failure reports the signal and the end of stdout and stderr (gvec writes its Fortran messages to stdout). -## Hand-written CUDA vs. code generation (decision for #690) +## Hand-written CUDA (decision for #690) -**Decision: the guiding-center and hybrid kernels (steps 3–5 of the [porting order](#porting-order)) are ported by -hand, as `_cuda.cu` next to the pyccel kernel, like every kernel so far.** None of the generators is adopted. -This is decided again if a condition under [What would change the decision](#what-would-change-the-decision) -comes true. +**Decision: all CUDA kernels are written by hand, as `_cuda.cu` next to the pyccel kernel, like every kernel +so far. Kernels are not generated from the pyccel source (no numba, `cupyx.jit` or other code generation).** To decide, one guiding-center kernel, `bstar_parallel_3form` (step 3, #7: drift-kinetic B*∥ as a 3-form; mapping -Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported in two ways: - -- **By hand:** `pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel, with - 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping (all analytic and four spline mappings), - with α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)`, a negative `eta_n`, a hole - and two output columns. This port stays in the code. -- **Generated:** a prototype generator, `pic/tests/codegen_spike/numba_codegen.py`, that translates the pyccel - source of the kernel *and of every helper it calls* into one numba.cuda module (`compare.py` runs it, - `test_numba_codegen.py` tests it; skipped without numba). It only exists for this comparison and can be - deleted. There is no `kernels` in the module names, so `struphy compile` does not pick it up. - -### Results - -| | Hand-written CUDA | Generated (numba.cuda) | -|---|---|---| -| Code for this kernel | 75 lines `.cu` (36 without comments), plus one new 5-line helper `eval_0form_spline_mpi` in `bsplines/evaluation_kernels_3d.cuh`; everything else is reused (`df` for 14 mappings, `det`, `get_spans`, `eval_spline_mpi_kernel`) | 0 lines by hand, but a 458-line generator. The generated module has 490 lines and contains its own copy of the whole helper chain (24 device functions); every generated kernel repeats it | -| Agreement with pyccel (14 cases) | CPU emulation of the real `.cu` source: all agree, max abs. difference 5.7·10⁻¹⁴. Removing the `mod` from the CUDA kernel makes the emulation test fail (115 of 3225 values) | numba.cuda **simulator**: all agree, max 6.3·10⁻¹³; the same translation compiled with `numba.njit`: the same values. The simulator interprets the Python one thread at a time and compiles nothing | -| Compiles for a GPU? | Not checked here (no GPU); the same NVRTC path as the 15 kernels before it | Not checked (numba.cuda needs NVVM from the CUDA toolkit, which macOS does not have). Not checked either: named tuples of local arrays in device code, and the 38 flat kernel parameters | -| Argument classes | Takes `MarkerArgs`, `DomainArgs`, `DerhamArgs` by value, like every kernel; cunumpy's `CudaKernel` checks the call against the signature | numba does not accept the C structs. The generator turns each argument object into one kernel parameter per struct field (38 parameters) and rebuilds the objects as named tuples inside the kernel. Calling it needs a second launch path next to cunumpy's (flattening, numba's array conversion, its own compile cache) | -| 1:1 with pyccel | Names, arguments and order as in pyccel, kept by review | Exactly 1:1 for the parts the generator understands; whatever it does not understand fails with `NotImplementedError` (see the last row) | -| Parity tests, emulation | The existing tests run unchanged: `PARITY_CASES` entry, GPU parity with `assert_kernels_agree`, CPU emulation, which compiles the `.cu` itself | The parity cases are reused, but the simulator checks the Python, not the code the GPU will run. cunumpy's emulation and `assert_kernels_agree` do not apply to numba kernels | -| Debugging | A C++ compiler and printf in emulation; `-lineinfo`, cunumpy's debug mode, compute-sanitizer and Nsight on the GPU; the source you read is the source you run | pdb in the simulator, where Python float arithmetic hides bugs specific to the GPU. On the GPU, errors point into generated code. The generator's own mistakes are a new kind of bug (see below) | -| Expected performance (no GPU numbers possible here) | Per-thread arrays with fixed size (`MAX_SPLINE_DEGREE`), `__device__ inline` helpers; atomics where needed (`cunumpy_atomic_add`) | About the same local memory (the same `MAX_LOCAL` bounds; `b_der_splines_slim` has the same 9×9 `values`), plus loops that fill `zeros` arrays. numba applies Python's wrap-around for negative indices to every array access (one extra compare/select each), passes shape, strides and data pointer per array, and has a higher launch overhead than `RawKernel`. The generator cannot tell which writes need atomics: a translated accumulation kernel would race | -| Dependencies | None beyond CuPy | numba + llvmlite + numba-cuda (≈175 MB installed), and numba only supports a range of NumPy versions (numba 0.68 was needed for NumPy 2.4). The numba.cuda that comes with numba is deprecated in favour of NVIDIA's numba-cuda | -| The other 59 pusher and accumulation kernels | Each is a transcription onto existing helpers; new helpers as needed (e.g. 1-/2-form evaluations exist, `eval_0form_spline_mpi` was added here) | Tried on all 60 entry kernels (`.claude/scripts/codegen_probe.py`, not committed): 31 translate, 29 do not (2D slice assignments such as `c[:, :] = ...` in the linear-algebra helpers, `numpy.abs` in the SPH helpers). "Translates" does not mean "is correct": none of them ran, and the accumulation kernels would race | - -What the generator had to change in the pyccel code, and the risks that come with it: - -- `empty`/`zeros` with a size known only at run time (`zeros(pn)`) become per-thread arrays of `MAX_LOCAL` - entries. This silently changes `len(x)`/`x.shape` of these arrays, e.g. `diff[:] = diff * pn` in - `b_der_splines_slim` now loops over 9 entries instead of `pn` (harmless there, because the extra entries are - zero). -- Array expressions on slices (`eta[:] = alpha[:3] * eta_k + ...`) become loops over the elements, because numba.cuda - cannot allocate the temporaries. Expressions with 2D slices are not supported yet. -- `numpy.mod` has to be translated exactly as pyccel compiles it (Fortran `MODULO`, `a - floor(a / b) * b`). -- The scratch arrays of `DerhamArguments` (`bn1`, ..., `bd3`), which every marker shares in pyccel, have to become - per-thread arrays; the hand-written `SplineScratch` does the same. -- The marker loop becomes one thread per marker; `continue` becomes `return`. - -**The other generators.** `cupyx.jit` (CuPy's experimental Python-to-CUDA transpiler) would need the same -translation of the pyccel source: its reference lists elementary scalar operations, thread indices, shared memory -and atomics, but no per-thread arrays and no struct arguments, and "CuPy's array operations ... are not supported". -It adds no dependency, but it was not run here (no CuPy). **pyccel 2.2.1 has no CUDA backend** (`--language` is -Fortran, C or Python). The CUDA work happens in the separate fork [pyccel-cuda](https://github.com/pyccel/pyccel-cuda) -(`@kernel`/`@device`, last commit January 2025); its open issues still include thread indexing, passing arrays as -arguments and device arrays, and nothing there handles class arguments like `MarkerArguments`. - -### Reasoning +Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported by hand +(`pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel) and through a +numba.cuda prototype generator. The prototype was removed after the decision. The reasons: 1. **The shared part is already written.** What a generator would save is mostly the helper chain (mappings, - metric, splines, linear algebra, boundary conditions), and PRs 10–19 have already ported it by hand and tested - it, every mapping included. What is left per kernel is its body: here 36 lines that read like the pyccel code. -2. **A generator is a second compiler to maintain.** The prototype has 458 lines for a subset of what the kernels - use, and its mistakes (wrong array lengths, missing atomics) are silent. pyccel already is the compiler for the - CPU; a second one written by ourselves only pays off for many more kernels than ~45. -3. **It does not fit the infrastructure.** Everything since PR 8 (structs, signature checks in `CudaKernel`, - compile cache, debug mode, `assert_kernels_agree`, CPU emulation of the real source, no host transfers) is - built for CUDA C++ source. numba kernels would need a second launch and test path in cunumpy and struphy. -4. **It is harder to test.** The CPU emulation compiles the `.cu` that will run on the GPU; the numba simulator - only interprets Python. The hand-written port is gated better without a GPU than the generated one. -5. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change - to a pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. - -### What would change the decision - -- **pyccel gets a CUDA backend** in a release (not the fork) that accepts our argument classes and arrays. Then the - `.cu` files could be generated by pyccel in place, behind the same `Kernel`, structs and tests; the parity cases - would check them unchanged. This is the most attractive single-source route; check pyccel's releases before - step 4. -- **Porting by hand becomes the bottleneck**, e.g. new models add many kernels, or the pyccel kernels change often. - Then the next step is a generator that writes **`.cu` text** (C++) instead of numba: the AST transformations of - the prototype (helper calls, local arrays, slice loops, marker loop to threads, `mod`) carry over, and the output - goes through the existing structs, `CudaKernel`, emulation and parity tests. Accumulation kernels would still need - their atomics marked by hand. -- **GPU measurements** that show a generated kernel to be as fast as the hand-written one would remove the - performance argument; they do not touch points 2–4. + metric, splines, linear algebra, boundary conditions), which PRs 10–19 have already ported and tested by hand. + What is left per kernel is its body: for `bstar_parallel_3form` 36 lines that read like the pyccel code, plus one + 5-line helper (`eval_0form_spline_mpi`). +2. **A generator is a second compiler to maintain**, and its mistakes are silent (array lengths that change when a + size is known only at run time; it cannot tell which writes need atomics, so accumulation kernels would race). + The prototype translated only 31 of the 60 pusher and accumulation entry kernels. +3. **It does not fit the infrastructure.** Structs, signature checks in `CudaKernel`, the compile cache, debug + mode, `assert_kernels_agree`, the CPU emulation of the real source and the transfer checks are all built for + CUDA C++ source. numba does not accept the C structs and would need a second launch and test path, plus about + 175 MB of dependencies. +4. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change to a + pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. + +The port of `bstar_parallel_3form` stays, with 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping +(all analytic and four spline mappings), α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)` +(Fortran `MODULO`, `eta - floor(eta)` in CUDA), a negative `eta_n`, a hole and two output columns. The CPU emulation +agrees with pyccel to 5.7·10⁻¹⁴. ## PR 10 implementation notes diff --git a/src/struphy/pic/tests/codegen_spike/__init__.py b/src/struphy/pic/tests/codegen_spike/__init__.py deleted file mode 100644 index d1b189c74..000000000 --- a/src/struphy/pic/tests/codegen_spike/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Prototype for issue #690 (hand-written CUDA vs. code generation); not used by struphy, see ``numba_codegen``.""" diff --git a/src/struphy/pic/tests/codegen_spike/compare.py b/src/struphy/pic/tests/codegen_spike/compare.py deleted file mode 100644 index ae72c1a6f..000000000 --- a/src/struphy/pic/tests/codegen_spike/compare.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Run the numba kernel generated from ``bstar_parallel_3form`` on its parity cases and compare it with pyccel. - -Usage (numba is not a struphy dependency):: - - NUMBA_ENABLE_CUDASIM=1 python -m struphy.pic.tests.codegen_spike.compare cuda # numba.cuda simulator - python -m struphy.pic.tests.codegen_spike.compare cpu # same translation, numba.njit - -The simulator interprets the generated Python thread by thread (it does not compile anything); ``cpu`` compiles the -same translation in nopython mode. Prints one line per case and a summary; exits with 1 if a case disagrees. -""" - -import sys -import time - -import cunumpy as xp -import numpy as np - -from struphy.pic.tests.codegen_spike.numba_codegen import flatten, load - -MODULE = "struphy.pic.pushing.kernels.bstar_parallel_3form.bstar_parallel_3form_kernels" -NAME = "bstar_parallel_3form" -BLOCK = 128 - - -def compare(target): - from struphy.pic.pushing.kernels.bstar_parallel_3form import bstar_parallel_3form - from struphy.pic.tests.cuda_parity_cases import PARITY_CASES - from struphy.pic.tests.test_cuda_emulation import arrays - - kernel, path = load(MODULE, NAME, target) - print(f"generated {path} ({len(path.read_text().splitlines())} lines)") - spec = PARITY_CASES[NAME] - failures, bitwise = 0, 0 - with xp.use_backend("numpy"): - for case in spec.cases: - reference_args, generated_args = spec.build(case), spec.build(case) - bstar_parallel_3form(*reference_args) - flat = flatten(generated_args) - start = time.perf_counter() - if target == "cuda": - n_markers = generated_args[2].n_markers - kernel[(n_markers + BLOCK - 1) // BLOCK, BLOCK](*flat) - else: - kernel(*flat) - seconds = time.perf_counter() - start - reference, generated = arrays(reference_args), arrays(generated_args) - diff = max( - float(np.max(np.abs(r.astype(float) - g.astype(float)), initial=0.0)) - for r, g in zip(reference, generated) - ) - same = all(np.array_equal(r, g, equal_nan=True) for r, g in zip(reference, generated)) - ok = all(np.allclose(g, r, rtol=spec.rtol, atol=spec.atol) for r, g in zip(reference, generated)) - failures += not ok - bitwise += same - print(f"domain {case[0]:2d}: max |diff| = {diff:.2e}, bitwise equal: {same}, {seconds:.2f} s") - print(f"{target}: {len(spec.cases) - failures}/{len(spec.cases)} cases agree, {bitwise} bitwise") - return failures - - -if __name__ == "__main__": - sys.exit(1 if compare(sys.argv[1] if len(sys.argv) > 1 else "cpu") else 0) diff --git a/src/struphy/pic/tests/codegen_spike/numba_codegen.py b/src/struphy/pic/tests/codegen_spike/numba_codegen.py deleted file mode 100644 index 062df877d..000000000 --- a/src/struphy/pic/tests/codegen_spike/numba_codegen.py +++ /dev/null @@ -1,458 +0,0 @@ -"""Prototype for issue #690: generate a ``numba.cuda`` kernel from the Python source of a pyccel kernel. - -Not used by struphy. It exists to compare code generation with the hand-written ``_cuda.cu`` files on one -real kernel (``bstar_parallel_3form``); the result of the comparison is in ``CUDA_STRATEGY.md`` (decision for -#690). The module name contains no ``kernels``, so ``struphy compile`` ignores it. - -The generator works on the source files (the imported kernel modules are usually the compiled pyccel extensions): - -- it parses the entry kernel's module, follows every call into struphy modules (``import a.b as x`` and - ``from a.b import f``) and translates each reached function once, as a numba device function named - ``__``; -- it strips the pyccel decorators and type annotations, maps the NumPy scalar functions to ``math``, and - ``numpy.mod`` to the Fortran ``MODULO`` pyccel generates; -- ``empty``/``zeros`` become ``cuda.local.array`` with a constant size (a run-time size such as ``zeros(pn)`` gets - ``MAX_LOCAL`` entries, like the hand-written ``MAX_SPLINE_DEGREE`` arrays); -- array expressions on slices (``eta[:] = alpha[:3] * eta_k + ...``, ``bn[:] = 0.0``) become element loops, since - numba.cuda cannot allocate the temporaries; -- the entry kernel's marker loop ``for ip in range(n): ...`` becomes ``ip = cuda.grid(1)`` and ``continue`` becomes - ``return``; -- argument objects (``MarkerArguments``, ``DomainArguments``, ``DerhamArguments``) are passed as the fields of - their CUDA structs (``CudaMarkerArguments.struct.fields``, one kernel parameter each) and rebuilt as named tuples - at the top of the kernel; ``DerhamArguments`` gets its scratch arrays ``bn1``, ..., ``bd3`` as per-thread local - arrays there, so the helpers can stay one to one with pyccel. - -With ``target="cpu"`` the same translation is compiled with ``numba.njit`` (NumPy allocations, the marker loop kept), -which checks that the generated code types in nopython mode; numba.cuda itself needs a GPU (or NVVM) to compile, and -its simulator (``NUMBA_ENABLE_CUDASIM=1``) only interprets the Python. -""" - -import ast -import importlib.util -import re -import sys -import tempfile -from pathlib import Path - -from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments - -MAX_LOCAL = 9 # MAX_SPLINE_DEGREE + 1 of bsplines_kernels.cuh - -STRUCT_CLASSES = { - "MarkerArguments": CudaMarkerArguments, - "DerhamArguments": CudaDerhamArguments, - "DomainArguments": CudaDomainArguments, -} -# pyccel scratch arrays of the argument classes: per-thread local arrays in the generated kernel -SCRATCH = {"DerhamArguments": ("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")} - -NUMPY_MATH = { - "cos": "cos", - "sin": "sin", - "tan": "tan", - "arctan": "atan", - "arcsin": "asin", - "arccos": "acos", - "arctan2": "atan2", - "sqrt": "sqrt", - "exp": "exp", - "log": "log", - "floor": "floor", -} -NUMPY_ALLOC = ("empty", "zeros") - - -def _source_file(module): - """The ``.py`` file of `module` (the import system would find the compiled pyccel extension first).""" - spec = importlib.util.find_spec(module) - return Path(spec.origin).parent / (module.rsplit(".", 1)[-1] + ".py") - - -class _Module: - """Functions and imports of one source module.""" - - def __init__(self, name): - self.name = name - self.short = name.rsplit(".", 1)[-1] - tree = ast.parse(_source_file(name).read_text()) - self.functions = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)} - self.module_aliases = {} # alias -> module - self.imported = {} # name -> (module, name) - for node in tree.body: - if isinstance(node, ast.Import): - for a in node.names: - self.module_aliases[a.asname or a.name] = a.name - elif isinstance(node, ast.ImportFrom): - for a in node.names: - self.imported[a.asname or a.name] = (node.module, a.name) - - def resolve(self, func): - """``(module, name)`` of a called function node, or ``None``.""" - if isinstance(func, ast.Name): - if func.id in self.functions: - return self.name, func.id - return self.imported.get(func.id) - if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): - module = self.module_aliases.get(func.value.id) - if module is not None: - return module, func.attr - return None - - -def _array_ndim(annotation): - """Number of dimensions of a pyccel array annotation such as ``"float[:,:]"``, or 0.""" - if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): - match = re.search(r"\[([:,\s]+)\]", annotation.value) - if match: - return match.group(1).count(":") - return 0 - - -def _struct_type(annotation): - if isinstance(annotation, ast.Constant) and annotation.value in STRUCT_CLASSES: - return annotation.value - if isinstance(annotation, ast.Name) and annotation.id in STRUCT_CLASSES: - return annotation.id - return None - - -def _expr(source): - return ast.parse(source, mode="eval").body - - -def _stmt(source): - return ast.parse(source).body - - -class _Elementwise(ast.NodeTransformer): - """Index a slice expression at element ``index``: ``a[lo:hi]`` -> ``a[lo + index]``, 1D array ``x`` -> ``x[index]``.""" - - def __init__(self, index, arrays): - self.index = index - self.arrays = arrays - - def _element(self, slc): - lower = slc.lower if slc.lower is not None else ast.Constant(0) - return ast.BinOp(lower, ast.Add(), ast.Name(self.index, ast.Load())) - - def visit_Subscript(self, node): - if isinstance(node.slice, ast.Slice): - return ast.Subscript(node.value, self._element(node.slice), node.ctx) - if isinstance(node.slice, ast.Tuple) and any(isinstance(e, ast.Slice) for e in node.slice.elts): - elts = [self._element(e) if isinstance(e, ast.Slice) else e for e in node.slice.elts] - return ast.Subscript(node.value, ast.Tuple(elts, ast.Load()), node.ctx) - return node - - def visit_Name(self, node): - if self.arrays.get(node.id) == 1: - return ast.Subscript(ast.Name(node.id, ast.Load()), ast.Name(self.index, ast.Load()), node.ctx) - if node.id in self.arrays and self.arrays[node.id] > 1: - raise NotImplementedError(f"array expression on the {self.arrays[node.id]}D array {node.id}") - return node - - -def _slice_length(target): - """Number of elements written by a sliced assignment target (one slice per target).""" - slices = target.slice.elts if isinstance(target.slice, ast.Tuple) else [target.slice] - positions = [i for i, s in enumerate(slices) if isinstance(s, ast.Slice)] - if len(positions) != 1: - raise NotImplementedError(f"assignment to {ast.unparse(target)}: exactly one slice is supported") - slc = slices[positions[0]] - if slc.step is not None: - raise NotImplementedError(f"assignment to {ast.unparse(target)}: slice steps are not supported") - if slc.upper is None: - upper = ast.Subscript(ast.Attribute(target.value, "shape", ast.Load()), ast.Constant(positions[0]), ast.Load()) - else: - upper = slc.upper - if slc.lower is None: - return upper - return ast.BinOp(upper, ast.Sub(), slc.lower) - - -def _has_slice(target): - if not isinstance(target, ast.Subscript): - return False - s = target.slice - return isinstance(s, ast.Slice) or (isinstance(s, ast.Tuple) and any(isinstance(e, ast.Slice) for e in s.elts)) - - -class _ContinueToReturn(ast.NodeTransformer): - """``continue`` of the marker loop -> ``return`` (nested loops keep theirs).""" - - def visit_Continue(self, node): - return ast.Return(None) - - def visit_For(self, node): - return node - - def visit_While(self, node): - return node - - -class Generator: - """Translate a pyccel entry kernel and every struphy function it calls into one numba module.""" - - def __init__(self, target="cuda"): - assert target in ("cuda", "cpu") - self.target = target - self.modules = {} - self.translated = {} # (module, name) -> mangled name - self.sources = [] # translated functions, callees first - self.counter = 0 - - def _module(self, name): - if name not in self.modules: - self.modules[name] = _Module(name) - return self.modules[name] - - def _is_struphy_function(self, ref): - if ref is None or not ref[0].startswith("struphy."): - return False - try: - return ref[1] in self._module(ref[0]).functions - except (ModuleNotFoundError, FileNotFoundError, AttributeError): - return False - - def _numpy_name(self, module, node): - """The NumPy name a call or name node refers to in `module`, or ``None``.""" - if isinstance(node, ast.Name): - ref = module.imported.get(node.id) - if ref is not None and ref[0] == "numpy": - return ref[1] - return None - - def _fresh(self, prefix): - self.counter += 1 - return f"{prefix}{self.counter}" - - # ------------------------------------------------------------------ allocations - - def _allocation(self, module, target, call, arrays): - kind = self._numpy_name(module, call.func) - shape = call.args[0] - ndim = len(shape.elts) if isinstance(shape, ast.Tuple) else 1 - dtype = "int64" - for kw in call.keywords: - if kw.arg == "dtype" and isinstance(kw.value, ast.Name) and kw.value.id == "float": - dtype = "float64" - if len(call.args) > 1 and isinstance(call.args[1], ast.Name) and call.args[1].id == "float": - dtype = "float64" - if not call.keywords and len(call.args) == 1: - dtype = "float64" - arrays[target] = ndim - if self.target == "cpu": - return _stmt(f"{target} = np.{kind}({ast.unparse(shape)}, dtype=np.{dtype})") - elts = shape.elts if isinstance(shape, ast.Tuple) else [shape] - sizes = [e.value if isinstance(e, ast.Constant) else MAX_LOCAL for e in elts] - size = sizes[0] if ndim == 1 else "(" + ", ".join(map(str, sizes)) + ")" - stmts = _stmt(f"{target} = cuda.local.array({size}, {dtype})") - if kind == "zeros": - loops = "".join(f"{' ' * d}for _z{d} in range({s}):\n" for d, s in enumerate(sizes)) - index = ", ".join(f"_z{d}" for d in range(ndim)) - stmts += _stmt(f"{loops}{' ' * ndim}{target}[{index}] = 0") - return stmts - - # ------------------------------------------------------------------ functions - - def translate(self, module_name, name, entry=False): - """Translate function `name` of `module_name`; returns its name in the generated module.""" - key = (module_name, name) - if key in self.translated: - return self.translated[key] - module = self._module(module_name) - mangled = f"{module.short}__{name}" - self.translated[key] = mangled - node = ast.parse(ast.unparse(module.functions[name])).body[0] # a fresh copy - - # parameters: array dimensions and argument objects from the pyccel annotations - arrays = {a.arg: _array_ndim(a.annotation) for a in node.args.args if _array_ndim(a.annotation)} - structs = {a.arg: _struct_type(a.annotation) for a in node.args.args if _struct_type(a.annotation)} - for a in node.args.args: - a.annotation = None - node.returns = None - node.decorator_list = [] - node.name = mangled - if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant): - node.body = node.body[1:] or [ast.Pass()] - - node.body = self._statements(module, node.body, arrays) - - if entry: - self._make_entry(node, structs) - decorator = "@cuda.jit" if self.target == "cuda" else "@njit" - else: - decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" - self.sources.append(f"{decorator}\n{ast.unparse(ast.fix_missing_locations(node))}\n") - return mangled - - def _statements(self, module, body, arrays): - out = [] - for stmt in body: - out += self._statement(module, stmt, arrays) - return out - - def _statement(self, module, stmt, arrays): - # local arrays - if ( - isinstance(stmt, ast.Assign) - and isinstance(stmt.value, ast.Call) - and self._numpy_name(module, stmt.value.func) in NUMPY_ALLOC - ): - (target,) = stmt.targets - return self._allocation(module, target.id, stmt.value, arrays) - # array expressions on slices -> element loops - if isinstance(stmt, (ast.Assign, ast.AugAssign)): - target = stmt.targets[0] if isinstance(stmt, ast.Assign) else stmt.target - if _has_slice(target): - index = self._fresh("_k") - length = _slice_length(target) - element = _Elementwise(index, arrays) - value = self._expressions(module, element.visit(stmt.value)) - new_target = element.visit(target) - if isinstance(stmt, ast.Assign): - inner = ast.Assign([new_target], value) - else: - inner = ast.AugAssign(new_target, stmt.op, value) - loop = ast.For( - ast.Name(index, ast.Store()), - ast.Call(ast.Name("range", ast.Load()), [length], []), - [inner], - [], - lineno=stmt.lineno, - ) - return [loop] - # compound statements: recurse into their bodies - for field in ("body", "orelse"): - if hasattr(stmt, field) and isinstance(getattr(stmt, field), list): - setattr(stmt, field, self._statements(module, getattr(stmt, field), arrays)) - for field in ("value", "test", "iter"): - if hasattr(stmt, field) and isinstance(getattr(stmt, field), ast.AST): - setattr(stmt, field, self._expressions(module, getattr(stmt, field))) - if isinstance(stmt, ast.Assign): - stmt.targets = [self._expressions(module, t) for t in stmt.targets] - if isinstance(stmt, ast.AugAssign): - stmt.target = self._expressions(module, stmt.target) - return [stmt] - - def _expressions(self, module, node): - generator = self - - class Calls(ast.NodeTransformer): - def visit_Call(self, call): - self.generic_visit(call) - ref = module.resolve(call.func) - numpy_name = generator._numpy_name(module, call.func) - if generator._is_struphy_function(ref): - call.func = ast.Name(generator.translate(*ref), ast.Load()) - elif numpy_name in NUMPY_MATH: - call.func = _expr(f"math.{NUMPY_MATH[numpy_name]}") - elif numpy_name == "mod": - call.func = ast.Name("_numpy_mod", ast.Load()) - elif numpy_name == "shape": - return ast.Attribute(call.args[0], "shape", ast.Load()) - elif numpy_name is not None: - raise NotImplementedError(f"numpy.{numpy_name} in {module.name}") - return call - - def visit_Name(self, name): - if generator._numpy_name(module, name) == "pi": - return _expr("math.pi") - return name - - return Calls().visit(node) - - def _make_entry(self, node, structs): - """Flatten the argument objects into struct fields, rebuild them, and map the marker loop to threads.""" - params, prologue = [], [] - for a in node.args.args: - kind = structs.get(a.arg) - if kind is None: - params.append(a.arg) - continue - fields = [f.name for f in STRUCT_CLASSES[kind].struct.fields] - params += [f"{a.arg}__{f}" for f in fields] - values = [f"{a.arg}__{f}" for f in fields] - if self.target == "cuda": - values += [f"cuda.local.array({MAX_LOCAL}, float64)" for _ in SCRATCH.get(kind, ())] - else: - values += [f"np.empty({MAX_LOCAL})" for _ in SCRATCH.get(kind, ())] - prologue += _stmt(f"{a.arg} = {kind}_T({', '.join(values)})") - node.args.args = [ast.arg(p) for p in params] - - body = prologue + node.body - if self.target == "cuda": - for i, stmt in enumerate(body): - if ( - isinstance(stmt, ast.For) - and isinstance(stmt.iter, ast.Call) - and getattr(stmt.iter.func, "id", None) == "range" - and len(stmt.iter.args) == 1 - ): - ip, n = stmt.target.id, ast.unparse(stmt.iter.args[0]) - loop_body = [_ContinueToReturn().visit(s) for s in stmt.body] - body = body[:i] + _stmt(f"{ip} = cuda.grid(1)\nif {ip} >= {n}:\n return") + loop_body - break - else: - raise NotImplementedError("no marker loop 'for ip in range(n)' in the entry kernel") - node.body = body - - def module_source(self): - decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" - header = [ - f"# Generated by struphy.pic.tests.codegen_spike.numba_codegen (target={self.target}); do not edit.", - "import math", - "from collections import namedtuple", - "", - "import numpy as np", - "from numba import cuda, float64, int64, njit", - "", - ] - for kind, cls in STRUCT_CLASSES.items(): - names = [f.name for f in cls.struct.fields] + list(SCRATCH.get(kind, ())) - header.append(f"{kind}_T = namedtuple({kind + '_T'!r}, {names!r})") - header += [ - "", - "", - decorator, - "def _numpy_mod(a, b):", - " # numpy.mod as pyccel compiles it (Fortran MODULO)", - " return a - math.floor(a / b) * b", - "", - "", - ] - return "\n".join(header) + "\n\n".join(self.sources) - - -def generate(module_name, name, target="cuda"): - """Source of the numba module for the pyccel kernel `name` in `module_name` (and every function it calls).""" - generator = Generator(target) - generator.translate(module_name, name, entry=True) - return generator.module_source() - - -def load(module_name, name, target="cuda", directory=None): - """Generate, write to `directory` (a temporary one by default) and import; returns the kernel function.""" - source = generate(module_name, name, target) - directory = Path(directory or tempfile.mkdtemp(prefix="struphy_numba_codegen_")) - stem = f"generated_{name}_{target}" - path = directory / f"{stem}.py" - path.write_text(source) - spec = importlib.util.spec_from_file_location(stem, path) - module = importlib.util.module_from_spec(spec) - sys.modules[stem] = module - spec.loader.exec_module(module) - return getattr(module, f"{module_name.rsplit('.', 1)[-1]}__{name}"), path - - -def flatten(args): - """Kernel arguments with every argument object replaced by the fields of its CUDA struct (host arrays).""" - from struphy.pic.tests.cuda_emulation import _struct_field - - flat = [] - for value in args: - kind = type(value).__name__ - if kind in STRUCT_CLASSES: - flat += [_struct_field(value, f.name) for f in STRUCT_CLASSES[kind].struct.fields] - else: - flat.append(value) - return flat diff --git a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py deleted file mode 100644 index 7648440b6..000000000 --- a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py +++ /dev/null @@ -1,29 +0,0 @@ -"""The numba kernel generated from ``bstar_parallel_3form`` agrees with pyccel (issue #690 prototype). - -Skipped unless numba (with numba-cuda for the simulator) is installed; numba is not a struphy dependency. -Runs in a subprocess because ``NUMBA_ENABLE_CUDASIM`` must be set before numba is imported. -""" - -import importlib.util -import subprocess -import sys - -import pytest - -from struphy.geometry.tests.test_domain import serial_child_env - -requires_numba = pytest.mark.skipif(importlib.util.find_spec("numba") is None, reason="numba is not installed") - - -@requires_numba -@pytest.mark.parametrize("target", ["cpu", "cuda"]) -def test_generated_kernel_agrees_with_pyccel(target): - """``cpu``: the translation compiled with numba.njit; ``cuda``: the numba.cuda kernel in the CUDA simulator.""" - extra = {"NUMBA_ENABLE_CUDASIM": "1"} if target == "cuda" else {} - result = subprocess.run( - [sys.executable, "-m", "struphy.pic.tests.codegen_spike.compare", target], - env=serial_child_env(**extra), - capture_output=True, - text=True, - ) - assert result.returncode == 0, result.stdout + result.stderr From 703a94717ab344da55617a04165cc37efe39e1d5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:13:19 +0200 Subject: [PATCH 3/3] Rename first_init_idx to first_pusher_idx in the code new on this branch Co-Authored-By: Claude Opus 5.5 --- .../bstar_parallel_3form/bstar_parallel_3form_cuda.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu index ec7389284..7a473ed1b 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -32,7 +32,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long double eta_k[3], eta_n[3], eta[3], dfm[9]; // get marker arguments - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; // only do something if particle is a "true" particle (i.e. not a hole) @@ -40,7 +40,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long for (int j = 0; j < 3; ++j) { eta_k[j] = args_markers.markers(ip, j) + args_markers.markers(ip, first_shift_idx + j); - eta_n[j] = args_markers.markers(ip, first_init_idx + j); + eta_n[j] = args_markers.markers(ip, first_pusher_idx + j); eta[j] = alpha[j] * eta_k[j] + (1.0 - alpha[j]) * eta_n[j]; // numpy.mod(eta, 1.0) (Fortran MODULO): the result has the sign of the divisor @@ -48,7 +48,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long } double v_k = args_markers.markers(ip, 3); - double v_n = args_markers.markers(ip, first_init_idx + 3); + double v_n = args_markers.markers(ip, first_pusher_idx + 3); double v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n; // evaluate Jacobian, result in dfm