PyTorch 源码学习：Dispatch Autograd Operators

对于 PyTorch 的动态计算图 (Dynamic Computation Graph) 模式来说，算子 (Operators) 注册、分发 (Dispatch) 机制和自动微分 (Autograd) 机制是至关重要的，了解这部分内容有助于更好地理解 PyTorch 动态计算图的运行机制。本文分享自己在学习 PyTorch 源码时阅读过的资料，重点关注与 PyTorch 动态计算图有关的算子 (Operators) 注册、分发 (Dispatch) 机制和自动微分 (Autograd) 机制。因为 PyTorch 不同版本的源码实现有所不同，所以笔者在整理资料时尽可能按版本号升序，版本号见标题前[]。最新版本的源码实现还请查看 PyTorch 仓库。更多内容请参考：

• Ubuntu 22.04 LTS 源码编译安装 PyTorch
• pytorch/CONTRIBUTING.md 机翻
• 深度学习框架与静态/动态计算图【笔记】
• PyTorch 源码学习：阅读经验 & 代码结构
• PyTorch 源码学习：从 Tensor 到 Storage-CSDN博客

PyTorch 动态计算图

以下内容来自：[1.0.1] 想读读PyTorch底层代码？这份内核机制简介送给你

PyTorch 使用一种称之为 imperative / eager 的范式，即每一行代码都要求构建一个图以定义完整计算图的一个部分。即使完整的计算图还没有完成构建，我们也可以独立地执行这些作为组件的小计算图，这种动态计算图被称为「define-by-run」方法。

以下内容来自：[1.13.0] PyTorch源码解析（1）- 整体预览 - Hurray’s InfoShare

一个主流的训练框架需要有两大特征：

1. 实现类似 numpy 的张量计算，可以使用 GPU 进行加速；
2. 实现带自动微分系统的深度神经网络。

而 PyTorch 能在一系列训练框架中脱颖而出成为今天的主流，主要原因是其原生支持动态图，具备对用户友好的特点。

在 TensorFlow1.x 中，我们如果需要执行计算，需要建立一个session，并执行session.run()来执行。其整个过程其实是将计算过程构成了一张计算图，然后运行这个图的根节点。这样先构成图，再运行图的方式我们称为静态图或者图模式。

而在 PyTorch 中，我们可以在计算的任意步骤直接输出结果。原因是，PyTorch 每一条语句是同步执行的，即每一条语句都是一个（或多个）算子，被调用时实时执行。这种实时执行算子的方式我们称为动态图或算子模式。

• 动态图的优点显而易见，可以兼容 Python 式的编程风格，实时打印编程结果，用户友好性做到最佳；
• 静态图则在性能方面有一定优势，即在整个图执行前，可以将整张图进行编译优化，通过融合等策略改变图结构，从而实现较好的性能。

PyTorch 除了有原生的动态图之外，当前也在静态图方面有诸多路线尝试：分别有torchscript(jit.script、jit.trace)、TorchDynamo、torch.fx、LazyTensor。这里用一幅图总结了各路线所处的层次：

在本文，我们重点关注上图左半边“动态图 - 标准执行流程”的内容。

首先 PyTorch 的动态图是从 Python 源码下降，拆分成多个 python 算子调用，具体调用到 tensor 的 OP。经过 pybind 转换到 c++，并通过 dispatch 机制选择不同设备下的算子实现，最终实现调用的底层设备实现（如 Nvidia cudnn、intel mkl 等）。

以下内容来自：[1.13.0]PyTorch源码分析（2）——动态图原理 - Hurray’s InfoShare

我们前面说过，训练框架最重要的特点是：

1. 支持类似 numpy 的张量计算，可以使用GPU加速；
2. 支持带自动微分系统的深度神经网络；

此外 PyTorch 的特征还包括原生支持动态图执行。

那么依据这三点，我们针对 PyTorch 提出三个问题：

1. PyTorch 如何支持 CPU、GPU 等诸多设备的？
2. PyTorch 如何实现自动微分的？
3. 动态图的原理是什么？

本文实际也是围绕这三个问题展开的，但与引用的资料所关注的重点有所不同。

这篇文章有一句话值得揣摩：“多设备的支持包括 tensor的存储 以及 tensor之间的操作（算子）两个方面。”

最后，用一张图总结一下本文的主要内容。本文主要介绍了 PyTorch 动态图的原理，主要分为三个部分。

