[论文解读] Petascale XCT: 3D Image Reconstruction with Hierarchical Communications on Multi-GPU Nodes

Mert Hidayetoğlu, Tekin Biçer+7 more|arXiv (Cornell University)|Sep 15, 2020

Medical Imaging Techniques and Applications参考文献 39被引用 24

一句话总结

本论文提出了一套面向多GPU节点的、针对大规模3D X射线计算机断层扫描（XCT）重建的Peta-scale系统，通过分层通信与混合精度计算进行优化。通过将面向内存的2D优化方法扩展至3D，减少节点内通信，并利用混合精度算术，该系统在Summit上使用24,576块GPU实现了9K×11K×11K小鼠脑体积的亚三分钟重建，达到65 PFLOPS（峰值性能的34%）。

ABSTRACT

X-ray computed tomography is a commonly used technique for noninvasive imaging at synchrotron facilities. Iterative tomographic reconstruction algorithms are often preferred for recovering high quality 3D volumetric images from 2D X-ray images, however, their use has been limited to small/medium datasets due to their computational requirements. In this paper, we propose a high-performance iterative reconstruction system for terabyte(s)-scale 3D volumes. Our design involves three novel optimizations: (1) optimization of (back)projection operators by extending the 2D memory-centric approach to 3D; (2) performing hierarchical communications by exploiting "fat-node" architecture with many GPUs; (3) utilization of mixed-precision types while preserving convergence rate and quality. We extensively evaluate the proposed optimizations and scaling on the Summit supercomputer. Our largest reconstruction is a mouse brain volume with 9Kx11Kx11K voxels, where the total reconstruction time is under three minutes using 24,576 GPUs, reaching 65 PFLOPS: 34% of Summit's peak performance.

研究动机与目标

• 解决用于TB级同步辐射数据集的迭代式3D断层成像重建中的计算瓶颈问题。
• 克服现有并行化技术在扩展至极端规模体积与GPU资源时的局限性。
• 实现对大型生物与材料样品的高性能、可扩展重建，具备高保真度与低噪声特性。
• 针对多GPU‘胖节点’架构优化通信与内存访问模式，以最小化延迟并最大化吞吐量。
• 在降低精度以减少内存占用与通信量的同时，保持重建质量，不牺牲收敛性或准确性。

提出的方法

• 通过在GPU节点间采用分层数据与批处理分区方式，优化（反）投影算子，将面向2D的内存中心化（MemXCT）方法扩展至3D。
• 设计针对XCT优化的稀疏矩阵-向量乘法（SpMM）内核，利用空间局部性，并重用共享内存与寄存器中的数据，以最大化GPU利用率。
• 实现分层通信模式：在节点间通信前，先在GPU节点内部执行归约操作，以减轻节点间带宽压力。
• 采用混合精度算术（半精度、单精度及混合精度），以减少内存占用与通信量，由于数值噪声底限远低于测量噪声，因此收敛性得以保持。
• 采用可配置的体积分发策略，支持跨进程与GPU节点的灵活负载均衡。
• 将这些优化整合进可扩展的迭代重建框架中，专为同步辐射设施产生的大规模断层成像数据集而设计。

实验结果

研究问题

• RQ1 能否利用现代多GPU超算系统将3D迭代式断层成像重建扩展至TB级规模的体积？
• RQ2 在‘胖节点’GPU架构上，分层通信模式如何降低大规模重建中的节点间通信开销？
• RQ3 在迭代XCT中，混合精度算术在不降低重建质量或收敛性的情况下，可被使用到何种程度？
• RQ4 面向极端规模3D体积，内存中心化优化技术的3D扩展能否保持高GPU利用率与高性能？
• RQ5 在Summit等Peta-scale系统上，采用这些优化技术后，全3D脑部重建在性能与效率方面可实现多大提升？

主要发现

• 该系统在Summit超算上使用24,576块GPU，实现了9K×11K×11K小鼠脑体积的重建，耗时不足三分钟。
• 重建性能达到65 PFLOPS，相当于Summit峰值性能的34%，展现出优异的强可扩展性效率。
• 混合精度计算在不显著损失重建保真度的前提下，有效减少了内存占用与通信量，同时保持了收敛性与图像质量。
• 通过将归约操作卸载至节点内GPU组，分层通信显著降低了节点间通信量，显著提升了可扩展性。
• XCT优化的SpMM内核通过重用共享内存与寄存器中的数据，最小化全局内存访问，实现了高GPU利用率。
• 由于测量噪声底限远高于数值噪声，所有精度类型（包括半精度与混合精度）的收敛性均保持稳定。