摘要

本文介绍了一种名为Flex-TPU的新型张量处理单元（TPU），该单元具有运行时可重配置的数据流架构。传统的TPU虽然在机器学习（ML）加速方面表现出色，但其固定的数据流架构限制了其在深度神经网络（DNN）推理中的性能。Flex-TPU通过在运行时动态调整每一层的数据流，显著提高了性能，同时仅增加了轻微的面积和功耗开销。这一创新为数据中心和边缘计算环境中的ML应用提供了更高的灵活性和效率。

原理

Flex-TPU的核心在于其处理元素（PE）的微架构改进，通过增加额外的寄存器和多路复用器（MUX），实现了运行时数据流的可重配置。这种设计允许PE在输入静止（IS）、输出静止（OS）和权重静止（WS）三种数据流模式之间动态切换，以适应不同DNN层的计算需求。通过配置管理单元（CMU）和数据流生成器，Flex-TPU能够为每一层选择最优的数据流策略，从而最大化数据重用并最小化数据传输，显著提升整体性能。

流程

Flex-TPU的工作流程包括以下几个关键步骤：首先，CMU根据DNN层的特性选择最优的数据流模式；其次，数据流生成器根据选定的模式生成内存读写地址；然后，主控制器管理内存和FIFO缓冲区之间的数据传输；最后，PE根据配置进行乘积累加（MAC）操作。例如，在ResNet-18模型的测试中，Flex-TPU能够根据每一层的最佳数据流模式动态调整，从而实现比传统TPU高达2.75倍的性能提升。

应用

Flex-TPU的设计不仅适用于数据中心的大型ML加速任务，也适用于边缘计算设备，如Google的Coral Edge TPU。其运行时可重配置的数据流架构使得Flex-TPU在处理多样化和高动态的ML工作负载时具有显著优势，预计将在未来的AI硬件加速器市场中占据重要地位。