TPU内存（一）

首先我们来看一下TPU的简要架构。

我们可以从ppt中了解到一个TPU中包含多个NPU（Neuron Processing Unit），主要由一个local memory和多个执行单元组成。前者用于存储要运算的数据，后者是TPU上最小的计算单元。每个NPU一次可以驱动它的所有EU做一个MAC操作。

就整体 TPU 内存而言，它由system memory和local memory组成。 system memory的主要部分是global memory，其实就是一块DDR。 有时根据 TPU 的特殊设计还会有其他组件，但我们不会在视频中提及这些部分，所以现在了解global memory就足够了。 而对于local memory，我们暂时只需要知道是一组Static RAM就可以了。 稍后我会进一步解释。

通常global memory很大，用于存储来自host端的整个数据块。

而local memory虽然有限但在计算速度上更有优势。

所以有时候对于一个很大的张量，我们需要把它切分成几个部分，送到local memory中进行计算，然后把结果存回global memory。

为了在 TPU 上执行这些操作，我们就需要用到指令。

指令主要有两种：

1. GDMA用于system memory和local memory间或system memory内的数据传输；
2. BDC用于驱动执行单元在NPU上做计算工作；
   另外，对于那些不适合并行加速的计算，比如NMS，SORT，我们还需要HAU指令，但是这意味着我们需要额外的处理器。

对于local memory的构成，它是由多个Static RAM组成的。每个 SRAM 称为一个bank。此外，我们将这些 SRAM 分成多个部分给同样数量的NPU，每个部分称为一个lane。

而对于每个NPU，它只能访问属于它的那部分local memory，这使得单个NPU的执行单元只能处理自己local memory上的那部分张量。

一旦我们调用单个 BDC 指令，所有 NPU 的执行单元将在每个 NPU 的相同位置执行相同的操作。 这就是 TPU 加速运算的方式。

此外，TPU 可以同时处理的数据数量取决于每个 NPU 上的执行单元数量。

对于一个特定的TPU，EU Bytes是固定的，所以对于不同类型的数据，EU的个数会有所不同。

例如当EU Bytes为64时，则表示一个NPU可以同时处理64个int8数据。

同理，我们可以根据数据的字节计算出对应的EU_NUM。

对于地址分配，假设我们的local memory由16个SRAM组成，总内存为16MB，有64个NPU，那么每个NPU的内存为256KB。

单个通道中每个bank的内存大小则为16KB，相当于16x1024 字节。

所以这个块的地址范围是从0到16x1024 – 1。

同理，NPU0中下个bank的地址从16x1024开始到32x1024-1

按照这个规则，我们就可以得到local memory上的所有地址。
 

　　审核编辑：汤梓红