Enabling GPU Portability into the Numba-JITed Monte Carlo Particle Transport Code MC/DC

TL;DR

This paper presents MC/DC, a portable Monte Carlo neutron transport code with GPU acceleration via JIT compilation, demonstrating significant speedups on various GPU architectures compared to CPU performance.

Contribution

We develop a GPU-compatible, JIT-compiled Monte Carlo neutron transport code, enabling cross-platform GPU acceleration from Python source code.

Findings

15x speedup on Nvidia V100 for benchmark problem

3-4x speedup on AMD GPUs for continuous-energy simulations

Effective GPU portability achieved across Nvidia and AMD hardware

Abstract

The Center for Exascale Monte Carlo Neutron Transport is developing Monte Carlo / Dynamic Code (MC/DC) as a portable Monte Carlo neutron transport package for rapid numerical methods exploration on CPU- and GPU-based high-performance computers. In this paper, we describe MC/DC's current event-based GPU algorithm as well as the just-in-time (JIT) compilation scheme we use to enable GPU operability on Nvidia and AMD GPUs from MC/DC's Python source. To analyze performance, we conduct runtime tests of the C5G7 k-eigenvalue benchmark problem and a continuous-energy infinite pin cell on Nvidia Tesla V100 GPU, AMD MI250X GPU, and the AMD MI300A APU and make comparison to a dual-socket Intel Xeon Sapphire Rapid CPU node. We found that for the multi-group C5G7 benchmark problem, we respectively see a 15$\times$, 0.7$\times$, 12$\times$ speedup on a V100, MI250X, and MI300A over 112 Intel Xeon CPU cores. For the continuous-energy infinite pin-cell benchmark, we found speedups of 5$\times$, 3$\times$, 4$\times$ on a V100, MI250X, and MI300A, respectively, over the same CPU node.