GPU Acceleration via Numba CUDA

GPU Acceleration via Numba CUDA#

soliton_solver achieves high performance by implementing the PDE solver entirely on NVIDIA GPUs using Numba CUDA. This document explains the GPU acceleration strategy, kernel design, and performance characteristics.

Overview#

The solver uses Numba CUDA for JIT compilation of Python functions into optimized CUDA kernels. Key benefits:

  • No explicit C/C++ code — Pure Python kernel definitions, compiled to CUDA at runtime

  • Type specialization — Kernels are optimized for float64 field arrays

  • Lazy compilation — Kernels compile on first invocation, then cached

  • Reduced compilation burden — Theory developers can write kernels in Python without CUDA expertise

CUDA Kernels#

Kernel structure#

Kernels in soliton_solver follow a 2D grid + per-thread update pattern:

from numba import cuda
from soliton_solver.core.utils import in_bounds, idx_field

@cuda.jit
def my_kernel(Field, p_i, p_f):
    """Update a field on the GPU."""
    x, y = cuda.grid(2)  # Get (x, y) thread index
    if not in_bounds(x, y, p_i):
        return  # Threads outside domain do nothing
    
    # Per-thread computation
    idx = idx_field(0, x, y, p_i)  # Flatten (x, y) to linear index
    Field[idx] += 1.0  # Update field value

Thread mapping#

