Performance portable batched linear algebra kernels for transport sweeps using Kokkos
Gabriel Suau, Thierry Gautier, Ansar Calloo, Rémi Baron, R. Le Tellier · 2025
This paper describes the development of performance portable batched linear algebra kernels for SN-DG neutron transport sweeps using Kokkos. We establish a new sweep algorithm for GPUs that relies on batched linear algebra kernels. We implement an optimized batched gesv solver for small linear systems that builds upon state-of-the-art algorithms. Our implementation achieves high performance by minimizing global memory traffic and maximizing the amount of computations done at compile-time. We assess the performance of the batched gesv kernel on NVIDIA and AMD GPUs. We show that our custom implementation outperforms state-of-the-art linear algebra libraries on these architectures. The performance of the new GPU sweep implementation is assessed on the H100 and MI300A GPUs. We demonstrate that it is able to achieve high performance on both architectures, and is competitive with an optimized multithreaded CPU implementation on a 128-core AMD Genoa CPU node.