GPU Kernel Programming for Custom Physics Simulation: Writing High-Performance Custom Kernels in Python

GPU acceleration in scientific computing often begins with high-level libraries. A researcher replaces NumPy with CuPy, uses an accelerated FFT, or calls a GPU-enabled linear algebra routine. This can produce a substantial improvement without requiring detailed knowledge of GPU hardware.

Custom GPU kernel programming goes one level deeper. Instead of calling a predefined operation, the developer writes the function that every GPU thread executes. This provides control over thread indexing, data layout, memory transfers, synchronization, shared memory, and the number of operations performed during each kernel launch.

This control is useful when a physical model contains access patterns that cannot be expressed efficiently through ordinary array operations. Particle interactions, finite difference stencils, collision rules, lattice-based fluid models, material updates, and geometry queries are common examples.

A custom kernel is not automatically faster than an optimized library. It becomes valuable when the algorithm exposes enough parallel work and when the developer can reduce memory traffic, combine operations, or adapt the computation to the GPU architecture.

What Is a GPU Kernel?

A GPU kernel is a function executed by many lightweight threads. Each thread normally handles one particle, grid cell, element, face, pixel, or entry in an array.

The host program launches the kernel with a grid of thread blocks:

kernel[blocks, threads_per_block](arguments)

Every thread determines which data it should process from its block and thread indices. A typical one-dimensional index is:

index =
    block_index * block_size
    + thread_index

Threads are organized into groups called blocks. Threads inside one block can synchronize and exchange data through shared memory. Separate blocks are generally expected to execute independently.

When Custom Kernels Are Worth Writing

High-level GPU libraries should usually be the first option. They already provide optimized matrix multiplication, FFTs, reductions, sparse operations, random number generation, and many element-wise functions.

A custom kernel becomes useful when:

• The simulation uses a non-standard memory-access pattern.
• Several small operations can be fused into one pass over memory.
• Neighboring threads repeatedly need the same local data.
• A specialized particle, collision, or stencil rule dominates runtime.
• Intermediate arrays consume too much device memory.
• The simulation needs a custom differentiable operation.
• An existing library cannot express the required boundary or material logic efficiently.

Before writing a kernel, profile the program. Optimizing a visually complex function that accounts for only a small part of runtime will not produce a meaningful overall speedup.

GPU Programming Models

Scientific software can access GPUs through several levels of abstraction.

The ENCCS introduction to GPU programming models provides a broader comparison of CUDA, HIP, OpenCL, SYCL, Kokkos, OpenMP, and related approaches.

CUDA

CUDA provides direct access to NVIDIA GPU programming through C++, Fortran, and several language bindings. It exposes grids, blocks, warps, shared memory, streams, events, device libraries, and hardware-specific optimization tools.

CUDA is often selected when maximum NVIDIA-specific performance and mature tooling are more important than hardware portability. The main costs are manual memory management, a lower-level programming model, and dependence on the NVIDIA platform.

HIP and ROCm

HIP provides a CUDA-like C++ programming model within AMD’s ROCm ecosystem. A large amount of CUDA-style code can be adapted to HIP, but source portability does not mean that one binary runs unchanged on every GPU.

Separate compilation, testing, and performance tuning may still be required for each target architecture. Details are available in the official HIP documentation.

HIP is a practical option when AMD GPUs must be supported or when a project wants to reduce dependence on a single accelerator vendor while retaining a CUDA-like programming style.

CuPy Custom Kernels

CuPy is best known for its NumPy-compatible arrays and accelerated numerical functions. It also supports custom GPU programming.

Developers can use:

• ElementwiseKernel for custom element-wise expressions
• ReductionKernel for custom reductions
• RawKernel for kernels written in CUDA C++
• cupyx.jit.rawkernel for Python-style JIT kernel definitions

A simple CuPy raw kernel can be defined as follows:

import cupy as cp

scale_kernel = cp.RawKernel(
    r'''
    extern "C" __global__
    void scale_values(
        const float* input,
        float* output,
        const float factor,
        const int size
    ) {
        int index =
            blockDim.x * blockIdx.x
            + threadIdx.x;

        if (index < size) {
            output[index] =
                factor * input[index];
        }
    }
    ''',
    "scale_values"
)

