SciPy 2023

Using Numba for GPU acceleration of Neutron Beamline Digital Twins
07-13, 10:45–11:15 (America/Chicago), Amphitheater 204

This talk will discuss how Numba was used to accelerate MCViNE, a software environment for building and running digital twins of neutron experiments via Monte Carlo ray tracing. Numba is an open-source JIT compiler for Python using LLVM to generate efficient machine code for CPUs and GPUs with NVIDIA CUDA. Python and Numba were used to create a GPU accelerated version of MCViNE utilizing an extensible object-oriented design that has achieved a speedup of up to 1000x over the CPU. The performance gain with Numba enables more sophisticated data analysis and impacts neutron scattering science and instrument design.


Motivation

MCViNE is a software package for creating digital twins of neutron scattering experiments using a Monte Carlo ray-tracing approach. These simulations are useful in performing advanced neutron data analysis as well as in designing novel neutron instruments and sample environments. Specifically, it has been used in the initial designs for instruments in the Second Target Station at the Spallation Neutron Source at Oak Ridge National Laboratory. Currently, MCViNE only runs on CPUs which is a bottleneck in large simulations with tens of billions of neutrons, and in modelling complex multiple scattering, with some simulations taking months to complete. Due to the massively parallel nature of Monte Carlo methods, bringing GPU acceleration to these simulations would offer superior performance and scalability. MCViNE is originally implemented in C++ and parallelized using MPI, and it has bindings to Python for user interaction; however, extensibility for the user can be very difficult.

Methods

To improve performance and to ease user contributions, Python and Numba were chosen to create a new package providing GPU acceleration of MCViNE. Numba is an open-source JIT (just-in-time) compiler for Python using LLVM to generate efficient machine code and supports GPUs using NVIDIA CUDA. Numba is designed for scientific computing and can support NumPy arrays and functions. Currently, we are only using Numba for its GPU capabilities.

Using Python and Numba for this application allowed several advantages such as utilizing an extensible object-oriented approach and polymorphism. Each MCViNE instrument is composed of several components, such as a neutron source, a guide, and a monitor. During the simulation, neutrons can travel through each component in the instrument. Each component has a method (“propagate”) defined for propagating the neutron through it. Additionally, sample environments are created using constructive solid geometry (CSG) with each primitive shape defined as a CUDA kernel. To each constructed shape, many CUDA kernels are available, each modeling a different type of scattering physics. Due to different component/scattering-kernel types and geometric shapes, using an object-oriented design was beneficial. Furthermore, this structure allowed for custom on-the-fly CUDA kernel generation for complex instrument/sample geometries and physics.

Results and Conclusions

Python and Numba was used to create a GPU accelerated version of MCViNE, which has so far achieved speedups of up to 1000x over the original CPU implementation. This performance gain enables more sophisticated data analysis for neutron scattering and impacts neutron scattering science and instrument design.
Using Python has helped increase the usability, extensibility, and maintainability of the codebase. Additionally, coupling Python with Numba allowed complex combinations of CUDA kernels to be generated at runtime, which would have been significantly harder to implement in other languages. The techniques used in this project could also be applied to other scientific computing applications.

Resources:

https://github.com/mcvine/acc
https://mcvine.ornl.gov/
https://github.com/mcvine/mcvine

Research Software Engineer in the Application Engineering group at Oak Ridge National Laboratory.