Skip to content

gpu.md

Latest commit

 

History

History
211 lines (151 loc) · 10.4 KB

gpu.md

File metadata and controls

211 lines (151 loc) · 10.4 KB
Raw

GPU acceleration guide

StructureFunctions.jl provides GPU structure-function kernels via StructureFunctionsGPUExt (loaded when KernelAbstractions.jl is available). This guide covers production APIs: single snapshots, workspace reuse, and time-slice batches.

For backend selection across serial / threaded / distributed / GPU, see backends.md. For CUDA validation and SLURM benchmark scripts, see gpu/README.md.

When to use the GPU

  • Problem size: GPU pair histograms pay off when N is roughly few×10³ and up (exact crossover depends on hardware; see benchmark figures).
  • Memory: Device arrays use layout (N_dims, N_points) or (N_dims, N_points, T) for batches. Prefer Float32 on GPU for bandwidth; Float64 is supported but often much slower.
  • Not a drop-in speedup on CPU: Without a GPU, KA.CPU() runs the same kernels on CPU and is slower than ThreadedBackend().

Single snapshot

using StructureFunctions: StructureFunctions as SF, Calculations as SFC
using KernelAbstractions: KernelAbstractions as KA
using CUDA: CUDA

backend = CUDA.functional() ? CUDA.CUDABackend() : KA.CPU()
N = 20_000
FT = Float32
x = CUDA.functional() ? CUDA.CuArray{FT}(rand(FT, 3, N)) : rand(FT, 3, N)
u = CUDA.functional() ? CUDA.CuArray{FT}(rand(FT, 3, N)) : rand(FT, 3, N)
bins = collect(FT, range(0.0f0, 1.5f0; length = 21))
sft = SF.LongitudinalSecondOrderStructureFunctionType()

result = SFC.gpu_calculate_structure_function(
    sft, backend, x, u, bins,
)  # returns a StructureFunctionSumsAndCounts (raw sums + counts)

distance_bins must have the same element type as x and u (e.g. Float32 fields → collect(Float32, edges)). The GPU API does not cast bin edges silently.

GPUSFWorkspace — reuse device histogram buffers

Repeated calls with the same bin layout should reuse a workspace to avoid reallocating device histogram buffers each launch (~10–15% end-to-end win at N=20k on A100):

ws = SFC.GPUSFWorkspace(backend, bins)
for _ in 1:10
    SFC.gpu_calculate_structure_function(
        sft, backend, x, u, bins; workspace = ws,
    )
end
SFC.release!(ws)  # optional explicit free

Time series — slice batch API

For T snapshots, stack data as (N_dims, N_points, T), upload once, and call the slice driver:

T = 100
x_batch = rand(FT, 3, N, T)   # host or CuArray
u_batch = rand(FT, 3, N, T)

sums = zeros(FT, length(bins) - 1, T)
counts = zeros(UInt32, length(bins) - 1, T)
ws = SFC.GPUSFWorkspace(backend, bins)

SFC.gpu_calculate_structure_function_batch!(
    sums, counts, sft, backend, x_batch, u_batch, bins; workspace = ws,
)

Avoid a naive loop that uploads a host slice each t — that pays H2D + allocation every step. The slice driver keeps the batch on device, reuses the workspace, and performs one final sync.

Public stubs with backend dispatch:

  • calculate_structure_function_batch!
  • calculate_structure_function_2d_batch!
  • calculate_structure_functions_single_pass_batch! (GPU-only for now)
  • calculate_structure_functions_single_pass_2d! — six (dist × value) histograms; GPU HTP-EJ when eligible

Production route table

Public call Shapes Dimensions Bin support GPU route
calculate_structure_function(sf, x, u, bins; backend=GPUBackend(...)) (D,N) D=2,3 linear, log, general vectors tiled128 1D histogram route; explicit error for unsupported D
same with shared positions x::(D,N), u::(D,N, auxiliary...) currently production optimized for D=2 auxiliary batch kernels linear distance bins for fused batch route fixed-position auxiliary batch route
same with varying positions x,u::(D,N, auxiliary...) currently production optimized for D=2 auxiliary batch kernels linear distance bins for fused batch route varying-position auxiliary batch route
calculate_structure_function(sf, x, u, bins, value_bins; backend=GPUBackend(...)) (2,N) or auxiliary variants D=2 typed and vector value-bin plans joint 2D shared-memory route when eligible, otherwise global route
calculate_structure_functions_single_pass(x, u, bins; backend=GPUBackend(...)) (D,N) or auxiliary variants D=2,3 point fields; D=2 optimized auxiliary batch kernels linear/log/general point fields; linear fused batch route tiled 2D fast path for eligible D=2, global fallback otherwise
calculate_structure_functions_single_pass_2d(x, u, bins, value_bins; backend=GPUBackend(...)) (D,N) or auxiliary variants D=2,3 point fields; D=2 optimized auxiliary batch kernels linear/log/general point fields; typed or vector value bins HTP-EJ shared/typeplane/direct strategy for eligible D=2; global fallback otherwise