size = 1_000_000
threads_per_block = 256
blocks = (
    size + threads_per_block - 1
) // threads_per_block

input_values = cp.arange(
    size,
    dtype=cp.float32
)

output_values = cp.empty_like(
    input_values
)

scale_kernel(
    (blocks,),
    (threads_per_block,),
    (
        input_values,
        output_values,
        cp.float32(2.0),
        cp.int32(size)
    )
)

CuPy is useful when most of the application already uses GPU arrays and only selected operations need custom kernels.

Numba-CUDA

Numba-CUDA compiles a restricted subset of Python into GPU kernels. Its programming model closely follows CUDA C. Developers define a function with a kernel decorator, calculate a thread index, and launch the function with a grid and block configuration.

The active project documentation is available at Numba-CUDA. Because its development status can change, teams building long-lived software should review the current maintenance and migration guidance before committing to the framework.

A Numba Particle Kernel

The following educational example calculates gravitational acceleration with direct all-pairs interactions:

import math
import numpy as np
from numba import cuda

@cuda.jit
def gravity_kernel(
    positions,
    masses,
    accelerations,
    gravitational_constant,
    softening_squared
):
    particle = cuda.grid(1)
    particle_count = positions.shape[0]

    if particle >= particle_count:
        return

    px = positions[particle, 0]
    py = positions[particle, 1]
    pz = positions[particle, 2]

    ax = 0.0
    ay = 0.0
    az = 0.0

    for other in range(particle_count):
        if other == particle:
            continue

        dx = positions[other, 0] - px
        dy = positions[other, 1] - py
        dz = positions[other, 2] - pz

        distance_squared = (
            dx * dx
            + dy * dy
            + dz * dz
            + softening_squared
        )

        inverse_distance = (
            1.0
            / math.sqrt(distance_squared)
        )

        inverse_distance_cubed = (
            inverse_distance
            * inverse_distance
            * inverse_distance
        )

        scale = (
            gravitational_constant
            * masses[other]
            * inverse_distance_cubed
        )

        ax += scale * dx
        ay += scale * dy
        az += scale * dz

    accelerations[particle, 0] = ax
    accelerations[particle, 1] = ay
    accelerations[particle, 2] = az


particle_count = 10_000
threads_per_block = 256

blocks = (
    particle_count
    + threads_per_block
    - 1
) // threads_per_block

host_positions = np.random.random(
    (particle_count, 3)
).astype(np.float32)

host_masses = np.ones(
    particle_count,
    dtype=np.float32
)

device_positions = cuda.to_device(
    host_positions
)

device_masses = cuda.to_device(
    host_masses
)

device_accelerations = cuda.device_array(
    (particle_count, 3),
    dtype=np.float32
)

gravity_kernel[
    blocks,
    threads_per_block
](
    device_positions,
    device_masses,
    device_accelerations,
    np.float32(1.0),
    np.float32(1e-4)
)

cuda.synchronize()

accelerations = (
    device_accelerations.copy_to_host()
)

This kernel has computational complexity of O(N²). It is suitable for explaining thread mapping, but it is not an efficient production algorithm for very large particle counts.

Large gravitational or particle systems often require a tree method, fast multipole method, spatial bins, or neighbor lists. A GPU cannot remove the scaling problem of an inefficient algorithm.

NVIDIA Warp

NVIDIA Warp allows developers to define strongly typed kernels with Python syntax. Warp compiles those functions into CPU or CUDA code and provides specialized primitives for geometry, simulation, sparse operations, finite elements, optimization, and automatic differentiation.

The current NVIDIA Warp product page and Warp GitHub repository include examples of particles, fluids, meshes, optimization, differentiable simulation, and tile-based GPU programming.

The Same Particle Kernel in Warp

import numpy as np
import warp as wp

wp.init()

