在TPU、GPU和CPU架构上对用于量子启发式均匀化的张量列车进行性能基准测试

《Future Generation Computer Systems》：Performance benchmarking of Tensor Trains for quantum-inspired homogenization on TPU, GPU, and CPU architectures

　　摘要
高分辨率CT成像技术的最新进展带来了新一代超高分辨率微观结构数据集，这类数据集对传统均质化方法的局限性提出了挑战。虽然基于FFT的先进均质化技术仍适用于中等分辨率的数据集，但随着分辨率的提高，其内存占用和计算成本会迅速上升，使得其在处理工业规模问题时效率逐渐降低。为应对这些挑战，近期开发的基于超快傅里叶变换（SFFT）的均质化算法利用了张量列车（TTs）的低秩表示形式，这种表示方式能够有效降低大规模均质化问题的存储和计算需求。该算法专为CPU设计，能够在假设数据表现良好的前提下高效处理高分辨率数据集。在本研究中，我们使用JAX框架在现代硬件加速器上测试了基本TT运算的性能。通过对比CPU、GPU和TPU的运行时间与计算效率，分析了不同架构下TT运算的优势与局限，为未来混合式方法的发展提供了依据。基于这些研究结果，我们对SFFT基均质化算法进行了适配，使其能够在3亿到700亿网格点的高分辨率下实现均质化处理，而现有的最佳GPU基FFT实现则无法达到这一水平。尽管其性能随几何尺度变化而有所不同，但这些结果展示了基于加速器的量子启发式均质化技术在高性能多尺度模拟中的潜力。

引言
均质化理论为预测具有复杂微观结构的材料的宏观行为提供了框架，这类材料包括复合材料、多晶材料以及纤维增强合金[1]–[9]。这一方法的核心是单元问题，即在一个能够反映微观结构特征的代表体积元上定义的偏微分方程。通过求解单元问题，可以确定材料的有效本构行为。然而，对于大多数实际材料而言，解析解并不可行，因此需要借助计算技术[10]–[11]。在这些技术中，基于FFT的均质化方法作为一种高效的替代方案，逐渐取代了传统的有限元方法（FEM）。该方法由Moulinec和Suquet在1998年提出[12]，它将单元问题转化为Lippmann–Schwinger方程，进而可通过快速傅里叶变换（FFT）进行迭代求解。通过在规则网格上离散化微观结构，基于FFT的方法避免了FEM中存在的网格划分和刚度矩阵构建等难题，从而能够快速且精确地处理周期性微观结构[13]。

随着CT成像技术的进步，如今已能够获取规模空前的微观结构数据集，其体素数量可达数百亿级[14]。尽管这类高分辨率数据能提供丰富的信息，但它们也给基于经典FFT的均质化方法带来了巨大的计算挑战——由于FFT的时间复杂度为O(dNdlogN)，因此在d维空间中离散化精度越高，每次迭代的计算时间就会急剧增加[12]–[15]。一种直接的解决方案是使用硬件加速器。不过，虽然计算负荷使得GPU和TPU成为大规模均质化处理的必要工具，但高分辨率数据极高的内存需求往往超出现代加速器的承载能力。虽然已有高效的GPU基FFT库存在[16]，但基于FFT的均质化算法在GPU上的应用通常仅限于小规模问题，或者需要通过批量处理策略来提升高分辨率数据的处理速度，所能带来的加速效果较为有限[17]–[20]。至于TPU，目前尚不存在完整的基于FFT的均质化实现，尽管早期的FFT实现已展现出不错的性能[21]。另一种可能的解决方案是利用量子硬件上的量子傅里叶变换（QFT）来实现指数级的加速，但由于即使是规模较小的系统也需要容错量子计算机，这类应用目前仍处于理论阶段[22]–[23]。

为克服传统基于FFT的均质化方法在计算能力和内存占用方面的限制，最近有人提出了基于张量列车（TT）的均质化算法。基于超快傅里叶变换（SFFT）的均质化算法[24]采用了QFT的TT版本，不仅打破了FFT的指数级增长瓶颈，还大幅降低了内存需求。不过，由于该技术仍处于发展初期，其适用范围仅限于相对简单的几何结构。在本文中，我们做出了两项主要贡献。首先，我们利用JAX框架，在现代硬件加速器，尤其是GPU和TPU上，开发了高效的硬件加速型基本TT代数实现。我们对这些加速后的TT运算的可行性与性能进行了严格评估，分析了不同平台的优势与局限，并给出了包括速度和可扩展性在内的详细基准测试指标。尽管目前已有一些用于GPU加速TT代数的库[25]，但针对TPU的系统性基准测试仍然缺乏；本研究首次完成了此类评估，并直接比较了GPU与TPU的性能，这一工作也建立在之前关于TPU适用于用于解决TT格式优化问题的DMRG算法的研究成果之上[26]。此外，我们还首次在现代硬件加速器上实现了基于极化的TT压缩技术，并对其进行了系统性评估，这项技术对于实现精确的低精度算术运算至关重要。其次，我们将上述开发的硬件加速型TT运算应用到了基于SFFT的均质化流程中，用于在GPU和TPU上处理高分辨率数据。我们的分析不仅探讨了在GPU和TPU上运行基于TT的算法的可行性，还直接比较了它们在计算效率和可扩展性方面的表现。重要的是，张量列车的应用使得原本在传统加速器上无法处理的庞大数据集也能实现均质化处理，这充分体现了低秩张量技术在克服工业规模模拟中的内存和计算瓶颈方面的巨大潜力。

综上所述，这些贡献为在现代加速器上实现高性能、大规模的基于张量的均质化处理奠定了基础，同时证明了基于TT的方法在处理高分辨率数据时的实际可行性和性能优势。需要强调的是，本研究并未对基于SFFT的均质化算法本身进行修改，而是侧重于利用张量列车技术在其现代硬件实现中的高效应用。

本文的其余部分安排如下：第2节介绍张量列车及其基本代数结构；第3节回顾基于FFT的先进均质化方法以及近期开发的基于SFFT的算法；第4节通过对基本TT运算进行严格的性能分析来为其建立基准；第5节讨论了将基于SFFT的均质化算法适配到GPU和TPU上的方法及应用；最后，第6节对全文内容进行总结，并指出主要研究结论。