Porting GLASS to the Python Array API: A Clear View to GPU Portability

Abstract

GLASS is a Python library designed to generate full-universe cosmological simulations for large-scale galaxy surveys. Historically, the use of NumPy has been essential to GLASS’ success, allowing rapid development and integration with a wide range of existing cosmological software. However, as simulation demands grow and high-performance computing shifts towards GPU-heavy architectures, GLASS must evolve to support these workloads. As NumPy was not designed for non-CPU devices, it now poses a significant bottleneck to this evolution. To overcome this, we are porting GLASS to the Python Array API, decoupling the library from NumPy and enabling native GPU execution. By adopting this standard, we are expanding GLASS to support all Array API-compliant libraries, broadening its ecosystem and utility. This transition involves managing varying levels of Array API maturity across libraries and optimising core scientific models for GPU parallelism. The resulting GPU-enabled GLASS provides a significant performance boost whilst remaining performant on CPUs, facilitating seamless integration with GPU-native workflows, and enabling advanced techniques like auto-differentiable simulations via JAX.