Presentation

· Contributors · Organizations · Search Program · My Schedule · Happening Now · Maps

This content is available for: Workshop Reg Pass. Upgrade Registration

High-Level GPU Code: A Case Study Examining JAX and OpenMP

Session2023 International Workshop on Performance, Portability, and Productivity in HPC (P3HPC)

DescriptionIn recent years, a new scientific software design pattern has emerged that pairs a Python interface with high-performance kernels in lower-level languages. The rise of general-purpose GPUs necessitates the rewriting of many such kernels, posing challenges in GPU programming and ensuring future portability and flexibility.

This paper investigates the use of high-level frameworks that abstract system architecture details, aiming for straightforward, portable yet performant GPU code. We focus on TOAST, a cosmology software framework designed to take full advantage of a supercomputer, and compare using the JAX Python library with OpenMP target offload compiler directives as porting strategies. While JAX allows kernel code to be written in pure Python, OpenMP target offload is a directive-based strategy that integrates seamlessly with our existing OpenMP-accelerated C++ kernels.

We port a dozen kernels, analyzing development cost, performance, and the viability of using either framework for complex numerical Python applications.

Author/Presenters