Flex-TPU: A Flexible TPU with Runtime Reconfigurable Dataflow Architecture

TL;DR

The paper introduces Flex-TPU, a reconfigurable dataflow architecture for TPUs that dynamically adapts to different layers, significantly boosting performance over traditional fixed dataflow TPUs with minimal overhead.

Contribution

Develops the first runtime reconfigurable dataflow TPU, enabling dynamic dataflow changes per layer to optimize performance.

Findings

Achieves up to 2.75x performance improvement over conventional TPU.

Maintains minimal area and power overheads.

Validates effectiveness across multiple ML workloads.

Abstract

Tensor processing units (TPUs) are one of the most well-known machine learning (ML) accelerators utilized at large scale in data centers as well as in tiny ML applications. TPUs offer several improvements and advantages over conventional ML accelerators, like graphical processing units (GPUs), being designed specifically to perform the multiply-accumulate (MAC) operations required in the matrix-matrix and matrix-vector multiplies extensively present throughout the execution of deep neural networks (DNNs). Such improvements include maximizing data reuse and minimizing data transfer by leveraging the temporal dataflow paradigms provided by the systolic array architecture. While this design provides a significant performance benefit, the current implementations are restricted to a single dataflow consisting of either input, output, or weight stationary architectures. This can limit the achievable performance of DNN inference and reduce the utilization of compute units. Therefore, the work herein consists of developing a reconfigurable dataflow TPU, called the Flex-TPU, which can dynamically change the dataflow per layer during run-time. Our experiments thoroughly test the viability of the Flex-TPU comparing it to conventional TPU designs across multiple well-known ML workloads. The results show that our Flex-TPU design achieves a significant performance increase of up to 2.75x compared to conventional TPU, with only minor area and power overheads.