Beschreibung
GPUs offer a way to massively parallelise certain types of physical models, allowing us to compare more models for the same compute budget (and hence do more science) whilst avoiding the use of emulators (which must be retrained every time you modify the underlying physics). More efficient forward models also allow us to tackle higher-dimensional problems (map/field-level inference) and more complex foreground systematics and nuisance parameters.
In particular, the so-called Einstein-Boltzmann equations that underlie models of the CMB are linear ordinary differential equations that are highly amenable to parallelisation on a GPU, across both wave modes and parameters. Meanwhile, modern cosmological analyses are burdened with the technical debt of low-level codes written in FORTRAN and C, which are difficult to modify when trying to test new models of physics.
We have therefore written a new CMB code in Python and JAX (DISCO-EB) which can take advantage of the massive computational throughput of modern GPUs, but is easy to modify with new physics. As part of this work, we have developed new ODE solvers that are better optimised for solving many small ODEs at once and are 100x faster than the current Python/JAX state-of-the-art (diffrax). The result is precision similar to that of existing codes but at vastly reduced $ inference cost. The ODE solvers themselves are packaged separately and could have broad application in the scientific community.
This deep dive into the world of GPU programming has taught us that simply "Jaxifying" one's code is not always the best solution, and that it pays to understand how a GPU works under the hood. Moreover, as GPU manufacturers begin to optimise their chips for a deep learning community that prefers lower precision matrix multiplication operations, the scientific community must adapt its codes to make use of GPUs that have fewer 64 bit precision cores but more "tensor cores". We present a number of examples of future-proofing to this effect.