[论文解读] Large-Scale Discrete Fourier Transform on TPUs

Tianjian Lu, Yifan Chen+3 more|arXiv (Cornell University)|Feb 9, 2020

Parallel Computing and Optimization Techniques参考文献 34被引用 11

一句话总结

本论文提出两种并行算法——基于Kronecker积与密集矩阵乘法的KDFT，以及基于Cooley-Tukey算法并引入相位调整的FFT——用于在TPU集群上进行大规模3D离散傅里叶变换。两种算法均在TensorFlow中实现，充分利用TPU的矩阵乘法能力及一种单次shuffle通信机制，在完整的2048核TPU Pod上实现了8192³ DFT仅需8.3秒，展现出极高的并行效率与极低的通信开销。

ABSTRACT

In this work, we present two parallel algorithms for the large-scale discrete Fourier transform (DFT) on Tensor Processing Unit (TPU) clusters. The two parallel algorithms are associated with two formulations of DFT: one is based on the Kronecker product, to be specific, dense matrix multiplications between the input data and the Vandermonde matrix, denoted as KDFT in this work; the other is based on the famous Cooley-Tukey algorithm and phase adjustment, denoted as FFT in this work. Both KDFT and FFT formulations take full advantage of TPU's strength in matrix multiplications. The KDFT formulation allows direct use of nonuniform inputs without additional step. In the two parallel algorithms, the same strategy of data decomposition is applied to the input data. Through the data decomposition, the dense matrix multiplications in KDFT and FFT are kept local within TPU cores, which can be performed completely in parallel. The communication among TPU cores is achieved through the one-shuffle scheme in both parallel algorithms, with which sending and receiving data takes place simultaneously between two neighboring cores and along the same direction on the interconnect network. The one-shuffle scheme is designed for the interconnect topology of TPU clusters, minimizing the time required by the communication among TPU cores. Both KDFT and FFT are implemented in TensorFlow. The three-dimensional complex DFT is performed on an example of dimension $8192 imes 8192 imes 8192$ with a full TPU Pod: the run time of KDFT is 12.66 seconds and that of FFT is 8.3 seconds. Scaling analysis is provided to demonstrate the high parallel efficiency of the two DFT implementations on TPUs.

研究动机与目标

• 为科学计算工作负载在TPU集群上实现高性能、大规模离散傅里叶变换（DFT）。
• 通过Kronecker积（KDFT）和Cooley-Tukey FFT（FFT）将DFT表述为密集矩阵运算，以充分利用TPU在矩阵乘法方面的优势。
• 设计一种针对TPU互连拓扑结构的通信高效型数据分解与核心间通信方案。
• 在数千个TPU核心上实现3D DFT的强可扩展性与高并行效率。
• 展示TPU上DFT加速在真实世界科学计算应用中的可行性与性能表现。

提出的方法

• 利用Kronecker积将3D DFT表述为输入数据与Vandermonde矩阵之间的矩阵乘法（KDFT），从而直接支持非均匀输入。
• 通过Cooley-Tukey算法结合相位调整实现FFT，将计算复杂度从O(N²)降低至O(N log N)。
• 应用数据分解策略，将密集矩阵乘法局部化在单个TPU核心内执行，从而在集群中实现完全并行化。
• 采用单次shuffle通信机制，使同一互连方向上的相邻TPU核心能够同时进行发送与接收操作，最大限度降低通信延迟。
• 将KDFT与FFT均映射至TensorFlow，实现在TPU Pod上的可移植、可扩展执行，利用einsum与collective_permute等硬件优化操作。
• 利用TPU的高带宽内存（128 GiB）与直接芯片间互连，最小化大规模数据传输中的瓶颈。

实验结果

研究问题

• RQ1 能否利用矩阵乘法原语在TPU集群上高效并行化基于Kronecker积的DFT表述？
• RQ2 在大规模3D问题上，基于FFT的DFT在TPU上的性能与直接矩阵乘法（KDFT）表述相比如何？
• RQ3 单次shuffle通信机制在大规模DFT计算中在多大程度上最小化了跨TPU通信开销？
• RQ4 在包含最多2048个核心的完整TPU Pod上，3D DFT的强可扩展性表现如何？
• RQ5 基于TPU的DFT实现能否在科学计算工作负载中实现高并行效率与低延迟？

主要发现

• 在完整的2048核TPU Pod上，8192×8192×8192的3D FFT运行时间为8.30秒，显著优于相同问题下KDFT的12.66秒运行时间。
• 对于2048×2048×2048的问题，KDFT与FFT的计算时间几乎相同（分别为0.120秒与0.118秒），表明在大规模下TPU的矩阵乘法性能主导了算法复杂度差异的影响。
• 强可扩展性分析显示，对于2048³问题的3D FFT，性能在128个核心以内接近线性扩展，观察到接近理想的加速比。
• 当核心数超过128个时，通信时间（主要由all_to_all与collective_permute操作主导）成为瓶颈，表明需要进一步优化大规模通信模式。
• 3D KDFT与FFT实现展现出高并行效率，当按算法复杂度归一化后，计算时间与问题规模近似呈线性关系。
• 单次shuffle通信机制通过在同一互连方向上同时支持发送与接收操作，有效最小化了跨TPU通信延迟。