深入浅出PyTorch（算子篇）

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来完成。