Kernel architecture and runtime data movement

AtmosTransport separates three concerns: the transport binary is a CPU-side storage format, Adapt.jl moves model data to the chosen backend, and KernelAbstractions.jl provides the CPU/CUDA/Metal implementations of hot loops. Understanding those boundaries is more useful than memorizing a benchmark from one machine.

Select a backend in the run config

[architecture]
use_gpu = true
backend = "cuda"   # "cpu", "cuda", "metal", or "auto"

"auto" tries CUDA and then Metal. "metal" supports Float32; CPU and CUDA support both configured precisions subject to the hardware. Start with the CPU quickstart, then change the backend without changing the physics configuration.

The resolved CPU() or GPU(:cuda|:metal) value is stored on the grid. Model construction adapts state, workspaces, surface sources, and forcing buffers to that backend. The runner verifies that a requested GPU run actually owns GPU arrays instead of silently falling back to CPU storage.

One kernel definition, specialized by types

Hot loops use KernelAbstractions.@kernel. A simplified sweep looks like:

@kernel function sweep!(new_mass, @Const(old_mass), @Const(face_flux))
    i, j, k = @index(Global, NTuple)
    @inbounds new_mass[i, j, k] =
        old_mass[i, j, k] + face_flux[i, j, k] - face_flux[i + 1, j, k]
end

The same Julia method is compiled for the concrete array backend at the call site. Scheme, limiter, mesh, precision, and operator choices are concrete types, so dispatch happens outside the cell loop. Backend-specific setup and synchronization live in the architecture layer rather than in separate copies of each physics algorithm.

Workspaces own reusable scratch arrays

Operators allocate their scratch storage during model construction. For example, the split-sweep advection workspaces contain typed ping-pong arrays; the cubed-sphere workspace owns halo-padded arrays and six-panel tuples. Adapt.adapt preserves the workspace structure while replacing Array storage with CuArray or MtlArray storage.

This design has two practical consequences:

• operator application does not rebuild large scratch arrays every step;

• the array type remains visible to Julia's compiler throughout the call.

If you add a field to a workspace, also add it to that type's Adapt.adapt_structure method and extend the operator's preflight shape checks.

Packed multi-tracer transport

CellState stores tracer mass in a fourth dimension, (horizontal..., level, tracer). The structured and cubed-sphere split-sweep paths process that tracer dimension inside each directional kernel. Air-mass updates and stencil indexing are therefore shared by all tracers in the launch.

For the six directional legs of one advection palindrome, this changes the launch structure from six launches per tracer to six packed launches, before any CFL subcycling. It does not make runtime independent of tracer count: each tracer still requires flux calculations and memory traffic. Lin-Rood has its own horizontal implementation and should be profiled separately from the split-sweep schemes.

What mmap does—and does not do

TransportBinaryReader memory-maps the version-4 payload as a read-only CPU vector. The map avoids repeated file-open/schema parsing and lets the operating system page in only the accessed bytes. It is not a zero-copy GPU data source.

When a window is requested, the loader:

1. computes section offsets from the parsed header;

2. copies required sections from the mmap into typed host arrays, converting precision when FT differs from the on-disk type;

3. adds cubed-sphere halos where required;

4. for a GPU run, copies the window into persistent backend buffers.

With multiple Julia threads, GPU runs can prefetch the next host window while the current one is being computed. Set ATMOSTR_DISABLE_PREFETCH=1 to disable that overlap when debugging. Linux runs can release already-used mmap pages between files through the driver's release_payload! path.

This distinction matters when diagnosing a regression: window_load_host, window_backend_copy, and the physics kernels are separate costs.

Measure your own run

Enable built-in section timers without changing the config:

ATMOSTR_TIMERS=1 julia --project=. scripts/run_transport.jl my_run.toml

The runner prints a timing breakdown and writes a sibling *.timings.csv when an output NetCDF path is configured. Useful sections include window loading, backend copying, forcing refresh, advection, diffusion, convection, and output. Allocation sampling can be added with ATMOSTR_ALLOC_TIMERS=1.

For CUDA tracing, scripts/run_transport.jl also supports:

ATMOSTR_PROFILE_MODE=full julia --project=. scripts/run_transport.jl my_run.toml

Use ATMOSTR_PROFILE_MODE=window with ATMOSTR_PROFILE_WARMUP_SEC and ATMOSTR_PROFILE_DUR_SEC for a bounded capture. Compare runs with the same binary, precision, backend, Julia thread count, and warmed compilation state; otherwise startup, disk, and compilation effects can dominate the result.

Source map

Concern	Primary source
Backend selection and adaptation	src/Architectures.jl
Runtime window copies and prefetch	src/Models/DrivenSimulation.jl
Section timing	src/Diagnostics/SectionTimer.jl
Structured packed sweeps	src/Operators/Advection/multitracer_kernels.jl
Cubed-sphere packed sweeps	src/Operators/Advection/CubedSphereStrang.jl
Memory-mapped binary reader	src/MetDrivers/transport_binary/reader.jl

Continue with Operators on top of the binary for the physics interfaces or The binary pipeline for the on-disk contract.