The public shape contract is always D = size(x, 1) = size(u, 1), N = size(x, 2) = size(u, 2). Trailing axes are independent auxiliary calculations; ndims(x) == 3 is not interpreted as 3D space.

Single-type joint 2D smem

GPUSFWorkspace for kind=:joint2d defaults to exact compile-time shared histogram width n_dist × n_val. Optional override: joint2d_compile_cells=joint2d_smem_max() or joint2d_smem_align256(n_dist, n_val). See gpu/GPU_2d_joint_sf_plan.md.

Six-invariant-type single-pass 2D (SP2D)

Production GPU path for calculate_structure_functions_single_pass_2d! with typed distance bins (LinearBinEdges / LogBinEdges) and GPUSFWorkspace(...; kind=:single_pass_2d). Histogram strategy (:shared / :typeplane / :direct) is frozen at workspace build from a 48 KiB shared-memory budget.

Rows are fixed to the six invariant/default quantities:

  1. S2 = |delta u|^2
  2. L2 = delta u_L^2
  3. T2 = |delta u_T|^2
  4. S3 = delta u_L |delta u|^2
  5. L3 = delta u_L^3
  6. LT2 = delta u_L |delta u_T|^2

Basis-dependent component diagnostics such as T3SF and L2T1SF are not included in default single-pass outputs. Those require an explicit transverse-basis convention rather than the invariant bulk six-row contract.

  • On-chip (:shared, :typeplane): shared histogram + joint-style flush to output (no merge).
  • Direct (:direct): partitioned global accumulation + merge when even one value plane does not fit in smem.

Benchmark on GPU: julia --project=gpu gpu/benchmark_2d_grid_scaling.jl
Design, gate, perf gaps: gpu/SP2D_HTP_EJ.md

The production gate is e2e SP2D < 6 × joint_2d. A naive “~2× digitize vs 1D” bound is too optimistic: SP2D performs six value digitizations per pair and may replay the full tile schedule n_type_passes times (typeplane). See the doc for a per-pair work table and future optimizations.

Testing tiers

Tier Command What it proves
1 — default CI Pkg.test() Kernel math, binning, workspace reset, slice logic via KA.CPU() (same @kernel source, no CUDA)
2 — CUDA smoke julia --project=gpu gpu/runtests.jl Device alloc, H2D/D2H, sync, Float32 on real GPU (skipped if !CUDA.functional())
3 — benchmarks julia --project=gpu gpu/benchmark_suite.jl Release-performance gates and timing JSON (run on GPU allocation)

Important: Tier 1 does not prove CUDA correctness. Always run tier 2 on a GPU node before trusting production CUDA runs.

Tier 2 tests live in gpu/runtests.jl and are not included in test/runtests.jl.

Benchmarks & figures

Shared bin layout and SF type with CPU benchmarks (benchmark/scaling_config.jl).

Study Script What varies What is fixed
CPU strong scaling benchmark/benchmark_scaling.jl threads N
CPU weak scaling same threads + N work/thread
GPU problem-size scaling gpu/collect_benchmark_assets.jl N 1 GPU, serial CPU
GPU slice-batch scaling same T (slices) N_SLICE, 1 GPU
GPU release gates gpu/benchmark_suite.jl route 1 GPU, shared bins
GPU strong/weak (multi-GPU) gpu/collect_multi_gpu_scaling.jl — not implemented

Problem-size scaling is the usual name for “one device, sweep input size.” It is not HPC strong or weak scaling.

The GPU collector always uses SerialBackend for the CPU reference (1 logical worker), independent of julia -t. That keeps doc assets reproducible on any GPU allocation. CPU thread scaling is only in benchmark/benchmark_scaling.jl (strong/weak figures); readers combine those plots with the GPU problem-size figure as needed.

The release benchmark suite is separate from docs asset generation:

julia --project=gpu gpu/benchmark_suite.jl

It writes gpu/benchmark_results/benchmark_suite_latest.json and a timestamped copy. The key ratios are workspace reuse, 6 * joint2D vs SP2D, shared-position auxiliary fusion vs explicit loops, and varying-position auxiliary fusion vs explicit loops. BENCH_BACKEND=kacpu is useful only as a smoke test; CUDA runs with representative N and BATCH are the performance signal.

Regenerate GPU doc assets (on GPU allocation)

julia --project=gpu gpu/collect_benchmark_assets.jl
julia --project=docs/generate_assets docs/generate_assets/generate_gpu_figures.jl

Outputs:

  • gpu/benchmark_results/assets_latest.json
  • docs/src/assets/gpu_problem_size_scaling.png
  • docs/src/assets/gpu_slice_batch_scaling.png

Parity figure (KA.CPU vs serial, no GPU):

julia --project=docs/generate_assets docs/generate_assets/generate_assets.jl

Produces docs/src/assets/sf_gpu_parity.png.

Figures (from latest collector run)

GPU problem-size scaling

GPU slice-batch scaling

Examples

See also