《智能计算系统》第五章 编程框架原理（上）课程视频链接：https://www.bilibili.com/video/BV1Ei421i7Rg

本文源自于B站国科大计算所 智能计算系统课程官方账号 所公开上传的视频，在原有视频之上，提取了关键帧、将音频转成了文字并进行了校正，以便学习使用。在此，也感谢国科大计算所 智能计算系统课程官方能够将该课程公开！

若侵权，请联系删除。

智能计算系统 第五章 编程框架原理（上）

大家好，下面我们来介绍智能计算系统第五章编程框架原理

这个是第五章的提纲，首先我们会介绍编程框架的设计原则，然后我们会围绕编程框架的主要的四大模块儿，分别去介绍编程框架中的计算图构建、计算图执行、深度学习编译以及分布式训练的主要的原理

然后最后是本章小结

概述

首先介绍为什么要了解编程框架原理，了解编程框架的原理对学习智能计算系统来说是有很大的意义的，它有助于程序员编写与框架底层更为契合、性能更优的代码，从而提升算法的实现效率

此外，在面对新的算法和硬件挑战的时候，程序员还能够具备定制化扩展培训框架以支持以提供支持的这个能力

那么在一个编程框架中，它主要包含了四个模块，分别是计算图构建模块、计算图执行模块、深度学习编译模块和分布式训练模块

其中计算图构建模块和计算图执行模块是必要的，而深度学习编译模块和分布式训练模块是为了追求更高的性能而需要的

那在本章中，我们首先介绍编程框架的设计原则，再介绍编程框架内部的整体架构；设计原则会抽象的指导编程框架的设计，而整体架构则是对这一指导的一个具体的细化实现

设计原则

不同的编程框架，它拥有各自的设计哲学

大体上总结可以分为三种，简洁性、易用性、高效性

简洁性是指由框架提供一套抽象的机制，那用户只需要关心宏观上的操作，而不需要去关心微观上的具体实现

像我们这个例子里面，我们要去实现一个把张量 b 和 a 去做加法，得到 c 这样的一个结果，其中张量 b 原本在 CPU 上，我们需要把张量 b 先搬运到 GPU 上，然后再去做张量 a 和 b 的这个加法的操作

那么在使用编程框架进行对这个功能进行编程的时候，我们只需要显示的去指定一个张量存放的位置和张量移动的这个时机；但是具体的这个张量应该怎么样存储，怎么样在设备间移动，是不需要用户去操心的，由可以交由编程框架去维护

我们再来看易用性；Pytorch，它自诞生就坚持 Python 优先，它提供了符合 Python 语言和库管理的一个交互方式，而不是一套需要重新学习的内嵌语言

那 Pytorch 的用户接口设计思路也一直忠于 Python 的语言惯例和开发习惯

高效性是指编程框架应进行充分的优化，从而尽量提高用户应用程序的运行效率

比如说像 Tensorflow 这个编程框架，他就提供了这个声明式的静态图编程方式，这样使得编程框架可以获得一个完整的计算图并进行全局的优化

那现有的编程框架大部都大部分都支持深度学习的编译技术，通过这个多层级的表示优化来充分的利用用户硬件的一个计算的能力

同时也基本上都支持多机多卡条件的一个分布式训练，从而能够高效支持大规模深度学习的任务

整体架构

编程框架的整体结构主要分成四大模块，分别是计算图构建模块、分布式训练模块、深度学习编译模块以及计算图执行模块

计算图构建模块，它的功能是完成从输入的用户程序到编程框架内部原始计算图的转换过程，是编程框架的入口模块

分布式训练模块，它主要是应对更大规模的神经网络，将训练推理任务从一台设备扩展到多台设备

深度学习编译模块，是对计算图分别进行图层级和图算子层级的编译优化，从而提升单设备上的执行效率

最后是这个计算图执行模块，它是将优化后的计算图中的张量和操作映射到指定的设备上去进行具体执行，并给出编程框架的输出结果

那么我们来看一下这张图，用户通过编程框架去编写用户程序，编写完了以后的程序，通过编程框架的计算图构建模块，构建出这个用户程序对应的这个原始的计算图