• 对于多设备支持，PyTorch 通过 dispatch 机制实现了不同设备下的算子映射。PyTorch 提供了一套算子注册方法，并在源码中利用codegen模块依据算子注册表native_functions.yaml自动生成算子的注册。
• 对于反向传播，PyTorch 的自动微分模块也是依赖codegen模块和微分注册表 (derivatives.yaml) 自动生成反向算子的实现、注册。微分注册表中主要是各类算子的反向公式，并且只包含最基础的算子的反向公式。
• 对于动态图原理，前向是在算子调用的时候直接执行的算子实现；而反向则是通过一套复杂的Execution Engine，利用线程、队列等，依次调用执行反向图中的grad_fn（也即反向算子）实现的执行过程。

以下内容来自：[1.6.0]一文详解pytorch的“动态图”与“自动微分”技术

众所周知，Pytorch 是一个非常流行且深受好评的深度学习训练框架。这与它的两大特性“动态图”、“自动微分”有非常大的关系。

• “动态图”使得 pytorch 的调试非常简单，每一个步骤，每一个流程都可以被我们精确的控制、调试、输出。甚至是在每个迭代都能够重构整个网络。这在其他基于静态图的训练框架中是非常不方便处理的。在静态图的训练框架中，必须先构建好整个网络，然后开始训练。如果想在训练过程中输出中间节点的数据或者是想要改变一点网络的结构，就需要非常复杂的操作，甚至是不可实现的。
• 而“自动微分”技术使得在编写深度学习网络的时候，只需要实现算子的前向传播即可，无需像 caffe 那样对同一个算子需要同时实现前向传播和反向传播。由于反向传播一般比前向传播要复杂，并且手动推导反向传播的时候很容易出错，所以“自动微分”能够极大的节约劳动力，提升效率。

用过 caffe 或者 tensorflow 的同学应该知道，在训练之前需要构建一个神经网络，caffe 里面使用配置文件 prototxt 来进行描述，tensorflow 中使用 python 代码来描述。训练之前，框架都会有一个解析和构建神经网络的过程。构建完了之后再进行数据读取和训练。在训练过程中网络一般是不会变的，所以叫做“静态图”。想要获取中间变量的输出，可以是可以，就是比较麻烦一些，caffe 使用 c++ 训练的话，需要获取 layer 的 top，然后打印，tensorflow 需要通过 session 来获取。但是如果想要控制网络的运行，比如让网络停在某一个 OP 之后，这是很难做到的。即无法精确的控制网络运行的每一步，只能等网络运行完了，然后通过相关的接口去获取相关的数据。

而 pytorch 的“动态图”机制就可以对网络实现非常精确的控制。在 pytorch 运行之前，不会去创建所谓的神经网络，这完全由 python 代码定义的 forward 函数来描述。即我们手工编写的 forward 函数就是 pytorch 前向运行的动态图。当代码执行到哪一句的时候，网络就运行到哪一步。所以当你对 forward 函数进行调试，断点，修改的时候，神经网络也就被相应的调试、中断和修改了。也就是说 pytorch 的 forwad 代码就是神经网络的执行流，或者说就是 pytorch 的“动态图”。对 forward 的控制就是对神经网络的控制。如下图所示：

正因为这样的实现机制，使得对神经网络的调试可以像普通python代码那样进行调试，非常的方便和友好。并且可以在任何时候，修改网络的结构，这就是动态图的好处。

以下内容来自：[unknown]理解 PyTorch 的即时执行模式

即时执行模式为 PyTorch 带来了高灵活性，可以在 python 环境下直接执行计算操作、控制计算过程，甚至构建动态计算图。

执行过程中，通过动态分派算法，根据输入张量的类型、设备、形状等动态特性，决定如何调用适当的函数来执行计算操作。这种灵活性对于深度学习框架至关重要，因为深度学习模型需要在各种条件下进行训练和推理。

更多关于 PyTorch 动态计算图的解读资料：

• [2.0.0]新手如何入门 PyTorch 动态图

PyTorch Dispatch

以下内容来自：[1.13.0]PyTorch源码分析（2）——动态图原理 - Hurray’s InfoShare

上图来自：Let’s talk about the PyTorch dispatcher : ezyang’s blog

我们可以将 Dispatch 机制看做一个二维的表结构。其一个维度是各类设备（CPU、CUDA、XLA、ROCM等等），一个维度是各类算子（add、mul、sub等等）。

PyTorch提供了一套定义（def）、实现（impl）机制，可以实现某算子在某设备（dispatch key）的绑定。例如上图m.impl()中就是对dispatch key为CPU时neg算子的实现绑定，其绑定了neg_cpu()这个函数。

以下内容来自：[unknown]理解PyTorch分发机制的内部工作原理

