The recent trend toward deep learning has led to the development of a variety of highly innovative AI accelerator architectures. One such architecture, the Cerebras Wafer-Scale Engine 2 (WSE-2), features 40 GB of on-chip SRAM, making it a potentially attractive platform for latency- or bandwidth-bound HPC simulation workloads. In this study, we examine the feasibility of performing continuous energy Monte Carlo (MC) particle transport on the WSE-2 by porting a key kernel from the MC transport algorithm to Cerebras's CSL programming model. New algorithms for minimizing communication costs and for handling load balancing are developed and tested. The WSE-2 is found to run 130 times faster than a highly optimized CUDA version of the kernel run on an NVIDIA A100 GPU -- significantly outpacing the expected performance increase given the difference in transistor counts between the architectures.
@article{arxiv.2311.01739,
title = {Efficient Algorithms for Monte Carlo Particle Transport on AI Accelerator Hardware},
author = {John Tramm and Bryce Allen and Kazutomo Yoshii and Andrew Siegel and Leighton Wilson},
journal= {arXiv preprint arXiv:2311.01739},
year = {2023}
}