那有了这个原始的计算图以后，通过编程框架中提供的这个分布式训练机制对原始计算图在多个计算设备上去进行拆分，得到不同计算设备上拆分后的计算子图

那对于这个拆分以后的计算子图，我们在应用深度学习编译这个模块，这个深度学习编译模块儿它会针对计算图分别进行图层级和算子层级的编译优化

得到优化后的计算图，针对这个优化后的计算图，我们使用计算图执行模块，将优化以后的计算图中的张量和操作映射到指定设备上进行具体执行，并且给出编程框架的输出结果

接下来我们就针对编程框架中的 4 个模块的原理去进行简单的介绍

计算图构建

首先介绍计算图构建模块

计算图它主要是由两个基本的元素构成，分别是张量以及张量操作，张量就是我们其中称谓的tensor，张量操作就是operation

计算图它是一个有向图，包含了有向的边，这些有向边就指明了张量的流动方向；

下面这张图它展示了一个神经网络，从源代码到构建正向和反向计算图的过程

正向传播会构建正向计算图，反向传播是利用自动求导的原理来去构建反向计算图

那么我们这一小节会首先从正向传播的计算图构件讲起，在介绍反向传播时计算图构件的具体原理

正向传播

在正向传播中， 输入张量经过搭建的神经网络层层计算传递，最终获得计算结果

那么计算图它的构建形式主要分成两种，一种是动态图，一种是静态图

动态图它的意思是说在执行函数的时候，我会按照函数顺序逐条语句的去生成节点，然后立即计算返回结果。它的优点就是易调试，但是它的性能优化空间有限

静态图是和动态图相对的，它是在执行计算之前已经构建好了所有图上的节点，然后在图运行的时候才计算整个计算图，并且返回最终的结果，那么它的优点和取它的优点就是性能好；那么它的缺点就是不易调试

动态图

在PyTorch 1.X的版本里主要使用的就是这个动态图的机制，动态图就是说计算图在函数运行过程中逐渐构建，那它执行的是一种称为立即模式(Eager)的这样的一个计算过程，就是每次调用语句就会立即执行计算，在 Pytorch 中使用的就是动态图的机制，就是每次执行一条语句就会重新构建这个计算图

那我们来看一下左下这张图的这个代码，前四行代码，它分别构建了四个张量，那对应的我们就会构建在计算图中构建相应的四个节点，Wh、 Wx 、x以及这个prev_h 这样的四个张量节点

接下来的两行语句分别是进行了两个矩阵乘法的计算，得到了这个 h2h 和 i2h 这样的两个张量，那么我们对应的也在计算图中添加了这个两个矩阵乘的节点，以及 h2h和 i2h 这样的两个张量的节点

再接下来的两行代码，我们去计算 next_h

然后最后一行代码是计算loss相应的，每执行完一条，语句都会在计算图中构建出对应的计算节点、张量节点以及这个对应的边

当我们的语句编写完成以后，我们的计算图也就构建完成

那么无论这一系列语句它是否在同一个函数里面，我们在下一次构建的时候，在下一次调用的时候都会重新构建，即使这个图的结构可能和这个图是完全相同的

静态图

那我们再来看静态图，静态图指的是整个网络的结构会在开始计算前就建立完成计算图，然后在框架执行时它会接收整个计算图，而不是单一的语句

在 Tensorflow 1.X中使用的就是静态图的机制，它使用若干基本控制流的算子，其实它一共使用了 5 个基本的控制流算子，分别是 Switch、Merge、Enter、Exit 和 NextIteration 这五个控制流算子，它通过这五个控制流的算子的这个不同的组合能够去实现一些，比如说像跳转、像循环等等这样一些比较复杂的控制流的场景

在 Pytorch 2.X的版本里面，同样也支持静态图的机制。在 Pytorch 2.0 中采取了图捕获TorchDynamo 的这个技术，将用户的动态图转化为静态图

反向传播

那我们再来看反向传播，在正向传播的时候，正向计算得到的这个结果和目标结果之间通常会存在一个损失函数的值

然后我们就会使用损失函数对我们的参数去进行求导，然后得到梯度，并且使用这个梯度去更新我们的参数

