[Paper Review] GX: a GPU-native gyrokinetic turbulence code for tokamak and stellarator design

Noah Mandell, W. Dorland+5 more|arXiv (Cornell University)|Sep 14, 2022

Magnetic confinement fusion research4 citations

TL;DR

GX is a GPU-native gyrokinetic turbulence code designed for fast, high-fidelity simulations of low-frequency turbulence in tokamaks and stellarators, using a pseudo-spectral Laguerre-Hermite velocity space formulation and radial flux-tube geometry to enable minutes-scale simulations on a single GPU. It achieves high performance and accurate convergence in both linear and nonlinear benchmarks, making it suitable for fusion reactor design and optimization workflows.

ABSTRACT

GX is a code for solving the nonlinear gyrokinetic system for low-frequency turbulence in magnetized plasmas, particularly tokamaks and stellarators. In GX, our primary motivation and target is a fast gyrokinetic solver that can be used for fusion reactor design and optimization along with wide-ranging physics exploration. This has led to several code and algorithm design decisions, specifically chosen to prioritize time to solution. First, we have used a discretization algorithm that is pseudo-spectral in the entire phase-space, including a Laguerre-Hermite pseudo-spectral formulation of velocity space, which allows for smooth interpolation between coarse gyrofluid-like resolutions and finer conventional gyrokinetic resolutions and efficient evaluation of a model collision operator. Additionally, we have built GX to natively target graphics processors (GPUs), which are among the fastest computational platforms available today. Finally, we have taken advantage of the reactor-relevant limit of small $ρ_*$ by using the radially-local flux-tube approach. In this paper we present details about the gyrokinetic system and the numerical algorithms used in GX to solve the system. We then present several numerical benchmarks against established gyrokinetic codes in both tokamak and stellarator magnetic geometries to verify that GX correctly simulates gyrokinetic turbulence in the small $ρ_*$ limit. Moreover, we show that the convergence properties of the Laguerre-Hermite spectral velocity formulation are quite favorable for nonlinear problems of interest. Coupled with GPU acceleration, which we also investigate with scaling studies, this enables GX to be able to produce useful turbulence simulations in minutes on one (or a few) GPUs. GX is open-source software that is ready for fusion reactor design studies.

Motivation & Objective

• To develop a fast, GPU-optimized gyrokinetic solver for fusion reactor design and optimization.
• To enable rapid simulation turnaround by prioritizing time-to-solution over full physical realism.
• To support smooth resolution transitions from gyrofluid-like to high-fidelity gyrokinetic simulations via a pseudo-spectral Laguerre-Hermite velocity space formulation.
• To validate GX’s accuracy and performance across tokamak and stellarator configurations using established benchmarks.
• To provide an open-source, production-ready tool for reactor-relevant turbulence studies.

Proposed method

• Uses a pseudo-spectral discretization in full phase space, including Laguerre-Hermite spectral basis for velocity space to enable smooth resolution scaling.
• Employs a radially-local flux-tube approximation to exploit the small ρ∗ limit (ρi/R ≪ 1), simplifying geometry and enhancing computational efficiency.
• Natively targets GPUs using CUDA/ROCm, with algorithmic design choices optimized for massive parallelism and memory coalescing.
• Implements hyper-viscosity and hyper-collisions to stabilize nonlinear simulations while preserving spectral convergence.
• Uses a model collision operator compatible with the spectral velocity space formulation for efficient computation.
• Supports both Boltzmann and kinetic electron models, with consistent normalization and flux scaling.

Experimental results

Research questions

• RQ1 Can a GPU-native gyrokinetic code achieve high-performance simulations of turbulence in tokamaks and stellarators with minimal latency?
• RQ2 How does the Laguerre-Hermite spectral velocity space formulation perform in nonlinear turbulence problems compared to traditional grid-based methods?
• RQ3 To what extent does GPU acceleration reduce simulation time while maintaining accuracy in reactor-relevant configurations?
• RQ4 How well does GX reproduce established results from GS2 and stella in both linear and nonlinear benchmarks?
• RQ5 Can GX enable rapid, high-fidelity turbulence simulations suitable for use in fusion reactor design and optimization pipelines?

Key findings

• GX achieves simulation turnaround in minutes on a single GPU for low-resolution cases, with higher-fidelity results in a few hours using multiple GPUs.
• The Laguerre-Hermite spectral velocity space formulation exhibits favorable convergence properties for nonlinear turbulence problems.
• Benchmark results for linear and nonlinear cases in both tokamak and stellarator geometries show excellent agreement with GS2 and stella codes.
• In the W7-X stellarator case, GX successfully reproduces nonlinear turbulence levels comparable to reference simulations using 128 parallel grid points and 192 binormal grid points.
• Scaling studies confirm strong GPU scaling efficiency, with performance gains proportional to GPU count and optimal memory access patterns.
• GX is open-source and production-ready, with full support for reactor-scale design studies and parameter space exploration.