GPU-Accelerated Spectral Methods in Python: A Case Study

Jamie Quinn

University College London

7th SeptembRSE 2021

Melvin.py

A pseudo-spectral, GPU-accelerated framework for numerically solving 2D problems in fluids.

github.com/jamiejquinn/melvin.py

Why Python?

  • Easy prototyping
  • Fast iteration cycle
  • Collective good practice
  • Access to accelerated libraries
  • Students want to learn & use Python

It’s a great language for prototyping research codes & teaching

Easy prototyping

  • spectral/FDM spatial discretisation
  • periodic/Dirichlet boundary conditions
  • 2nd/4th order spatial derivatives for nonlinear advection operator
  • 2nd/4th order Adams-Bashforth time integration
  • Implicit handling of diffusion operator

Designed to be extensible

Fast iteration cycle

  1. Jupyter notebook
  2. Collection of scripts
  3. Fully-fledged Python package

Permits refinement of methods and interfaces

Collective good practices

  • Version control with git & Github
  • Parameterisation with JSON
  • Community contributions through Github
  • Automated with Github Actions
    • unit testing with pytest
    • static analysis with Flake and Black

Acceleration: CuPy

#import numpy as np
import cupy as np

spectral_arr = np.fft(physical_arr)
  • GPU-accelerated Numpy
  • Access to multi-GPU
  • Easy integration with other libraries
  • Can profile using NVIDIA’s Nsight

Alternatives to CuPy

  • Jax - Google’s numpy replacement (differentiation!)
  • Legate - NVIDIA’s numpy replacement
  • CUDA Python - NVIDIA’s low-level
  • PyCUDA - ind low-level
  • Numba - JIT compiler (+ CUDA kernels)

Either interchangeable or interoperable

Results