在这个反向传播的过程中， PyTorch 就会根据正向传播的计算图自动生成对应的反向计算节点，并且把它加入到计算图中去，构建成反向传播计算图

反向传播过程中一个必不可少的步骤就是要计算导数，那么自动微分就是一种计算导数的方法，目前常用的几种求导的方法包括了手动求导法、数值求导法、符号求导法和自动求导法

• 手动求导法就是指使用链式法则求解出梯度公式，然后根据公式去编写代码，代入数值计算得到梯度结果
• 数值求导法的含义是直接代入数值来近似求解
• 符号求导法是直接对代数表达式进行求解，最后代入问题数字，它会有表达式膨胀的问题
• 自动求导法就是用户只需要描述前向计算的过程，然后由编程框架自动的推导反向计算图，它会先建立表达式，再带入数值计算，那自动求导法是当前的这个编程框架都会支持的求导的方法

手动求导法

好，我们先来看手动求导法，它就是指手动用链式法则求解除梯度公式，然后代入数值得到最终的梯度值，这是一种非常直观的这个求导的方法，那么它的缺点就是对于大规模的深度学习算法，手动用链式法则进行梯度计算，并且转换成计算程序是非常困难的

我们需要手动去编写梯度的求解代码，当模型发生变化的时候，算法也需要随之修改

数值求导法

那数值求导法就是指利用导数的原始定义去进行求解，就是我们下面给出来的一个公式，我们使用一个非常小的一个数h，然后把它带入到这个公式中，就能求得某一点处的这个函数的这个导数，那么数值求导法它的优点就是易操作，并且可以对用户隐藏求解的过程，那它的缺点就是计算量大，并且它的求解速度比较慢，同时它可能会引起舍入误差和截断误差

符号求导法

符号求导法是利用求导规则来对表达式进行自动操作，从而获得导数

我们会把常见的一些求导规则，比如说加法求导规则、减法求导规则、乘法除法这些求导的规则带入到表达式上，获得这个求导以后的表达式，然后最后再带入我们的这个数值，那么它的问题就是很容易会带来表达式膨胀的这个问题，最终导致求解的速度变慢

我们以下面这张图为例来去介绍表达式膨胀的含义啊

那么我们有一个\(l_n\)，这个\(l_n\)是和\(x\)有关的，随着\(n\)的变化，它\(l_n\)的表达式会变得更复杂

那么我们对\(l_n\)去计算\(l_n\)相对于\(x\)的导数，这个第3列表示的是使用符号求导法求得的这个表达式的结果，这个第4列表示的是这个使用手动求导法去计算这个导数的这个结果

那我们会发现当\(n\)等于 4 的时候，使用符号求导法计算出来的这个表达式相比于手动求导法求得的这个结果来说已经是非常的复杂了。所以随着这个\(n\)逐渐增大，这个求导的结果就会出现表达式膨胀的问题

自动求导法

目前编程框架中使用的基本上都是自动求导法，那么根据我们前面的介绍，数值求导法，它强调一开始直接带入数值去做近似的求解；符号求导法，它强调直接对代数表达式去进行求解，然后才代入问题数值，那么它们都有各自的一些问题。

自动求导法，它是介于数值求导和符号求导方法之间的一种方法，那它首先将符号求导法应用于最基本的算子，比如说幂函数、指数函数、对数函数、三角函数等等，然后再代入数值，保留中间的计算结果，最后再应用于整个函数

整个的计算，它分成两步执行：首先第一步是根据原始函数建立计算图，数据正向传播，然后我们会计算出所有的中间节点，并且记录计算图中的节点依赖关系

第二步我们会反向遍历计算图，计算输出对每一个节点的导数

那对于前项计算中一个数据连接多个输出数据的情况，在自动求导法里会把这些输出数据相对于该数据的导数去进行一个累加

我们这里根据一个例子来去介绍这个自动求导法，它的这个计算过程，我们有一个想要去计算\(f(x_1, x_2)=(e^{x_1}+x_2)(x_2+1)\)这样的一个函数，那么首先第一步我们会根据这个函数的内容去创建计算图

