GPU Computing

Tarang uses CUDA.jl through a Julia package extension. Load CUDA before creating GPU fields:

using CUDA
using Tarang

@assert CUDA.functional()

Quick start

Choose the architecture on the Distributor; fields and solvers inherit it.

coords = CartesianCoordinates("x", "y")
dist = Distributor(coords; dtype=Float64, device=GPU())

xb = RealFourier(coords["x"]; size=256, bounds=(0.0, 2π))
yb = RealFourier(coords["y"]; size=256, bounds=(0.0, 2π))
domain = Domain(dist, (xb, yb))
u = ScalarField(domain, "u")

u["g"] .= CUDA.rand(Float64, size(u["g"])...)
forward_transform!(u)
backward_transform!(u)

@assert get_grid_data(u) isa CuArray

Use CPU() instead of GPU() to construct a CPU simulation.

Strict GPU execution

A GPU field is never downloaded to run an unavailable CPU implementation. Supported operations remain on the selected device; unsupported operations raise an error. This applies to transforms, resampling, derivatives, solver vectors, MPI communication, and stochastic phase generation.

Host transfers still occur when explicitly requested, for example:

host_data = on_architecture(CPU(), u["g"])

Output, checkpoints, gathers, and explicit architecture conversions are data movement APIs, not computational fallbacks.

Transforms

GPU transforms run on the device for every array size.

set_gpu_fft_mode!(u, :auto)  # default; device-resident
set_gpu_fft_mode!(u, :gpu)   # explicitly require the GPU backend

set_gpu_fft_mode!(u, :cpu) is rejected for GPU fields. The legacy global FFT size threshold applies only to CPU-field preferences and cannot move a GPU field to FFTW.

Current single-GPU transform support is:

Unsupported basis and layout combinations fail before entering the CPU transform chain.

Complete 2D Fourier–Chebyshev path

The single-GPU 2D path supports RealFourier or ComplexFourier paired with ChebyshevT, with either basis first. Scaled Chebyshev grids are transformed at their nodal length and truncated or zero-padded on the device; Fourier scaling used by 3/2-rule nonlinear products also remains device-resident. The focused validation covers transforms, derivatives along both axes, dealiased products, and a nonlinear wall-bounded RK222 InitialValueProblem with tau boundary conditions and CUDA sparse subproblem solves.

Run the strict validation on an NVIDIA node from the repository root:

julia --project=. test/run_gpu_fc_2d.jl

This command prints CUDA device information and fails if CUDA is unavailable, scalar indexing is attempted, a CPU/GPU value comparison fails, or the warmed InitialValueProblem step performs a fresh device allocation.

The current boundary of this validation is deliberate:

• scaled pure-Chebyshev fields remain unsupported on the device;
• multi-rank 2D FC execution is not part of the single-GPU path (use the separately documented distributed interfaces for supported 3D layouts); and
• transform scratch is cached per device and shape for serial use. Run one transforming Julia task at a time on each CUDA device; concurrent same-shape transforms can share scratch and are not supported yet.

2D time stepping and solves

Pure-Fourier GPU IVPs use field-native stepping: explicit Runge-Kutta and the matrix-free multistep field path run entirely on device arrays. When the left-hand side has an implicit diagonal Fourier operator, select a diagonal IMEX scheme — DiagonalIMEX_RK222(), DiagonalIMEX_RK443(), or DiagonalIMEX_SBDF2() — so the operator is applied per mode in spectral space on the device; the operator is read from the equation (Laplacian, hyper- or fractional Laplacian, constant damping, derivatives of the stepped field) or from an attached SpectralLinearOperator. Every other scheme refuses a left-hand-side operator on a single GPU with an error that names these three, because a pure-Fourier GPU solver assembles no global matrix to treat it with. The full per-architecture table is in Time Steppers.