PyTorch 的动态性源自内部的调度器（dispatcher），它可以根据不同的输入类型自动选择正确的运算方式。当调用 Python 函数时，调度器会根据传入的参数类型选择正确的操作实现，这个过程称为分派（dispatch）。

例如，当执行矩阵乘法torch.matmul(a, b)时，调度器会根据输入张量a和b的类型（dtype、shape、device等）选择正确的 BLAS 库（CPU 还是 CUDA，float 还是 half，是否批量计算）来进行计算。对于 PyTorch 来说，模型的执行过程就是将各个操作（op）分派给本地方法（native function）执行的过程。

上图来自：Let’s talk about the PyTorch dispatcher : ezyang’s blog

dispatcher 为每个 op 都维护了一张跳转表（它有点像 C++ 实现多态用的虚表），如上图所示，表中每个条目存储了一个本地方法，

• 有些方法和输入张量所属的设备有关，比如 XLA/CUDA/CPU，
• 有的和 requires_grad 有关，比如 Autograd。

当 op 被执行时，e.g. aten::addmm，调度器会在它的跳转表中找出一个方法来执行，而且一个 op 执行过程可能会调用多个方法，例如，输入张量需要求导（requires_grad = true），那会

• 先调用 Autograd 方法来构建反向图，
• 再调用 backend（CPU/CUDA/XLA）的方法来运算。

文中还详细讨论了分派规则和分派流程。更多内容请参考原文。另外，作者还在下面两篇博客中，讨论了算子注册和执行。

• PyTorch Internal：算子注册
• PyTorch Internal：算子执行

以下内容来自：[unknown]PyTorch源码学习系列 - 3. 算子

当我们调用函数时，我们就会接触到 PyTorch 的调度模块。当我们调用torch.add函数时，我们总共会经历两次调度：

• 第一次调度会根据 tensor 的**设备（device）和布局（layout）**动态选择对应的实现函数，比如<CPU, Strided> Tensor，<CPU, Sparse> Tensor或者<GPU, Strided> Tensor。不同设备布局的实现可能会编译在不同的动态链接库里。
• 第二次调度则会根据 Tensor 元素的数据类型通过 switch 分支的方式进行一次轻量级的静态选择，最终选出正确的 Kernel 来执行对 Tensor 的操作。

PyTorch 的调度模块非常复杂，幸运的是在我们从调用函数开始直到深入 Kernel 之前所有的代码都是由 PyTorch 自动生成的，对于初学者而言只需要了解大概的过程就可以了。

更多关于 dispatch 机制的解读资料：

• [1.5.0]PyTorch & MMCV Dispatcher 机制解析
• [unknown]【Pytorch 源码 Detail 系列】Pytorch 中 dispatch 机制及其实现
  • [unknown]【Pytorch 源码 Detail 系列】DispatchKey & DispatchKeySet

PyTorch Operators & Autograd

以下内容来自：[unknown]PyTorch源码学习系列 - 3. 算子

算子是什么？

• 狭义的算子（Kernel）。对 Tensor 执行的基本操作集合，包括四则运算，数学函数，甚至是对 Tensor 元数据的修改，如维度压缩（squeeze），维度修改（reshape）等
• 广义的算子（Function）。PyTorch 中算子模块的具体实现，涉及到调度模块，Kernel模块，求导模块以及代码自动生成模块。

在后续的内容中，

• 将狭义的算子称之为核（Kernel），在 PyTorch 架构图里，C++ 实现层里的 Operator 指的就是这里的 Kernel，这里的 Kernel 并不支持自动梯度计算（Autograd）模块。
• 广义的算子我们将其称之为函数或方法，这也是我们平时经常接触到的 PyTorch API，包括 Python API 和 C++ API，其配合 PyTorch Autograd 模块后就可以支持自动梯度求导计算。

Kernel 主要是算子的计算模块，算子还包含求导模块。

• 计算模块主要定义了 Kernel 的计算步骤，我们需要先在aten/src/ATen/native/native_functions.yaml中声明 Kernel 计算模块的函数签名，然后在native/目录下实现该函数。在前面的函数调用中，我们主要就通过 Kernel 对 Tensor 进行操作。
• 求导模块主要是对计算模块的一个反向求导，我们需要直接在tools/autograd/derivatives.yaml中声明定义求导的过程，剩下的就可以交给代码生成模块帮我们自动生成对应的代码。

PyTorch 的函数是一个非常复杂核心的模块，其大部分代码都是由 PyTorch tool 根据模板文件自动生成。如果我们想要查看其源代码，我们无法直接在 GitHub 代码库中搜索到，必须要将代码下载到本地进行编译。

文中还详细讨论了算子声明、算子实现和算子求导。更多内容请参考原文。