创建完计算图以后，我们会让输入数据\(x_1\)和\(x_2\)去进行一个正向的传播，计算出所有的中间节点，\(x_3\), \(x_5\), \(x_4\), \(x_6\) 以及\(y\)（下图右半部分），并且记录计算图中的节点依赖关系

那么我们在做前向计算的时候，会把每一步的中间计算、中间结果都存下来，因为这个在反向计算的时候是要用到的

那第二步我们就要去进行反向计算，要计算输出对每一个节点的导数，那我们就看这个左边下边的这张图，我们从输出开始，向输入的方向逐个节点去计算它们的这个导数。

备注：反向传播的图是怎么来的呢？在下面会有PyTorch的代码，个人理解是：对于正向传播中的每一条边，都会有相应的反向传播的节点；而不是正向传播中的每一个节点对应反向传播的节点，这个图虽然画的很清楚，但是感觉还是会产生一定误导的=_=

右边这张图我们从下往上看，首先计算第一个节点，也就是 \(y\)相对于\(x_6\) 这个节点的导数，因为 \(y\) 它等于\(x_6\)，所以这个 \(y\) 相对\(x_6\)这个节点的导数就为1

然后我们再看\(x_5\)，再看这个 \(y\) 对 \(x_5\) 这个节点的导数

根据链式法则， \(y\) 对 \(x_5\) 这个节点的导数，它等于 \(y\) 对 \(x_6\) 这个节点的导数，乘以 \(x_6\) 对\(x_5\)这个节点的导数（\(\frac{\partial y}{\partial x_5}=\overline{x_5}=\frac{\partial y}{\partial x_6}\cdot\frac{\partial x_6}{\partial x_5}\)），那么 \(y\) 对 \(x_6\)这个节点导数在上一步中已经计算求得了（\(\frac{\partial y}{\partial x_6}=\overline{x_6}\)），那我们只要计算出 \(x_6\)对 \(x_5\)这个节点的导数，它等于\(x_4\)（\(\frac{\partial x_6}{\partial x_5}=x_4\)）；那我们就把之前向传播时保存的\(x_4\)节点的计算值带入到这个公式中，就求得了这个 \(y\) 对 \(x_5\) 这个节点的导数

小结：我们通过链式法则逐步的就计算出输出对每一个节点的导数。

那这里需要特殊关注的，就是\(x_2\)这个节点；它的特殊之处就是：它有两个分支，分别去连向\(x_5\)这个节点以及\(x_4\)这个节点。那么我们在计算导数的时候，我们就会分开分别计算：先去计算\(x_5\)这个分支，从这个分支上 \(y\)结点对于\(x_2\)的导数，得到这个\(\overline{x_2}^1\)，然后再去计算\(x_4\)这条分支上， \(y\)节点对这个\(x_2\)的导数，这个\(\overline{x_2}^2\)；然后我们再去把求得的这两个分支上的导数进行相加（\(\overline{x_2}^1+\overline{x_2}^2\)），就得到了最终\(y\)对\(x_2\)这个结点的导数。

那这样的一种自动求导的方法，它实际上是兼具了数据求导法和符号求导法的优点，它很容易操作，并且不会有表达式膨胀的问题。

像神经网络这种模型，它通常输入是上万到上百万维，那输出的损失函数只有一维，那这样的模型我们只需要一遍自动求导的过程，就可以求出输出对各个输入的导数

各种求导方式对比

我们对四种求导方法去进行了一个比较

• 从精度上来看，手动求解法、符号求导法和自动求导法都有比较高的精度

• 从对图的遍历次数上来看，自动求导法它只需要对图遍历\(N_O+1\)，这里的\(N_O\)指的就是神经网络层的这个输出个数。（备注：+1加的应该是正向传播；有几个输出，就要反向传播几次，所以是\(N_O\)）

所以使用自动求导法，它对于输入维度较大的情况下，它的性能优势是非常明显的

PyTorch中的自动求导

在 Python 之中使用的是AutoGrad这个自动微分引擎，用户只需要一行代码，也就是tensor.backward()就可以调用它来去自动的计算梯度，并且进行反向传播

在 PyTorch 中的这个 AutoGrad模块里面，这个 backward 函数它的实现步骤主要有三步：第一步是正向图的解析（获取反向图根节点和梯度节点），第二步是构建反向计算图的节点（获取反向图的边），第三步是进行反向梯度的传播