Fourier–Chebyshev IVPs use per-mode coupled subproblems, batched across Fourier modes by default on the device (batched_modes=false opts out). With GPU fields, matsolver=:auto, :gpu, and :hybrid resolve to the concrete CUDA sparse solver; CPU-only solvers are rejected and solver failures are not retried on CPU. A rank-deficient per-mode system — duplicate tau lifts, an under-constrained mode — is a factorization error on the device, not a least-squares fallback.

Most of this is covered without hardware: test/test_gpu_timesteppers_jlarray.jl drives all twenty schemes on the single-GPU dispatch path using JLArray device fields and a CPU stand-in for cuFFT (see Testing).

Linear GPU boundary-value solves select the CUDA sparse solver automatically. Their right-hand side and solution buffers stay on the field's device backend. An explicit solver selection is also supported:

solver = BoundaryValueSolver(problem; matsolver=:cuda_sparse)

GPU NonlinearBoundaryValueProblem and EigenvalueProblem solves are not device-native yet and raise an unsupported operation error.

Multi-GPU execution

Assign one device to each MPI rank before allocating fields:

using MPI

MPI.Init()
comm = MPI.COMM_WORLD
rank = MPI.Comm_rank(comm)
gpu_id = rank % CUDA.ndevices()
CUDA.device!(gpu_id)

arch = GPU(device_id=gpu_id)
coords = CartesianCoordinates("x", "y", "z")
dist = Distributor(coords; comm=comm, dtype=ComplexF64, device=arch)

Distributed GPU communication requires CUDA-aware MPI. Tarang raises an error when device buffers cannot be passed directly; it does not stage them through host memory.

@assert check_cuda_aware_mpi()

Pure complex-Fourier domains transform through the ordinary call. A field on a GPU architecture with more than one rank is given TransposableFieldStorage when it is constructed, and forward_transform! routes it through the explicit transposes automatically:

bases = (
    ComplexFourier(coords["x"]; size=128, bounds=(0.0, 2π)),
    ComplexFourier(coords["y"]; size=128, bounds=(0.0, 2π)),
    ComplexFourier(coords["z"]; size=128, bounds=(0.0, 2π)),
)
field = ScalarField(Domain(dist, bases), "u")
forward_transform!(field)
backward_transform!(field)

The transpose workspace (buffers, counts, and two MPI sub-communicators) is cached on the Distributor and shared by every field of the same shape and element type, so allocating many fields does not allocate many communicators. close(dist) frees them collectively.

Driving the wrapper by hand still works and is what to use if you need the lower-level control:

tf = TransposableField(field)
distributed_forward_transform!(tf)
distributed_backward_transform!(tf)

These require the field to be in the matching layout first (:g for forward, :c for backward); they raise rather than silently reading the stale buffer.

The distributed DCT-I path supports selected three-dimensional Fourier–Chebyshev layouts. It requires at least one Fourier and one Chebyshev axis. RealFourier is allowed only on the first axis and cannot be combined with another Fourier axis. Other layouts, including pure Chebyshev, are rejected.

Custom kernels

Use the public architecture layer instead of importing implementation details from TarangCUDAExt:

using KernelAbstractions: @kernel, @index

@kernel function scale_kernel!(y, x, α)
    i = @index(Global)
    @inbounds y[i] = α * x[i]
end

scale = KernelOperation(scale_kernel!) do y, x, α
    length(y)
end

arch = GPU()
x = ones(arch, Float64, 1024)
y = similar(x)
scale(arch, y, x, 2.0)

KernelOperation accepts an architecture or an array as its first argument; passing arch makes the launch target explicit.

Memory and profiling

Tarang uses CUDA.jl's memory pool. Prefer preallocated work arrays and use CUDA.jl directly for inspection and profiling:

CUDA.memory_status()
CUDA.@profile begin
    forward_transform!(u)
    backward_transform!(u)
end
CUDA.reclaim()

See the GPU API reference for the public architecture, transform, and distributed interfaces.