But What Is Inside an AI Accelerator?
Source: tai_blog

Heterogeneous computing refers to machines with more than one kind of computing core. The computing cores can be CPUs  GPUs  TPUs  and many other accelerators that are being developed every day. These specialized cores can also be called ASIC an abbreviation for Application-Specific Integrated Circuit.   This is how ARM defines ASIC    An application-specific integrated circuit is an integrated circuit (IC) thats custom-designed for a particular task or application. Unlike FPGA boards that can be programmed to meet a variety of use case requirements after manufacturing  ASIC designs are tailored early in the design process to address specific needs.   Since the release of ChatGPT and the subsequent release of other large language models (LLM)  there has been a growing demand for computing power that is required to train these models (with billions of parameters) and also generate results  which is called inferencing. This is precisely where AI Accelerators come to the rescue!   An overview of what lies ahead in this article   In this article  I will go over a small introduction to AI accelerators and how they differ from CPUs and GPUs. Then I will dive into systolic array architecture and how it works! I also peek inside the Google TPU and end the article with possible future research directions.   Introduction to AI AcceleratorsAI accelerators are specialized hardware designed to enhance the performance of artificial intelligence (AI) tasks  particularly in machine learning and deep learning. These accelerators are designed to perform large-scale parallel computations (read matrix multiplications) as required by many deep learning models efficiently as compared to traditional CPUs.   Some key characteristics that differentiate AI Accelerators from CPUs and GPUs are:   They are a type of ASIC specifically designed for deep learning workloads. In contrast  CPUs and GPUs can also be used for general-purpose programming and rendering graphics  respectively. NVIDIA GPUs in fact  started out as ASIC for handling computer graphics-related operations and then transitioned into being used in scientific computing (with the help of CUDA). Sometime later  around 2015  the focus of CUDA transitioned towards supporting neural networks.Massive parallel processing power  GPUs and accelerators are designed to execute many operations in parallel (high throughput)  whereas CPUs are designed to perform sequential operations in the shortest time (low latency). Accelerators are meant to offload deep learning workloads from CPUs so as to perform these operations more efficiently.Systolic ArraysSystolic array is a simple and energy-efficient architecture for accelerating general matrix multiplication (GEMM) operations in hardware. They provide an alternative way to implement these operations and support parallel data streaming to improve memory access and promote data reuse. This architecture forms the basis of many commercial accelerator offerings like the Google TPU (tensor processing unit)  Intel NPU (neural processing unit)  IBM AIU  etc.   These arrays comprise MAC (multiply-and-accumulate) units that perform the actual operations. Serving the MAC units are the row and column SRAM buffers that feed these units with data. Each MAC unit will save the incoming data in an internal register and then forward the same data to the outgoing connection in the next cycle.   This behavior results in significant savings in SRAM read requests and can exploit data reuse opportunities. For example  filter weights are something that remains stationary during a convolution operation as the filter map is convolved over the image. This can be exploited by storing the weights in the MAC array whereas the row buffer loads in the different windows of the input image. This reduces the read requests to load the weights  hence freeing up bandwidth to read from off-chip memory sources like DRAM or HBMs.   There are different techniques to exploit data reuse  which are referred to as dataflow or mapping schemes discussed in the next section.   Data Flow TechniquesAlthough there are no hard and fast rules to specify what kind of mapping is to be used with a systolic array architecture  here I will discuss one of the three strategies as specified in the Scale-Sim paper. The three strategies are named Output Stationary (OS)  Weight Stationary (WS)  and Input Stationary (IS). The word stationary here depicts what part of the computation spends the most amount of time being stored in the systolic array.   The output stationary dataflow is depicted in the figure above. Output stationary means that each MAC unit will be responsible for calculating the output pixel. All the required operands are fed from the left and top edges of the systolic array. Each row (IFMAP) consists of elements of one convolution window and one column (FILTER) entering from the top represents the unrolled filter. Elements of one row and one column are multiplied and accumulated to calculate one pixel of the output feature map (OFMAP).   Timing Model for a Systolic ArrayHere we try to calculate the number of cycles that a systolic array will take to perform a matrix multiplication. Here we assume that there are no stalls during the operation due to memory bandwidth (make sure that SRAM buffers are filled with data to perform the compute) and also assume that we have unlimited MAC units available to perform the required computation.   Sr  Sc are the dimensions of the systolic array and in this case  is equivalent to the number of rows and columns of the IFMAP and FILTER respectively. T is the temporal dimension which in the case of the output stationary represents the convolution window size.   As described by the figure above  we can conclude that the number of cycles for the systolic array to perform a matrix multiplication is:   Obviously  in the real world  we do not have unlimited MACs. In that case  we divide the workload by the number of available MAC units and therefore get the following expression for timing:   Here  we assume that R and C are the actual dimensions of the systolic array and Sr and Sc are the required dimensions. To decrease this time  we can increase the number of MAC units  a process we can call scaling up. Another approach is to have multiple MAC array units that perform the compute in parallel  which can be called scaling out. This further reduces the time needed to complete the operation.   A look inside Google TPUOriginsBack in 2013  a projection at Google showed that if people searched using voice even for 3 minutes a day  it would result in doubling the computing demand of Googles datacenters. Speech recognition models that used DNN were very expensive to perform inference using traditional CPUs. Therefore  they started working on a custom ASIC (application-specific integrated circuit) that would perform inference efficiently. The goal was 10x performance over GPUs. The outcome of this effort was the Google Tensor Processing Unit. Google TPU was based on the systolic array architecture.   TPU v1As you are now aware systolic array-based AI accelerators are composed of MAC units. Googles original TPU implementation consisted of 256x256 MAC units (see Matrix Multiply Unit in the figure above) that could perform 8-bit multiply-and-adds on signed or unsigned integers. The 16-bit products were then collected in 4 MiB of 32-bit Accumulators below the matrix unit. Then there are other components like the activation pipeline that could perform activation functions on the resulting matrix.   For more details about the Google TPU that was released in 2017 read this very interesting paper where they discuss in detail the TPUs design and performance!   In-datacenter performance analysis of a tensor processing unit U+007C IEEE Conference Publication U+007C IEEE Xplore   TPU v2 and v3Improving upon the design of TPU v1 Google released the specifications of TPU v2 and v3 as well with some major changes:   Interconnect  A critical element of any chip design is the interconnect which decides how fast is the inter-chip communication. An on-device switch called Interconnect Router (see above figure) provides deadlock-free routing. It enables a 2D torus topology of interconnect.Memory  A major performance bottleneck in TPU v1 was the limited memory bandwidth of DRAM. This problem was somewhat solved using the HBM (High Bandwidth Memory) DRAM in TPU v2. It offers 20 times the bandwidth of TPU v1 by using an interposer substrate that connects the TPU v2 chip via thirty-two 128-bit buses to 4-stacks of DRAM chips.Multiple smaller MXU units per chip  While TPUv1 featured a MXU of the size 256x256  it was reduced to 128x128 in TPUv2 onwards and has multiple MXUs per chip. Larger MXUs require more memory bandwidth for optimal chip utilization. Google analyzed that convolutional model utilization ranged between 37%-48% for 128x128 MXUs  which was 1.6x of a single 256x256 MXU (22%-30%). The reason that Google has come up with this is that some convolutions are naturally smaller than 256x256  which leaves parts of the MXU unused.For more details regarding Google TPU v2 and v3:   A Domain Specific Supercomputer for Training Deep Neural Networks U+007C ACM   AI and Memory WallThe amount of computing needed to train modern deep learning models and perform inference using them is growing at a large rate. This trend prompted research into AI accelerators with a focus on increasing computing power. This has been achieved sometimes at the expense of neglecting memory hierarchies and bandwidth thus creating a memory bottleneck. In this section  I have briefly summarized what this very interesting paper [Gholami et al.  2024] talks about and which points toward future research avenues in the realm of AI accelerators.   But what is a memory wall?   Memory wall refers to the problem where the compute is faster than the rate at which data can be fetched from off-chip DRAM which limits the overall compute that can be performed. The time to complete an operation is dependent both on the speed of performing compute and also on how fast the data can be fed to the arithmetic units of hardware.   As can be seen in the graph above  the peak compute has increased 60000x in the last 20 years  whereas the DRAM and interconnect bandwidth have increased only 100x and 30x  respectively. This huge deficit results in aggravating the problem of memory wall especially with growing model sizes.   As depicted in figure (a) above the number of parameters in the SOTA transformer models has increased at a rate of 410x every two years  whereas the AI accelerator memory capacity (green dots) has only been scaled at a rate of 2x every 2 years. Figure (b) depicts the amount of compute  measured in Peta FLOPs  needed to train SOTA models for different computer vision (CV)  natural language processing (NLP)  and Speech models  along with the different scaling of Transformer models (750x/2yrs).   This problem opens up many research avenues where progress can be made. Techniques like quantization and model pruning are being actively investigated to reduce model size. One of the major breakthroughs in AI accelerators has been the successful adoption of half-precision (FP 16) instead of single precision enabling a 10x increase in hardware compute capability. Another possible solution that the author proposes worth investigating is revisiting the organization of the cache hierarchy of AI Accelerators that has been simplified to prioritize computing power.   Do check out the paper by the author for a more detailed analysis and discussion on this topic: [2403.14123] AI and Memory Wall (arxiv.org)   Further ReadingDNN Accelerator Architecture  SIMD or Systolic? U+007C SIGARCHArchitecting Chips For High-Performance ComputingHow To Build A Better Blackwell GPU Than NVIDIA DidExtending Dataflow Techniques from Dense to Sparse Accelerators U+007C SIGARCHReferencesJouppi  N. P.  Young  C.  Patil  N.  Patterson  D.  Agrawal  G.  Bajwa  R.   & Yoon  D. H. (2017  June). In-datacenter performance analysis of a tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture (pp. 112).Jouppi  N. P.  Yoon  D. H.  Kurian  G.  Li  S.  Patil  N.  Laudon  J.   & Patterson  D. (2020). A domain-specific supercomputer for training deep neural networks. Communications of the ACM  63(7)  6778.Gholami  A.  Yao  Z.  Kim  S.  Hooper  C.  Mahoney  M. W.  & Keutzer  K. (2024). AI and memory wall. IEEE Micro.Samajdar  A.  Joseph  J. M.  Zhu  Y.  Whatmough  P.  Mattina  M.  & Krishna  T. (2020  August). A systematic methodology for characterizing scalability of dnn accelerators using scale-sim. In 2020 IEEE International Symposium on Performance Analysis of Systems and Software (ISPASS) (pp. 5868). IEEE.