那我们先来看第一步正向图的解析；这个里面列出来了PyTorch框架中的源代码部分，对于正向图解析这一部分，这个代码在这部分里面我们会根据输入配置了roots和 grads这两个变量，其中roots是反向传播的根节点的集合，同时也是前向传播的输出节点的集合，grads它是反向传播需要的梯度节点集合

第二步是构建反向计算图的节点；我们通过遍历正向图中的节点，构建反向计算图中的边output_edges：对inputs中的每个元素都获取它对应的张量，检查这个张量是否是（正向图中的）叶节点，也就是说它是否有梯度函数grad_fn；它如果不是叶子结点的话，就表明这个结点存在梯度函数grad_fn，那么我们就创建一条从梯度函数到该节点的边；这样遍历完所有的输入，我们就构建出了反向计算图中所有边的集合output_edges

第三步就是进行反向梯度传播，此时反向计算图已经构建完成了，那么我们就调用这个engine.execute(...)来去进行反向的梯度传播

计算图执行模块

下面我们来介绍编程框架中的计算图执行模块

那在这一节里面，我们会讲述如何将给定计算图中的张量以及操作，映射到给定设备上具体执行的整个过程

我们会首先介绍在编程框架中设备管理的方法，然后介绍张量的实现方法，最后讲解计算图中的算子是如何完成执行的

设备管理

首先来看设备管理，设备是编程框架中计算图执行时的硬件实体，每个设备都具体负责计算子图中的张量存放和算子计算

那么常见的设备包括通用的处理器，也就是CPU，以及我们领域专用的处理器，比如GPU，或者是深度学习处理器DLP

那通用处理器 CPU 的管理方法是比较简单的，这里我们就不进行介绍了

在编程框架的开发里面，主要需要添加的是对领域专用的处理器，也就是 DLP 的这个设备管理的支持，它主要包含包括了三个模块，分别是设备操作、执行、流管理以及设备管理的模块

PyTorch中的设备被直接按照类型去进行分类，比如说CPU，CUDA， DLP 等等

一个设备是由一个类型和一个设备索引或者是序列来去唯一标识的它；前者指定了机器的类型，比如说是 CPU 或者是 GPU 等等，那后者是在有多个特定类型的计算设备时去标识特定的计算设备

下图中parse_type的实现在：pytorch/c10/core/Device.cpp at v2.4.0-rc7 · pytorch/pytorch (github.com)

为了支持设备管理，Pytorch它在DeviceGuardImplInterface.h里面定义了抽象的设备管理类，叫做DeviceGuardImplInterface。这个抽象的管理类，它提供了对设备管理统一的一个抽象接口，那在这个DeviceGuardImplInterface类里面提供了对设备操作、执行流管理和事件管理的函数接口设计

执行流是设备上抽象出来的管理计算任务的软件概念，用于在领域专用处理器上的异构编程模型下，完成设备上任务执行的下发和同步的操作。那具体来说，下发到同一个执行流中的任务具有串行性，下发到不同执行流上的任务能并发执行；因此在编程时候，用户是可以创建多个执行流，并将计算任务分配到不同的执行流中的，从而达到任务并发执行的效果；典型的执行流管理操作包括了执行流创建、执行流同步和执行流销毁等

事件也是设备在软件层面抽象出来的概念，它主要是用来表示设备上任务运行的状态和进展，比如记录事件之间的时间间隔，从而计算设备运行时间等；事件管理它主要包括计算创建事件创建、时间记录和事件销毁等基本操作

那我们来看一下在这个DeviceGuardImplInterface这个类里面，我们就可以看到它提供了对设备操作、执行流管理和事件管理的函数接口设计

https://github.com/pytorch/pytorch/blob/v2.4.0-rc7/c10/core/impl/DeviceGuardImplInterface.h#L57

张量实现

逻辑视图与物理视图

再来看张量实现，张量是神经网络算法里面使用的基本的数据结构，它也是计算图里面的核心概念之一，对应了计算图中不同张量操作之间传递流动的数据

