深入浅出PyTorch（算子篇）

来源：马士兵AI程序员　2023-04-26 17:31:44

Tensor

自从张量（Tensor）计算这个概念出现后，神经网络的算法就可以看作是一系列的张量计算。所谓的张量，它原本是个数学概念，表示各种向量或者数值之间的关系。PyTorch的张量（torch.Tensor）表示的是N维矩阵与一维数组的关系。

【资料图】

torch.Tensor的使用方法和numpy很相似（https://pytorch.org/…tensor-tutorial-py），两者唯一的区别在于torch.Tensor可以使用GPU来计算，这就比用CPU的numpy要快很多。

张量计算的种类有很多，比如加法、乘法、矩阵相乘、矩阵转置等，这些计算被称为算子（Operator），它们是PyTorch的核心组件。

算子的backend一般是C/C++的拓展程序，PyTorch的backend是称为"ATen"的C/C++库，ATen是"A Tensor"的缩写。

Operator

PyTorch所有的Operator都定义在Declarations.cwrap和native_functions.yaml这两个文件中，前者定义了从Torch那继承来的legacy operator（aten/src/TH），后者定义的是native operator，是PyTorch的operator。

相比于用C++开发的native code，legacy code是在PyTorch编译时由gen.py根据Declarations.cwrap的内容动态生成的。因此，如果你想要trace这些code，需要先编译PyTorch。

legacy code的开发要比native code复杂得多。如果可以的话，建议你尽量避开它们。

MatMul

本文会以矩阵相乘–torch.matmul()为例来分析PyTorch算子的工作流程。

我在深入浅出全连接层（fully connected layer）中有讲在GPU层面是如何进行矩阵相乘的。Nvidia、AMD等公司提供了优化好的线性代数计算库–cuBLAS/rocBLAS/openBLAS，PyTorch只需要调用它们的API即可。

Figure 1是torch.matmul()在ATen中的function flow。可以看到，这个flow可不短，这主要是因为不同类型的tensor（2d or Nd, batched gemm or not，with or without bias，cuda or cpu）的操作也不尽相同。

at::matmul()主要负责将Tensor转换成cuBLAS需要的格式。前面说过，Tensor可以是N维矩阵，如果tensor A是3d矩阵，tensor B是2d矩阵，就需要先将3d转成2d；如果它们都是>=3d的矩阵，就要考虑batched matmul的情况；如果bias=True，后续就应该交给at::addmm()来处理；总之，matmul要考虑的事情比想象中要多。

除此之外，不同的dtype、device和layout需要调用不同的操作函数，这部分工作交由c10::dispatcher来完成。

Dispatcher

dispatcher主要用于动态调用dtype、device以及layout等方法函数。用过numpy的都知道，np.array()的数据类型有：float32, float16，int8，int32，… 如果你了解C++就会知道，这类程序最适合用模板（template）来实现。

很遗憾，由于ATen有一部分operator是用C语言写的（从Torch继承过来），不支持模板功能，因此，就需要dispatcher这样的动态调度器。

类似地，PyTorch的tensor不仅可以运行在GPU上，还可以跑在CPU、mkldnn和xla等设备，Figure 1中的dispatcher4就根据tensor的device调用了mm的GPU实现。

layout是指tensor中元素的排布。一般来说，矩阵的排布都是紧凑型的，也就是strided layout。而那些有着大量0的稀疏矩阵，相应地就是sparse layout。

Figure 2是strided layout的演示实例，这里创建了一个2行2列的矩阵a，它的数据实际存放在一维数组（a.storage）里，2行2列只是这个数组的视图。

stride充当了从数组到视图的桥梁，比如，要打印第2行第2列的元素时，可以通过公式：1 ∗ s t r i d e ( 0 ) + 1 ∗ s t r i d e ( 1 ) 1 * stride(0) + 1 * stride(1)1∗stride(0)+1∗stride(1)来计算该元素在数组中的索引。

除了dtype、device、layout之外，dispatcher还可以用来调用legacy operator。比如说addmm这个operator，它的GPU实现就是通过dispatcher来跳转到legacy::cuda::_th_addmm。