@wp.kernel
def gravity_kernel(
    positions: wp.array(dtype=wp.vec3),
    masses: wp.array(dtype=wp.float32),
    accelerations: wp.array(dtype=wp.vec3),
    gravitational_constant: wp.float32,
    softening_squared: wp.float32
):
    particle = wp.tid()
    particle_count = positions.shape[0]

    position = positions[particle]
    acceleration = wp.vec3(0.0, 0.0, 0.0)

    for other in range(particle_count):
        if other != particle:
            displacement = (
                positions[other]
                - position
            )

            distance_squared = (
                wp.dot(
                    displacement,
                    displacement
                )
                + softening_squared
            )

            inverse_distance = (
                1.0
                / wp.sqrt(distance_squared)
            )

            inverse_distance_cubed = (
                inverse_distance
                * inverse_distance
                * inverse_distance
            )

            acceleration += (
                gravitational_constant
                * masses[other]
                * inverse_distance_cubed
                * displacement
            )

    accelerations[particle] = acceleration


particle_count = 10_000

host_positions = np.random.random(
    (particle_count, 3)
).astype(np.float32)

host_masses = np.ones(
    particle_count,
    dtype=np.float32
)

positions = wp.array(
    host_positions,
    dtype=wp.vec3,
    device="cuda"
)

masses = wp.array(
    host_masses,
    dtype=wp.float32,
    device="cuda"
)

accelerations = wp.zeros(
    particle_count,
    dtype=wp.vec3,
    device="cuda"
)

wp.launch(
    kernel=gravity_kernel,
    dim=particle_count,
    inputs=[
        positions,
        masses,
        accelerations,
        wp.float32(1.0),
        wp.float32(1e-4)
    ],
    device="cuda"
)

wp.synchronize()

Warp determines a suitable launch configuration internally for an ordinary kernel launch. Advanced users can still work with block dimensions, command graphs, tiled execution, and specialized device operations when required.

Grid-Stride Loops

A kernel does not need one permanently assigned thread for every element. A grid-stride loop allows each thread to process several entries:

from numba import cuda

@cuda.jit
def scale_with_stride(
    input_values,
    output_values,
    factor
):
    index = cuda.grid(1)
    stride = cuda.gridsize(1)

    for position in range(
        index,
        input_values.size,
        stride
    ):
        output_values[position] = (
            factor
            * input_values[position]
        )

This pattern separates the number of launched threads from the total data size. It is useful when processing very large arrays or reusing a fixed launch configuration.

A practical introduction to this pattern is available in The Data Frog’s Numba and CUDA tutorial.

Memory Coalescing

GPU kernels are often limited by memory bandwidth rather than arithmetic throughput. Threads inside a warp should ideally access nearby memory addresses so the hardware can combine their requests.

Consider particle data stored as:

particle_0: x, y, z, mass
particle_1: x, y, z, mass
particle_2: x, y, z, mass

This array-of-structures layout may be convenient for object-oriented code. A structure-of-arrays layout stores separate contiguous arrays:

x_positions[]
y_positions[]
z_positions[]
masses[]

The second layout may provide better coalescing when every thread reads the same field for a different particle. The best layout still depends on which fields are accessed together.

Vector types do not automatically correct a poor access pattern. Developers should inspect the actual addresses requested by neighboring threads.

Shared Memory

Shared memory is a small, low-latency memory area accessible by threads in the same block. It can reduce repeated reads from global memory.

A one-dimensional stencil can load a block of values and its halo into shared memory:

import numpy as np
from numba import cuda, float32

BLOCK_SIZE = 256

@cuda.jit
def three_point_stencil(
    input_values,
    output_values
):
    shared = cuda.shared.array(
        shape=BLOCK_SIZE + 2,
        dtype=float32
    )

    local_index = cuda.threadIdx.x
    global_index = cuda.grid(1)
    size = input_values.size

    center = local_index + 1

    if global_index < size:
        shared[center] = (
            input_values[global_index]
        )
    else:
        shared[center] = 0.0

    if local_index == 0:
        left_index = global_index - 1

        shared[0] = (
            input_values[left_index]
            if left_index >= 0
            else 0.0
        )

    if local_index == BLOCK_SIZE - 1:
        right_index = global_index + 1

        shared[BLOCK_SIZE + 1] = (
            input_values[right_index]
            if right_index < size
            else 0.0
        )

    cuda.syncthreads()

    if global_index < size:
        output_values[global_index] = (
            shared[center - 1]
            + shared[center]
            + shared[center + 1]
        ) / 3.0

