A multi-GPU design for large size cryo-EM 3D reconstruction

Three-dimensional (3D) reconstruction of cryo-electron microscopy (cryo-EM) is a powerful method to determine the structures of macromolecules at near-atomic resolution. Recently, larger size with finer resolution 2D images has been collected, which can improve the reconstruction resolution. However, large size data incurs high computation and huge memory overhead. Current implementations fail to perform the complete reconstruction workflow on a multi-GPU cluster for large size data. Because of no effective parallel method for 3D convolution and the huge memory demanding, large size data can not be efficiently reconstructed, which impede the resolution improving 3D reconstruction. To enable cryo-EM 3D reconstruction with large size data on multi-GPU, in this work, we propose a new parallel framework called OML-Relion. In OML-Relion, we first adopt a stride based Fourier transform and eliminate data dependence to parallelize the 3D convolution on multi-GPU. Considering the input size varying in each iteration, we next use an auto-tuning model to optimize 3D convolution performance. Finally, guaranteeing the whole reconstruction on a multi-GPU cluster for large size data, we design a novel lossless data compression algorithm to reduce memory overhead on each GPU further. The experiment shows that OML-Relion can efficiently handle large size cryo-EM 3D reconstruction on multi-GPU. The reconstruction module, including 3D convolution operation, achieves 225-330x times speedup for 200-800 pixel size particles. The compression algorithm significantly reduces memory overhead approaching 70%. Moreover, the whole workflow with OMLRelion can achieve 54-65x speedup compared with Relion using two large size datasets.