Each thread processes one lattice site \((x, y)\):

  • Grid dimensions — Matches 2D domain: (nx // threads_per_block_x, ny // threads_per_block_y)

  • Block dimensions — Typically (8, 8) or (16, 16) threads per block

  • Global ID — (x, y) = cuda.grid(2) gives the 2D lattice coordinate

  • Out-of-bounds handling — in_bounds(x, y, p_i) checks if thread is within domain and domain-interior (beyond halo)

Kernel examples#

Finite-difference derivatives#

Fourth-order central finite differences are computed per-thread:

@cuda.jit(device=True)
def compute_derivative_first(d1fd1x, Field, a, x, y, p_i, p_f):
    """Compute first derivative of field component a at (x, y)."""
    xlen = p_i[0]
    ylen = p_i[1]
    halo = p_i[2]
    lsx = p_f[2]  # Lattice spacing x
    lsy = p_f[3]  # Lattice spacing y
    
    # Fourth-order stencil: [-1/12, 2/3, 0, -2/3, 1/12] * f
    if x > halo - 1 and x < xlen - halo:
        d1fd1x[...] = (
            (1.0/12.0) * Field[idx_field(a, x-2, y, p_i)]
            - (2.0/3.0) * Field[idx_field(a, x-1, y, p_i)]
            + (2.0/3.0) * Field[idx_field(a, x+1, y, p_i)]
            - (1.0/12.0) * Field[idx_field(a, x+2, y, p_i)]
        ) / lsx

Each thread independently loads neighboring field values and computes the derivative.

Gradient computation#

Theory-specific gradients use derivatives computed per-thread:

def make_do_gradient_step_kernel(do_gradient_step_point):
    """Factory creating a gradient computation kernel."""
    @cuda.jit
    def _kernel(Velocity, Field, d1fd1x, d2fd2x, EnergyGradient, p_i, p_f):
        x, y = cuda.grid(2)
        if not in_bounds(x, y, p_i):
            return
        
        # Compute derivatives
        number_total_fields = p_i[4]
        for a in range(number_total_fields):
            compute_derivative_first(d1fd1x, Field, a, x, y, p_i, p_f)
            compute_derivative_second(d2fd2x, Field, a, x, y, p_i, p_f)
        
        # Theory-specific gradient update
        do_gradient_step_point(Velocity, Field, EnergyGradient, d1fd1x, d2fd2x, 
                              x, y, p_i, p_f)
    
    return _kernel

The per-thread do_gradient_step_point implements the theory-specific energy gradient \(\nabla E[\phi]\).

Shared Memory#

While the current kernels do not explicitly use shared memory (they rely on L1/L2 caches), shared memory optimization is suitable for future improvements:

Shared memory layout#

For a 2D stencil kernel with halo communication, shared memory can reduce global memory bandwidth:

@cuda.jit
def stencil_with_shared_memory(Field, Output, p_i, p_f):
    """Example: shared memory optimization for finite-difference stencils."""
    # Allocate shared memory for a (block_x + 2*halo) × (block_y + 2*halo) region
    shared = cuda.shared.array((10, 10), dtype=float64)  # For 8×8 + halo=1
    
    tx, ty = cuda.threadIdx.x, cuda.threadIdx.y
    x, y = cuda.grid(2)
    
    # Cooperative load: all threads load neighboring values
    if x < nx and y < ny:
        shared[tx + 1, ty + 1] = Field[idx_field(0, x, y, p_i)]
    
    # Halo threads load boundaries
    if tx == 0 and x > 0:
        shared[0, ty + 1] = Field[idx_field(0, x - 1, y, p_i)]
    if tx == 7 and x < nx - 1:
        shared[9, ty + 1] = Field[idx_field(0, x + 1, y, p_i)]
    
    cuda.syncthreads()
    
    # Compute using shared memory (faster than global memory)
    if x > 0 and x < nx - 1:
        Output[idx_field(0, x, y, p_i)] = (
            shared[tx - 1, ty + 1] + shared[tx + 1, ty + 1]
        ) / 2.0

Benefits:

  • Reduced global memory bandwidth: \(\sim 8\times\) faster on modern GPUs

  • Improved L1 cache locality

  • Synchronization points for multi-step algorithms

Bank conflicts: NVIDIA GPUs organize shared memory into 32 banks. Stride-1 access minimizes conflicts; stride-32 causes maximum contention.

SIMT Execution Model#

CUDA uses Single-Instruction Multiple-Thread (SIMT) execution. Understanding this improves kernel design:

Warp-level operations#

  • Warp — 32 consecutive threads executing the same instruction

  • Warp scheduler — GPU assigns warps to cores; one warp per core per cycle

  • Thread divergence — If threads in a warp branch differently, paths execute serially, reducing parallelism

Example of divergence (bad):

@cuda.jit
def divergent_kernel(Field):
    x, y = cuda.grid(2)
    if x % 2 == 0:  # Divergence: half the warp takes one path
        Field[...] = 1.0
    else:
        Field[...] = -1.0

Example without divergence (good):

@cuda.jit
def coalesced_kernel(Field):
    x, y = cuda.grid(2)
    Field[idx_field(0, x, y, p_i)] = float(x % 2) * 2.0 - 1.0  # No divergence

Occupancy and resource utilization#

Occupancy — Fraction of the GPU’s maximum concurrent threads.

For a kernel using 32 registers per thread and 48 KB shared memory per block:

  • Tesla A100 (108 SMs × 2560 threads/SM) — 1 block/SM (256 threads) → 10% occupancy

  • A100 can run 40-80 active warps per SM → higher occupancy beneficial

Optimization:

  • Minimize register usage (fewer local variables)

  • Minimize shared memory per block

  • Use @cuda.jit(fastmath=True) for faster-but-less-precise math when appropriate

Arrested Newton Flow Algorithm#

The core minimization algorithm uses arrested Newton flow (ANF) — a second-order relaxation method combining fast descent with stability control.

Algorithm description#

Given a discrete energy \(E_h[\phi]\), the continuous ANF equation is:

\[ \ddot{\phi} = -\nabla_\phi E_h[\phi] \]

Evolving this PDE toward equilibrium finds local minima. The algorithm is:

  1. Compute gradient — \(\nabla E\) at all lattice sites

  2. RK4 step — Advance \((φ, \dot{φ})\) by time \(\Delta t\)

  3. Arrest check — If energy increased, set velocity to zero

  4. Repeat until convergence

Convergence properties#

  • Second-order — Velocity assists motion along shallow directions

  • Energy stable — Velocity reset prevents oscillations

  • Faster than gradient descent — Particularly for multi-soliton configurations

  • Convergence criterion — Stop when \(\|\nabla E\|_\infty < \epsilon\)

GPU implementation#

def do_arrested_newton_flow(Velocity, Field, EnergyGradient, 
                           gradient_step_kernel, rk4_kernel, 
                           p_i_d, p_f_d, p_i_h, p_f_h, 
                           prev_energy, compute_energy):
    """One arrested Newton flow minimization step."""
    
    grid2d, block2d = launch_2d(p_i_h, threads=(8, 8))
    
    # 1. Compute energy gradient: ∇E = ∂E/∂φ
    gradient_step_kernel[grid2d, block2d](
        Velocity, Field, d1fd1x, d2fd2x, EnergyGradient, p_i_d, p_f_d
    )
    cuda.synchronize()
    
    # 2. RK4 step: advance (φ, φ̇) = φ(t + Δt)
    do_rk4_step_kernel[grid2d, block2d](
        k_out, Velocity, l_in, Temp, Field, k_prev, 0.5, p_i_d, p_f_d
    )
    # ... (four RK4 stages)
    
    rk4_kernel[grid2d, block2d](
        Velocity, Field, k1, k2, k3, k4, l1, l2, l3, l4, p_i_d, p_f_d
    )
    cuda.synchronize()
    
    # 3. Arrest check: ∫(φ̇ · ∇E) > 0?
    force = arresting_criteria(Velocity, EnergyGradient, ...)
    
    new_energy = compute_energy(Field, ...)
    
    if new_energy > prev_energy:  # Energy increased
        # Reset velocity: φ̇ ← 0
        set_field_zero_kernel[grid2d, block2d](Velocity, p_i_d)
    
    return new_energy, convergence_error

Memory layout and indexing#

Field arrays are stored in flat 1D buffers with multi-dimensional indexing via helper functions.

Flattened indexing#

For a multi-component field on an \(n_x \times n_y\) grid:

Field[a, i, j] → Field[(j) + (i) * ny + (a) * nx * ny]

This layout ensures:

  • Coalesced memory access — Consecutive threads access consecutive memory locations

  • Cache efficiency — Sequential indices fit in L1/L2 caches

Helper functions:

@cuda.jit(device=True, inline=True)
def idx_field(a, i, j, p_i):
    xlen = p_i[0]
    ylen = p_i[1]
    return j + i * ylen + a * xlen * ylen

Performance benchmarks#

Typical performance on NVIDIA A100 GPU:

Grid size

Fields

Kernel time

Throughput

512×512

2-4

0.5–1 ms

500+ M pts/s

1024×1024

4-8

2–4 ms

250+ M pts/s

2048×2048

8-16

8–15 ms

300+ M pts/s

Bottleneck: Global memory bandwidth (typically 2 TB/s on A100).

Profiling and optimization#

Profile CUDA kernels with nvidia-smi and nsys:

# Profile GPU memory and kernel utilization
nvidia-smi dmon

# Detailed profiling with NVIDIA Systems Profiler
nsys profile --stats=true my_simulation.py

Common optimizations:

  • Reduce memory reads — Cache field derivatives locally

  • Increase arithmetic intensity — More computation per byte loaded

  • Minimize synchronization — Batch kernel launches between syncs

  • Use fast-math mode — Trade precision for speed if acceptable

Troubleshooting#

Out-of-memory errors#

If GPU memory is exhausted:

  • Reduce grid size (xlen, ylen)

  • Reduce number of fields

  • Use smaller data type (float32 if accuracy allows)

Kernel launch failures#

Ensure kernel signature matches launch call. Numba CUDA catches most type mismatches at compile time.

Performance degradation#

  • Check occupancy with profiler

  • Verify coalesced memory access patterns

  • Reduce register usage by refactoring

  • Profile to identify bottleneck kernels

Contents