以下内容来自：[1.13.0]PyTorch源码分析（2）——动态图原理 - Hurray’s InfoShare

将 PyTorch 算子总结为下面的图：

在我们调用合成算子的前向时，实际是被拆成多个基础算子被执行的，因此生成的反向图也是基础算子的反向。

PyTorch 算子流程：

整个流程总结如上图。即 PyTorch 中只需要实现native_functions.yaml[参考]这个算子配置文件，以及算子具体实现的函数。其他的都可以利用codegen模块自动完成整个流程的注册实现。最终将实现 Tensor 的算子成员函数，供用户使用。

• 前向算子的配置文件(native_functions.yaml) → aten/src/ATen/native/native_functions.yaml
• 反向算子的配置文件(derivatives.yaml) → tools/autograd/derivatives.yaml

我们可以总结：

1. 前向计算算子执行时，每个tensor会绑定一个grad_fn，也就是反向算子；
2. 反向算子是由反向公式表(derivatives.yaml)自动生成的。

而算子直接之所以可以通过图的顺序依次通过反向算子计算，其主要是由链式法则决定的。

以下内容来自：[1.6.0]这可能是关于Pytorch底层算子扩展最详细的总结了！

pytorch 直接扩展底层 C++ 算子主要有三种方式，

• native_functions.yaml：pytorch 的原生算子很多都是使用这种方式组织的。
• C++ extension：pytorch 提供的另外一种更加简便的方式。具体内容参考原文。
• OP register：pytorch 提供的一种更加强大的算子扩展能力。具体内容参考原文。

在native_functions.yaml中有关于各个算子的说明，然后在同级目录下面有这些算子的实现。

使用该方式添加新的算子，主要用在已经支持的硬件上面。例如 pytorch 本身已经支持了 CPU 和 GPU，此时需要一些新的算子，该算子只需要在 CPU 或者 GPU 上面运行，那么这种方式就非常适合。

只需要定义新算子的 kernel 实现，然后添加配置信息，就可以自动生成：torch.xxx()、torch.nn.functional.xxx()以及tensor.xxx()方法，而不用去关注算子与 pytorch 是如何衔接，以及如何把算子添加到 tensor 的属性中等其他细节。

native_functions.yaml文件对于每个算子的描述，包括几个主要字段：

• func字段：表示算子的名称以及输入输出参数类型。
• variants字段：表示需要自动生成的高级方法。
  • function表示自动生成torch.func()方法，
  • method表示生成 tensor 的func()方法，即可以定义一个tensor a，然后可以执行a.func()方法。
• dispatch字段：表示分发的设备类型对应的op方法。
  • CPU 指的是该算子支持 CPU 设备，
  • CUDA 指的是当前算子支持 GPU 设备。

native_functions.yaml 方式的优缺点

• 优点：可以比较方便的增加或者修改算子。
• 缺点：与 pytorch 的耦合度过高。由于在 pytorch 的源码中直接修改，那么每次增加或者修改算子都需要重新编译 pytorch。

具体分析添加算子的流程：leakly_relu

1. 在native_functions.yaml中添加算子的说明，包括反向传播函数。
2. 在配置文件derivatives.yaml中添加算子和反向算子的对应关系。
3. 在aten/src/Aten/native/目录下面通过 C++ 实现相关的算子流程。这里定义的实现只是一个封装，没有真正的实现。
4. 在后端实现，CPU 端的实现流程都在aten/src/Aten/native/cpu/Activation.cpp中，GPU 端的实现流程都在aten/src/Aten/native/cuda/Activation.cu中。

以下内容来自：[1.6.0]一文详解pytorch的“动态图”与“自动微分”技术

Pytorch 能够实现精确的反向传播，其最大的奥秘就藏在 tensor 的grad_fn属性里面。

Pytorch 中的 tensor 对象都有一个叫做grad_fn的属性，它实际上是一个链表，实现在 pytorch 源码的autograd下面。该属性记录了该 tensor 是如何由前一个 tensor 产生的。

文中详细介绍了 pytroch 中的 leaf tensor 和非 leaf tensor，以及 grad_fn 属性。更多内容请参考原文。

更多关于 Operators & Autograd 的解读资料：

• [1.8.0] Overview of PyTorch Autograd Engine | PyTorch → 翻译
  • [1.8.0]How Computational Graphs are Constructed in PyTorch | PyTorch → 翻译
  • [1.11.0]How Computational Graphs are Executed in PyTorch | PyTorch → 翻译
• [unknown] PyTorch源码学习系列 - 4. Autograd
• [unknown] PyTorch源码浅析(4)：Autograd | NIUHE（内容稍微旧了些）