Google的TPU的向量化内存的读取规格——单次读取/写入的数据量

异构加速设备：

GPU、TPU、NPU

这几种设备，除了GPU公开了部分硬件设计原理和软件编程范式以外，所有的TPU和NPU的资料都是不公开的，都是被其所属公司保留的，然后这些公司会自己在这些硬件之上和这些软件编程接口之上作出高级编程框架提供给用户，而最后外界可以看到的只有最上层的编程框架，如Google的TPU和Tensorflow，在这中间的细节和硬件的具体特性都是不公开的。但是，我们可以根据一些零星的公开资料知道一些硬件的基本特性，本文就是介绍TPU的内存的读写特性。

在Google的TPU中，内存非两种类型，一种是向量计算所用的内存，另一种是标量计算所用的内存。

本文介绍TPU的向量内存的部分读写特性：

资料地址：

https://jax.readthedocs.io/en/latest/pallas/tpu.html

TPU的向量内存，对32位的数据类型的切片读取、写入操作没有限制，也就是float32、int类型的array都是可以随意操作的，比32位短的数据类型，如float16、int16等可能存在限制。

TPU向量内存单次读取的内存大小为（8， 128）规格，因此使用TPU的向量内存读写时最后使array的最后两个axis的size是（8， 128）的倍数，因为这样读取效率最高。