The call to cuda.syncthreads() ensures that every thread finishes loading its values before the stencil calculation begins.

Shared memory should not be used automatically. Excessive shared-memory allocation can reduce occupancy, add synchronization cost, and make the kernel slower. Profiling is required.

Kernel Fusion

Separate array operations often create several kernel launches and intermediate arrays:

velocity += dt * acceleration
position += dt * velocity
energy = compute_energy(position, velocity)

A fused kernel may calculate all three updates while the necessary values remain in registers. This reduces launch overhead and global memory traffic.

Fusion is most beneficial when operations are simple and memory-bound. Fusing too much work can increase register use, reduce occupancy, and make the kernel difficult to maintain.

Choosing Block Size

Values such as 128 or 256 threads per block are reasonable starting points, not universal optima.

The best block size depends on:

• Registers used per thread
• Shared memory used per block
• Branch divergence
• Instruction mix
• Memory-access behavior
• The target GPU architecture

A rule such as launching at least twice as many blocks as streaming multiprocessors can be a useful initial experiment, but it does not guarantee maximum performance. Occupancy calculators and profiling tools should guide the final configuration.

Differentiable Physics

Differentiable simulation calculates how an output changes with respect to inputs such as material properties, forces, geometry, or initial conditions.

Warp can record supported kernel operations and execute reverse-mode automatic differentiation. This can be used for inverse problems, design optimization, parameter estimation, and integration with machine learning workflows.

A simplified pattern uses a tape:

with wp.Tape() as tape:
    wp.launch(
        kernel=simulation_kernel,
        dim=element_count,
        inputs=[state, parameters],
        outputs=[result],
        device="cuda"
    )

    wp.launch(
        kernel=loss_kernel,
        dim=element_count,
        inputs=[result, target],
        outputs=[loss],
        device="cuda"
    )

tape.backward(loss)

Automatic differentiation is not guaranteed for every kernel. In-place overwrites, nondeterministic atomics, external native code, discontinuous logic, and unsupported operations may require reformulation or custom gradients.

This capability connects directly with applications such as adjoint optimization and physics-informed neural networks.

When a GPU May Not Be Faster

There is no universal problem-size threshold at which a GPU becomes faster than a CPU. The crossover depends on hardware, precision, data movement, algorithm structure, compiler quality, and how often the same device-resident data is reused.

The correct workflow is to measure the CPU version, GPU-library version, and custom-kernel version with representative data. The scientific computing performance profiling guide explains how to identify the real bottleneck before optimizing it.

Avoiding Transfer Overhead

Repeated transfers between host and device memory can eliminate the benefit of GPU computation.

An inefficient loop may follow this pattern:

1. Copy data to the GPU.
2. Run one small kernel.
3. Copy the result to the CPU.
4. Modify it on the CPU.
5. Copy it back to the GPU.

A better design keeps the simulation state on the device for many time steps and transfers only the output required for visualization, checkpointing, or analysis.

Asynchronous streams, pinned host memory, overlapping communication, and unified-memory features can help, but they should be introduced only after ordinary transfer costs have been measured.

The XLB Case Study

Autodesk Research’s XLB project provides a useful example of Python-native GPU simulation. XLB is an open-source lattice Boltzmann library with several computational backends, including NVIDIA Warp.

In benchmark configurations reported by Autodesk Research and NVIDIA, the Warp backend reached performance close to the compared C++/OpenCL FluidX3D implementation for a specific lid-driven cavity case. A separate comparison reported an approximate eightfold speedup over XLB’s JAX backend on selected hardware and configurations.

The team also demonstrated an out-of-core approach on an eight-node GH200 cluster with a domain of approximately 50 billion lattice cells. These results apply to the reported solver, problem, hardware, and implementation choices. They should not be treated as general speedup guarantees for Python kernels.

