Put proj_on_basis on_gpu

[Maxwell-Rosen]

Closes #27 (wow such an old issue)

My simulations feature large grids (8x400x32x64), and the initial condition projection gets very slow (on the order of minutes). This prompted me to investigate why, and I found that proj_on_basis is entirely a CPU operation. I gave Claude the problem, and I had it write a plan first. Here is what my commit says, which I'm very grateful for, recording detailed commit messages because this was a while ago

Add GPU support for projection on basis in gyrokinetic app

• Introduced device-callable function pointers for evaluation and computational-to-physical (c2p) mappings in the projection on basis functionality.
• Enhanced gkyl_proj_on_basis structure to include GPU-related fields and logic for managing device memory.
• Implemented CUDA kernels for evaluating projections on the GPU, including memory management for device arrays.
• Updated the gk_species_projection functions to handle projections on both host and device, allowing for seamless integration of GPU acceleration.
• Added device-specific implementations for position and velocity mappings to ensure compatibility with GPU execution.
• Refactored existing code to support device memory allocation and copying, ensuring efficient data handling between host and device.
• Introduced new private headers for managing device contexts related to c2p mappings, ensuring encapsulation and clarity in the codebase.

Plan: proj_on_basis on GPUs and GPU projection in gk_species_projection_calc_proj_func

Status: implemented, NOT tested on a GPU (no GPU available on the development
machine). CPU builds compile and pass existing tests. All CUDA code follows
established in-repo patterns; the "Testing on a GPU machine" section below lists
exactly what to run once a GPU is available.

Goal

Port core/zero/proj_on_basis.c to GPUs so projections of user-specified
functions (antenna sources, BGK operators, time-dependent sources, etc.) can run
directly on the device, instead of projecting on the host and copying to the GPU.
Then use it to perform the proj_on_basis update in
gk_species_projection_calc_proj_func (gyrokinetic app, GKYL_PROJ_FUNC
projections) on the GPU.

Key design facts (established from the codebase)

1. Only .cu files (and the generated-kernel directories) are compiled as
   CUDA. With CC=nvcc, NVCC_FLAGS = -x cu -dc -rdc=true is applied to
   %.cu rules and to the kernel-directory %.c rules; every other .c
   file (zero/, apps/, unit/, creg/) is compiled by nvcc as host C. In
   those host TUs __NVCC__ is still defined, so gkyl_util.h includes
   cuda_runtime.h (making GKYL_HAVE_CUDA and the host CUDA API
   available) and GKYL_CU_DH/GKYL_CU_D expand to attributes the host
   compiler ignores. Consequences:

   • GKYL_CU_DH static inline functions in headers are safe everywhere:
     plain inline functions in host TUs, device-callable when the header is
     included from a .cu TU.
   • __global__ kernels, <<<>>> launches, and taking a device function
     address can only appear in .cu files. Device functions used by GPU
     projections must therefore live in headers (as GKYL_CU_DH inlines) or
     in .cu translation units — not in plain .c app/input files.
   • Helpers like gkyl_range_idx, gkyl_sub_range_inv_idx,
     gkyl_rect_grid_cell_center, gkyl_array_(c)fetch are header inlines,
     already usable in kernels.

2. Device function pointers cannot be produced by host code. Taking &func
   on the host yields the host address. The repo's existing pattern (see
   dg_lbo_gyrokinetic_drag_cu.cu: "Doing function pointer stuff in here avoids
   troublesome cudaMemcpyFromSymbol") is to assign function pointers inside a
   __global__ kernel, where the function name resolves to its device address.
   For user-supplied functions the library cannot name the function at compile
   time, so we provide a macro (GKYL_DEFINE_CU_DEV_FUNC_GETTER, in
   gkyl_util.h, available only under __CUDACC__) that is instantiated in a
   .cu file next to the function; it defines a tiny setter kernel plus a host
   getter that returns the device address (expose it to C code with an
   extern "C" wrapper).

3. The context must live on the GPU. The library cannot deep-copy an opaque
   void *ctx; the user must place their context in device memory
   (gkyl_cu_malloc + gkyl_cu_memcpy) and pass that pointer.

