基于 AI 硬件和 GPU 的逆重叠矩阵高效混合精度矩阵分解

近年来，一种新型加速硬件在人工智能（AI）和机器学习（ML）社区中广受欢迎，这些硬件支持在降低精度下实现极高性能的张量收缩，用于深度神经网络计算。在本文中，我们利用 Nvidia Tensor 核心——这种 AI/ML 硬件的典型示例——开发了一种用于计算电子结构理论中逆重叠矩阵 $S^{-1}$ 的混合精度密集矩阵分解方法。该 $S^{-1}$ 的分解形式为 $ZZ^T=S^{-1}$，用于将通用矩阵特征值问题转化为标准矩阵特征值问题。我们提出了一种混合精度迭代细化算法，其中 $Z$ 通过矩阵-矩阵乘法递归计算，可在 Tensor 核心上实现高性能。为了解 Tensor 核心的性能和准确性，将其与仅使用 GPU 的单精度和双精度实现进行比较。此外，我们提出一种非参数停止准则，能够鲁棒地应对低精度浮点运算。该算法在我们有一个关于 $Z$ 的良好初始猜测时尤其有用，例如来自量子力学分子动力学模拟中前一步的时间步骤或几何优化中的前一次迭代。