The full case study is available in Autodesk Research Brings Warp Speed to Computational Fluid Dynamics on NVIDIA GH200. The current project code is available in the Autodesk XLB repository.

How to Benchmark a Kernel Correctly

GPU operations are usually asynchronous. Measuring only the Python function-call duration can report the time required to enqueue the kernel rather than the time required to execute it.

A basic benchmark should:

1. Run the kernel several times to trigger compilation and warm-up.
2. Synchronize before starting the timer.
3. Run several measured iterations.
4. Synchronize before stopping the timer.
5. Report average and variation across repetitions.
6. Separate data-transfer time from kernel-execution time.
7. Verify that CPU and GPU versions produce equivalent results.

import time
from numba import cuda

# Warm-up and JIT compilation
kernel[blocks, threads](*arguments)
cuda.synchronize()

start = time.perf_counter()

for _ in range(100):
    kernel[blocks, threads](*arguments)

cuda.synchronize()

elapsed = time.perf_counter() - start
average = elapsed / 100

print("Average kernel time:", average)

The comparison should use a realistic simulation size and include the full workflow when transfer or preprocessing costs matter.

Work-precision methodology can also compare runtime against numerical error. The SciML benchmark documentation provides examples of this evaluation style.

A Practical Development Workflow

1. Implement and verify a clear CPU reference version.
2. Profile the application to locate the dominant operation.
3. Try an optimized GPU library before writing a kernel.
4. Keep frequently reused data on the device.
5. Write the simplest correct kernel.
6. Validate it against the CPU result.
7. Measure memory transfers and execution separately.
8. Inspect coalescing, occupancy, divergence, and register use.
9. Test shared memory or fusion only when profiling supports them.
10. Benchmark several problem sizes and GPU architectures.

Choosing a Framework

Common Kernel Programming Mistakes

• Moving data between the CPU and GPU inside every time step
• Ignoring out-of-bounds threads
• Using an inefficient algorithm and expecting hardware to fix its scaling
• Accessing memory with an uncoalesced layout
• Adding shared memory without measuring whether it helps
• Launching many tiny kernels instead of considering fusion
• Using excessive registers or shared memory per block
• Benchmarking asynchronous code without synchronization
• Comparing outputs without checking numerical precision
• Using fixed performance claims across different GPUs and workloads
• Assuming portable source provides portable performance
• Applying automatic differentiation to unsupported in-place operations

Related Guides

• GPU-Accelerated Scientific Computing: CuPy, Numba, and CuDF Compared — Compare high-level Python tools for GPU workloads.
• GPU Acceleration for FiPy Simulations: CuPy and Numba Integration Guide — Explore possible GPU components around a FiPy workflow.
• Physics-Informed Neural Networks for Scientific Simulations — Review the relationship between physical models and differentiable learning.
• When to Use FEM, FVM, and FDM — Select a spatial discretization before optimizing its implementation.
• Performance Profiling: Identifying Bottlenecks in Scientific Code — Locate CPU, memory, transfer, and I/O limitations.
• HPC Python Workflows: From Laptop to Supercomputer — Plan the progression from local development to large systems.

Conclusion

Custom GPU kernels allow scientific developers to translate particles, grid cells, material rules, and other physical operations directly into massively parallel functions. They provide control over indexing, memory access, synchronization, kernel fusion, and device-specific optimization.

CuPy offers both high-level array operations and custom kernel interfaces. Numba-CUDA provides a CUDA-like model in Python, while NVIDIA Warp adds simulation primitives, tile operations, and automatic differentiation. CUDA and HIP provide lower-level C++ control, and performance-portable frameworks support broader hardware targets.

The largest performance gains rarely come from changing syntax alone. They come from choosing a parallel algorithm, keeping data on the device, reducing memory traffic, using a suitable data layout, and eliminating unnecessary intermediate operations.

A custom kernel should be developed only after profiling identifies a real bottleneck. It should be validated against a reference solution and benchmarked with synchronization, representative problem sizes, and complete accounting of memory-transfer costs.

When those conditions are met, Python-based kernel tools can support serious physics simulation without forcing researchers to move every part of the application into low-level C++.