张量它有一些基本的属性，像形状、布局、步长、偏移量、数据类型和设备等等，这些基本的属性可以统称为张量数据结构的逻辑视图；张量数据结构的逻辑视图是编程框架使用者在软件层面上直接控制和表达的一些基本属性

那对编程框架开发者来说，我还需要去维护张量数据结构的另外一种视图，就是物理视图。张量数据结构的物理视图，它主要包括在设备上的物理地址、空间大小、指针、数据类型等属性；物理视图是编程框架底层需要维护的基本属性，它对编程框架使用者来说是不可见的

那我们下面这张表就是张量数据结构的逻辑视图和物理视图的一个基本属性的对比

张量数据结构(逻辑视图、物理视图示例)

在逻辑视图里主要是通过两个关键变量，也就是偏移量和步长，来去最终确定它所对应的张量在物理视图中物理地址空间的寻址方法

我们先来看下左图，当我们使用A[:, 0]来去索引张量A的第一列数据的时候，这个切片以后的张量的逻辑视图就变成了一个形状为[2]的张量[1, 3]，然后它的步长为2，它的数据类型是int32

备注：上面的这个形状是我拿PyTorch试出来的，下面的行切片也是，和课程中老师说的不太一样。但神奇的是，如果在选切片时，使用的是A[: , 0:1]，出来的形状就是[2,1]了...

>>> a = torch.randn(2,2)
>>> a[:, 0].shape, a[1, :].shape
(torch.Size([2]), torch.Size([2]))
>>> a[:, 0:1].shape, a[1:2, :].shape
(torch.Size([2, 1]), torch.Size([1, 2]))

那这样的一个切片操作，它其实并没有隐式的去创建一个新的张量并且拷贝，而只是提供了原本物理视图下的一个新的逻辑视图；那它的物理视图仍然是这个物理地址空间从0x10这个位置开始连续存储的一块数据，但是由于它的步长等于2，所以在进行物理地址空间寻址的时候，每访问一个元素，我都要跳跃两个元素，从而对应了是0x10和0x18两个位置的数据

同理，这个下图右边的这个图，当我们使用张量的索引方法 A[1, :]来去索引张量的第二行数据的时候，这个张量的逻辑视图就变成了一个形状为[2]的这个张量；它的步长为1，数据类型仍然是int32；这个时候我们就需要在逻辑视图中额外引入偏移量的属性，来记录这个新的逻辑视图对应的张量在物理视图上数据实际开始的位置

PyTorch中的张量抽象

我们再来看一下编程框架中对张量的这个支持。右边这张图显示了 PyTorch 编程框架中实现张量的数据调用关系

那么张量在PyTorch 之中，它存在一个和张量对应的类Tensor，这个 Tensor 它是继承了它的基类 TensorBase而来；TensorBase类再进一步去调用相关的结构体，最终支持张量在不同的设备类型，像CPU， GPU 和 DLP 上得到支持

在 Python 中是通过张量的抽象类Tensor以及这个存储抽象类Storage来去分别表示张量数据结构中的逻辑视图和物理视图

其中TensorImpl类，它是张量抽象的实现，包含了维度信息、步长信息、数据类型、设备布局等逻辑视角的张量信息

StorageImpl类，它是张量的一个存储实现，包含了内存、指针、数据总数等物理视角的张量信息，它会调用结构体allocator去进行张量数据空间的一个分配

那么我们来看一下底下这张图，就显示了在PyTorch编程框架中张量实现的流程

首先Tensor类是由TensorBase类继承而来的，在TensorBase类中包含了唯一的成员变量impl_，这个相当于一个指向TensorImpl的指针，并且它表达了前述张量数据结构中的逻辑视图

TensorImpl又利用结构体Storage来去表示和张量存储相关的信息

一个Storage代表一个张量的底层，支持数据缓冲区(原视频也不太清楚(?TODO))，并且唯一的拥有一个指针storage_impl_， Storage（还是Impl?(?TODO)）会调用结构体allocator去进行张量数据空间的分配

根据后端设备的不同， allocator的代码实现也会有不同

张量内存分配

张量的数据结构，包含了逻辑视图和物理视图两种表示；从逻辑视图到物理视图的转换需要完成对张量的内存，也就是对张量进行内存管理