4. Updater device clones. The standard pattern (e.g. velocity_map_cu.cu)
   is: fill the host struct, create device mirrors of its arrays, make a
   temporary host copy of the struct whose pointers are replaced by
   array->on_dev pointers, gkyl_cu_memcpy it into a gkyl_cu_malloced clone,
   and store the clone as up->on_dev (with on_dev = up itself on CPU).

Changes

1. Core: gkyl_proj_on_basis GPU support

• core/zero/gkyl_proj_on_basis.h

  • struct gkyl_proj_on_basis_inp gains bool use_gpu. When true:
    • eval (and c2p_func, if provided) must be device function pointers
      (obtained via GKYL_DEFINE_CU_DEV_FUNC_GETTER).
    • ctx / c2p_func_ctx must point to device-resident memory.
    • c2p_func == 0 still means the identity map (set up on the device
      internally).
  • gkyl_proj_on_basis_new() keeps its signature (CPU only,
    use_gpu = false).

• core/zero/gkyl_proj_on_basis_priv.h (new)

  • Holds the struct gkyl_proj_on_basis definition (moved out of the .c so
    the .cu file can see it), the log_to_comp helper (now GKYL_CU_DH), and
    declarations of the CUDA-side functions
    (gkyl_proj_on_basis_cu_dev_new, gkyl_proj_on_basis_advance_cu).
  • Struct gains: bool use_gpu, device mirrors
    ordinates_cu, weights_cu, basis_at_ords_cu, and
    struct gkyl_proj_on_basis *on_dev (self on CPU).
  • The host-side ordinates/weights/basis_at_ords arrays are retained in
    all cases so gkyl_proj_on_basis_quad() and
    gkyl_proj_on_basis_fetch_ordinate() keep working on the host.

