Performance Portable Monte Carlo Neutron Transport in MCDC via Numba

TL;DR

This paper presents a portable, high-performance Monte Carlo neutron transport application built with Python and Numba, capable of running efficiently on CPUs and GPUs, with demonstrated scalability and competitive performance.

Contribution

The authors introduce a novel portability scheme using Numba for Python, enabling rapid development and execution of Monte Carlo neutron transport on diverse hardware.

Findings

Achieved scalable performance up to 256 nodes.

Demonstrated competitive GPU and CPU performance compared to C++ codes.

Enabled flexible execution modes with a unified Python-based framework.

Abstract

Finding a software engineering approach that allows for portability, rapid development, and open collaboration for high-performance computing on GPUs and CPUs is a challenge. We implement a portability scheme using the Numba compiler for Python in Monte Carlo / Dynamic Code (MC/DC), a new neutron transport application for rapidly developing Monte Carlo. Using this scheme, we have built MC/DC as an application that can run as a pure Python, compiled CPU, or compiled GPU solver. In GPU mode, we use Numba paired with an asynchronous GPU scheduler called Harmonize to increase GPU performance. We present performance results (including weak scaling up to 256 nodes) for a time-dependent problem on both CPUs and GPUs and compare favorably to a production C++ code.