PyKokkos: Performance Portable Kernels in Python

As modern supercomputers have increasingly heterogeneous hardware, the need for writing parallel code that is both portable and performant across different hardware architectures increases. Kokkos is a C++ library that provides abstractions for writing performance portable code. Using Kokkos, programmers can write their code once and run it efficiently on a variety of architectures. However, the target audience of Kokkos, typically scientists, prefers dynamically typed languages such as Python instead of C++. We demonstrate a framework, dubbed PyKokkos, that enables performance portable code through Python. PyKokkos transparently translates code written in a subset of Python to C++ and Kokkos, and then connects the generated code to Python by automatically generating language bindings. PyKokkos achieves performance comparable to Kokkos in ExaMiniMD, a ~3k lines of code molecular dynamics mini-application. The demo video for PyKokkos can be found at https://youtu.be/1oFvhlhoDaY.