那根据设备的类型不同，张量管理的方式也有不同，对于 CPU 来说，主要采用即时分配的这个方式；对GPU来说，主要采用的是内存池分配的方式

即时分配

我们先来看 CPU 中采用的即时分配方式，当需要分配张量内存的时候，就立刻从系统中申请一块合适大小的内存，这个就是即时分配；核心代码部分就是malloc分配空间以及free释放空间

内存池分配

内存池分配，它是一种预先分配一块固定大小的内存池，然后在需要时从内存池中分配内存的策略；分配的内存它来自事先分配好的内存块，而不是每次都向系统申请新的内存

内存池分配不仅仅是提供了简单的内存的预分配，它还具备一个自我维护的能力，能够灵活地处理内存块的拆分和合并等等操作

那在领域专用处理器，比如 GPU 或者是 DLP 上，需要使用处理器的运行时接口来去管理设备端内存，通常就会通过内存池的方式去管理设备端的张量

这种做法相比于调用系统 API 直接分配内存来讲，它不仅减少了系统调用开销，同时还具备两个重要的优势，一个是节约了设备内存的使用，并且它减少了设备内存的碎片化问题

分配好内存以后，就可以建立张量。建立张量时需要对张量进行初始化，那这个初始化的步骤是首先根据固定内存标志去选择相应的分配器allocator；这个 allocator可以是即时分配，也可以是内存池分配。

在我们下面这个例子里，使用的是 CPU 内存分配，也就是即时分配

然后我们会调用这个_empty_generic去进行具体的实现；这个_empty_generic它实现了通用的张量创建逻辑，也就是说我首先去分配存储空间，然后去创建StorageImpl类，再去创建TensorImpl类

在其他的硬件后端，比如说 GPU 或者 DLP 上去进行张量初始化的时候，我们只需要去选择相应的分配器，并且设置对应的参数，然后再调用_empty_generic就可以实现

算子执行

我们再来看算子的执行计算图的过程，它可以被分解为每个算子单独执行的过程

• 首先通过计算图生成一个执行序列，这个序列会确定算子的执行顺序来确保正确的数据流和依赖关系
• 然后我们针对每个算子去进行算子实现，包括前端定义、后端实现和前后端绑定三个步骤
• 最后进行分派执行，包括查找适合给定输入的算子实现并调用相应的实现来执行具体的计算任务

执行序列

我们先来看执行序列的确定。计算图它描述了算子之间的依赖关系，通过分析计算图节点之间的依赖关系就可以获得算子的执行序列

这个序列一般可以使用拓扑排序的算法来获得，它的具体的流程是：首先我要计算每一个节点的入度；然后我使用拓扑排序算法对计算图的节点进行一个排序，这个通常会采用深度优先搜索或者是广度优先搜索来去遍历图中的节点；最后我们要检查结果，检查这个执行序列的长度是否是等于图中的这个节点数

像我们这个例子中，我有一个计算图的节点，这个里面包含了 5 个算子，那么我们就使用拓扑排序方法，给出了这 5 个算子的一个执行序列

那么需要注意的是，对同一个计算图，我们可能会得到多个不同的结果；这是因为在求拓扑序列的过程中，可能有多个入度为0的节点，我们可以从中任意选择下一个执行的节点

算子实现

我们再来看算子实现，那在获得算子的执行序列以后，每个算子都需要在编程框架中完成对应的算子实现

在我们的编程框架里面，用户的接口（也就是前端），和具体的实现（也就是后端）一般会采用不同的编程语言，比如说在PyTorch里面，它的用户前端一般是使用Python作为这个前端的编程语言，然后用C++来去作为它的后端的编程语言

算子实现的流程首先是做前端定义，就是指在编程框架中配置算子信息，这里边包含了算子的输入输出以及相关的接口定义，最后再生成前端的接口

第二步是后端实现，就是使用C++或者是其他高级的编程语言来去编写算子的底层实现代码，完成算子的计算逻辑部分的实现

最后是前后端的一个绑定，在这一步骤里面，我们编程框架把前端定义的算子和后端的具体实现去进行一个绑定

前端定义

