Software Stack for General-Purpose Tensor Processing Units

点击翻译按钮获取中文摘要

The widespread adoption of neural networks in many application areas has created the market conditions for the development of specialized architectural solutions for the faster computation of tensor algebraic operations. These operations are at the core of both training and inference for deep neural networks. The industry has responded with a diverse offering of architectural solutions. Known as Tensor Processing Units (TPUs), Neural Network Processors, Deep-Learning Accelerators, Artificial Intelligence Chips, such architectural solutions are often based on a multi-processing-unit topology called systolic arrays and have the common goal of accelerating multi-dimensional array computations. A key desirable feature is to increase the reuse of data that has been brought from memory into the processing units. The premise of this innovative research program is that the architectural constructs developed for TPUs can also be deployed for general-purpose numerical computing applications. Similarly to what happened in the evolution of GPU design, the broader utilization of TPUs will lead to the evolution of these accelerators. This evolution will, in turn, make them also better-suited for neural networks themselves. While many programming solutions -- languages and libraries -- and compilers exist for TPUs, the development of a software stack for neural processing is still incipient. Recently Chris Lattner, the creator fo the widely used LLVM compilation infrastructure, has proposed the Multi-Level Intermediate Representation (MLIR) that has the potential of enabling more efficient and more effective optimizations at a higher level in the compilation process. Google has an R&D team working on the development of the MLIR infrastructure that is needed for its adoption. However, all the initial effort is focused on a replacement solution for TensorFlow, the current software stack offered by Google. The proposed research expands the software stack for TPUs by creating compilation paths from traditional High-Performance Computing (HPC) programming languages, such as C or Fortran, potentially augmented with OpenMP directives, for execution in TPUs. Main goals: 1. Develop program analysis to discover program segments (loop nests) in numerical computing that can benefit from execution in the architectural blocks available in TPUs. 2. Create profitability analysis that can be used to estimate the performance, and power consumption, when a computation is executed in a TPU. These profitability analysis will also be used to determine which code transformations could be applied to the code to improve its performance in a TPU. 3. Develop compile-time code transformations that can make a given loop next more suitable for execution in TPU hardware. 4. Extend existing IRs to better improve the mapping of numerical computations to TPUs.