• core/zero/proj_on_basis.c

  • _inew: quadrature setup unchanged (computed on the host); when
    use_gpu, create the *_cu device mirrors (copy host data over) and call
    gkyl_proj_on_basis_cu_dev_new() to build the device clone. The identity
    c2p on GPU is installed by a <<<1,1>>> setter kernel that assigns the
    device-side identity function and points its ctx at the clone's own grid
    member.
  • _advance: dispatches to gkyl_proj_on_basis_advance_cu() when use_gpu
    (guarded by #ifdef GKYL_HAVE_CUDA).
  • _release: releases device mirrors and gkyl_cu_frees the clone.

• core/zero/proj_on_basis_cu.cu (new)

  • proj_on_basis_c2p_identity_cu (GKYL_CU_D) + setter kernel.
  • gkyl_proj_on_basis_cu_dev_new(): device clone per the standard pattern.
  • gkyl_proj_on_basis_advance_cu() + kernel: one thread per cell of the
    update range (range->nblocks, range->nthreads launch specifiers, standard
    stride loop with gkyl_sub_range_inv_idx). Each thread mirrors the CPU
    per-cell algorithm exactly (same loop order, hence same FP behavior up to
    FMA contraction): cell center → for each quadrature node: logical→comp
    coords, c2p, eval → accumulate w_q * f_q * psi_k into the output cell.
  • The per-cell function values f_q (num_ret_vals doubles) are staged in a
    device scratch array with one slot per cell, allocated/released per advance
    call — this mirrors the CPU code (which also allocates fun_at_ords per
    call) and avoids any hard cap on num_ret_vals.

• core/zero/gkyl_util.h

  • New macro (under #ifdef __CUDACC__, i.e. .cu files only):

    GKYL_DEFINE_CU_DEV_FUNC_GETTER(func, func_type, getter)

    defines static func_type getter(void) returning the device address of
    func by launching a <<<1,1>>> kernel that stores func (resolved on the
    device) into device memory and copying it back.

  Usage sketch — in a .cu file (e.g. a companion to an input file, or a
  library .cu):

  GKYL_CU_DH static void eval_source(double t, const double *xn, double *fout, void *ctx) { ... }
  GKYL_DEFINE_CU_DEV_FUNC_GETTER(eval_source, evalf_t, eval_source_getter);
  extern "C" evalf_t eval_source_dev_ptr(void) { return eval_source_getter(); }

  and from C code:

  .eval = eval_source_dev_ptr(), .ctx = ctx_on_gpu, .use_gpu = true,

  Because plain .c input files are host TUs even in GPU builds, a C
  regression test that wants GPU projection needs its device eval function in
  such a companion .cu (or in an app/library .cu). This is the "the user
  will have to ensure the context/functions are on the GPU" part of the task.

2. Velocity map: device-callable c2p evaluation

• vlasov/zero/gkyl_velocity_map.h: new GKYL_CU_DH static inline
  function gkyl_velocity_map_eval_c2p_dev(). Identical logic to the host
  gkyl_velocity_map_eval_c2p(), but it uses the vmap array and the
  vmap_basis pointer (instead of the vmap_ho/vmap_basis_ho host copies).
  On the device clone (gvm->on_dev) those members are device pointers
  (velocity_map_cu.cu builds the clone that way), so this function is valid in
  kernels; on CPU builds it also works with the host object.

3. Position map: device copy + device-callable evaluation

• gyrokinetic/zero/gkyl_position_map.h

  • struct gkyl_position_map gains struct gkyl_array *mc2nu_dev (device
    mirror of the mapping's DG coefficients) and
    struct gkyl_position_map *on_dev (self on CPU).
  • gkyl_position_map_eval_mc2nu() moves into the header as
    GKYL_CU_DH static inline (VLAs replaced with GKYL_MAX_CDIM arrays, since
    VLAs are not valid in device code). Given the host object it behaves as
    before; given gpm->on_dev inside a kernel it evaluates on the device
    (the clone's basis holds device eval_expand etc. pointers, and its
    mc2nu points at the device mirror).
  • New gkyl_position_map_make_cu_dev(): creates the mirror + clone. The
    clone's basis member is overwritten with a device-to-device copy from a
    gkyl_cart_modal_{serendip,tensor}_cu_dev_new() basis so its function
    pointers are device addresses (the temporary device basis container is
    released afterwards; the code addresses stay valid). Host-only members
    (maps, ctxs, bmag_ctx, constB_ctx, xpt_ctx) are nulled in the
    clone. No kernels are needed, so this lives in position_map.c under
    #ifdef GKYL_HAVE_CUDA (CPU builds get an assert(false) stub so the
    symbol always links).
  • gkyl_position_map_set_mc2nu() also syncs the device mirror when present.
  • gkyl_position_map_free() releases the mirror and clone.

• gyrokinetic/apps/gyrokinetic.c: after
  gkyl_position_map_set_mc2nu(...) (post-geometry), call
  gkyl_position_map_make_cu_dev(app->position_map) when app->use_gpu, so
  every species' projection can use app->position_map->on_dev.

4. Gyrokinetic app: GKYL_PROJ_FUNC on the GPU

• gyrokinetic/apps/gkyl_gyrokinetic.h: struct gkyl_gyrokinetic_projection
  (the func/ctx_func branch) gains:

  void (*func_on_dev)(double t, const double *xn, double *fout, void *ctx);
  void *ctx_func_on_dev;

  When running on GPUs, the user may set func_on_dev to the device pointer
  of their (device-callable) function and ctx_func_on_dev to a GPU-resident
  context. The projection then runs fully on the device. If func_on_dev is
  not set, the app falls back to the previous behavior (project on host, copy
  to device), so all existing input files keep working unchanged.

• gyrokinetic/apps/gkyl_gyrokinetic_priv.h: struct gk_proj (FUNC branch)
  gains struct gk_proj_on_basis_c2p_func_ctx *proj_on_basis_c2p_ctx_dev, the
  device copy of the c2p context.

• gyrokinetic/apps/gkyl_gk_proj_on_basis_c2p_priv.h (new): holds
  struct gk_proj_on_basis_c2p_func_ctx (moved out of
  gkyl_gyrokinetic_priv.h, which now includes this header) plus the
  declaration of the device-address getter, so the small .cu below does not
  need to compile the whole app private header as C++.

• gyrokinetic/apps/gk_species_projection_cu.cu (new): the device c2p
  function and its address getter:

  GKYL_CU_D static void gk_proj_on_basis_c2p_phase_func_cu(xcomp, xphys, ctx)
    -> gkyl_position_map_eval_mc2nu(ctx->pos_map /* on_dev */, ...)
       gkyl_velocity_map_eval_c2p_dev(ctx->vel_map /* on_dev */, ...)
  extern "C" proj_on_basis_c2p_t gk_species_projection_c2p_phase_func_cu_dev_ptr(void);

• gyrokinetic/apps/gk_species_projection.c

  • gk_species_projection_init (FUNC branch): if app->use_gpu && inp.func_on_dev, build a device-resident c2p context
    {cdim, vdim, pos_map->on_dev, vel_map->on_dev} and create the
    gkyl_proj_on_basis with use_gpu = true; otherwise keep the host
    updater + proj_host staging array. Per the app design rules, the
    conditional is resolved once at init by choosing between two
    projection_calc implementations:
    • gk_species_projection_calc_proj_func — advances directly into f
      (CPU builds, and GPU builds with a device-side updater);
    • gk_species_projection_calc_proj_func_host — advances into
      proj_host and copies to the device (GPU fallback).
      Both then multiply by the bmag (gyrocenter) and velocity-space Jacobians
      on whatever device f lives on (those ops already support GPUs).
  • gk_species_projection_release: release proj_host only if allocated;
    free the device c2p context if allocated.

5. Unit tests

• core/unit/ctest_proj_on_basis_cu.cu (new): the device eval/c2p
  functions and extern "C" getters for their device addresses (unit .cu
  files are globbed into the library in GPU builds, like ctest_array_cu.cu).
• core/unit/ctest_proj_on_basis.c: GPU tests (declared/registered under
  #ifdef GKYL_HAVE_CUDA), each comparing the device projection
  coefficient-by-coefficient (1e-14) against the host projection of the
  matching host function:
  • test_1_cu: 1x p1 projection of x^2, no ctx, identity c2p (set on the
    device internally).
  • test_2d_2c_cu: 2x p2, num_ret_vals = 2, device ctx allocated with
    gkyl_cu_malloc — exercises multi-component projection and device contexts.
  • test_c2p_1d_cu: projection with a nontrivial device c2p mapping vs. the
    CPU projection with the equivalent host c2p — the same machinery the
    gyrokinetic app uses.

Out of scope (follow-ups)

• eval_on_nodes (the nodal evaluator): the same treatment applies —
  mirror basis_at_nodes on the device, clone the updater, evaluate the device
  eval fptr at nodes and do the nodal→modal transform per cell in a kernel.
  Not needed for gk_species_projection_calc_proj_func; deferred to keep this
  change reviewable.
• Neutral species (gk_neut_species_projection.c): same recipe as
  gk_species_projection.c once this lands.
• Other proj_on_basis app call sites (Maxwellian moments projections, etc.):
  these are init-only and can migrate incrementally by supplying device
  fptrs/ctxs.
• Lua layer: Lua-defined functions cannot run on the device; Lua apps keep
  the host+copy path.

Risks / review notes

• cudaFree after kernel launch (scratch array in advance_cu): safe —
  cudaFree synchronizes with prior work on the default stream (and the repo
  frees kernel-visible memory this way elsewhere).
• Device basis embedded by value (position map clone): the D2D copy of
  struct gkyl_basis copies device code addresses; releasing the temporary
  container does not invalidate them.
• FP differences: GPU FMA contraction can make GPU results differ from CPU
  at the last-ulp level; the unit tests compare to 1e-14, matching other GPU
  tests in the repo.
• num_quad: quadrature tables are precomputed on the host and copied;
  no device-side Gauss-Legendre generation needed.

Testing on a GPU machine

# configure with a machine file that sets CC=nvcc, then:
make core-unit -j 8 && ./cuda-build/core/unit/ctest_proj_on_basis
make gyrokinetic -j 8

# any gyrokinetic C regression test with GKYL_PROJ_FUNC ICs, adapted to set
# .func_on_dev/.ctx_func_on_dev, run with -g; compare frame 0 to the CPU run:
./cuda-build/gyrokinetic/creg/rt_gk_sheath_2x2v_p1 -g -s1

compute-sanitizer --tool memcheck --leak-check full ./cuda-build/core/unit/ctest_proj_on_basis

On CPU-only machines: make core-unit && ./build/core/unit/ctest_proj_on_basis
plus valgrind must remain clean (GPU code paths compile out).