PyTorch提供了一种高效的模式用于管理整个算子实现模块，我们称为native_function。那在使用native_function这个模式进行算子实现的时候，需要修改配置文件native_functions.yaml以添加算子配置信息

这个native_functions.yaml这个函数中包含了几个字段：

第一个字段是func，func字段定义了算子的名称和输入输出的参数类型

第二个字段是variants，这个字段表示需要自动生成的高级方法

第三个字段是dispatch，这个字段表示该算子所支持的后端类型和对应的实现函数

那下面这段代码示例就显示了这个native_function函数的格式，我们可以看到里面包含了这个 func字段、 variants字段以及dispatch字段

我们以 PyTorch 中的PReLU这个算子实现为例，说明使用 native function 模式在 PyTorch 中实现一个以 CPU 为后端运行的算子的流程

那我们这张代码，它展示了 PyTorch 中PReLU算子的配置文件，也就是 native 函数，包含了PReLU算子实现部分：PReLU正向传播函数实现和PReLU反向传播函数实现

这个操作，它接受两个张量作为输入，其中self表示输入张量， weight表示该操作的可学习的参数

dispatch字段，它表示PReLU支持的后端类型和对应的实现函数

PReLU算子的前端实现代码如图所示，这段代码实现了PyTorch中PReLU类的定义，包括PReLU类的相关描述以及正向计算逻辑和可学习参数

然后需要在配置文件中添加算子正向传播函数和反向传播函数的对应关系

我们这段示例代码就表明了正向传播函数_prelu_kernel，它对应的反向传播函数是_prelu_kernel_backward

后端实现

后端实现步骤包含了算子的表层实现和底层实现两部分

表层实现

表层实现可以看作不同设备之间的抽象函数接口；底层实现可以看作具体到某个设备上的实际代码实现，表层实现和底层实现的代码中均需要包含正向传播函数和反向传播函数，二者的结构是类似的

首先我们创建一个空的返回对象，然后通过配置TensorIteratorConfig来去创建一个TensorIterator用于迭代张量操作

这个Iterator，它提供了统一的计算抽象，封装了正向计算的输入权重以及反向计算的梯度，接下来调用prelu_stub函数进行实现，完成PReLU算子的运算，最后返回计算结果

这里需要注意的是，这个代码中定义的实现只是一个封装，它没有完成真正的实现，还需要根据后端的硬件来去编写对应的这个底层实现

底层实现

底层实现指的就是具体到某个设备上的实际代码实现，比如说在 CPU 上的实现，或者是 GPU 上的实现，或者是 DLP 上的实现

那底层实现中的这个PReLU kernel 函数和表层实现中的 PReLU stub函数会在前后端绑定中去完成对应，这种将表层实现和底层实现解耦的设计使得为多种后端注册内核函数时可以复用相关的接口

前后端绑定

我们再来看前后端绑定，在前后端绑定的步骤中，我们为每个后端编写了相应的底层实现步骤，以实现各种不同的硬件和软件平台，并且为每一种输入情况都提供了对应的后端实现的重载版本

在前后端绑定中， dispatch它在分派机制中扮演着调度和控制的角色，确保在不同的后端环境中能选择正确的实现方法

这个dispatch它会维护一个分派表，分派表的表项记录着算子到具体的后端实现的对应关系

分派表可以视为一个二维网格，纵轴表示PyTorch支持的算子，横轴表示支持的分派键；这个分派键是和后端相应的标识符；这个表初始是为空的，那当添加一个算子到后端实现的对应关系的时候，就需要编写TORCH_LIBRARY_IMPL这个函数去进行注册

分派执行

在获得算子的执行序列并且实现了对应的算子以后，就需要对算子进行分派执行；它指的是在运行时根据输入张量的类型和设备类型，查找并调用合适的算子实现方式

在分派执行过程中， Dispatcher会首先根据输入张量和其他的信息计算出对应的分派键，然后由该分派键找到相应的内核函数

这里涉及到几个概念，算子，指的是 Dispatcher 的调度对象，它代表了具体的计算任务；分派键，它是根据输入张量和其他信息计算，它可以简单的理解为与硬件平台相关联的标识符；内核函数，指的就是特定硬件平台上实现算子功能的具体代码