ATB算子实现原理解读

1 前言

从前文ATB是什么？ - 知乎 (zhihu.com)可知，在Ascend Transformer Boost加速库（ATB）中存在三种类型的算子，那么这三种类型的算子，它们的执行流程是什么，和其它的CANN中的算子有什么区别。带着这些疑问，进入到本文内容。

2 实现一个ATB算子

阅读本文之前，可以先行看下前文

https://zhuanlan.zhihu.com/p/720148106

对atb的开发流程，有个大概的感知。其大概分为以下10步：

ATB算子实现步骤

每一步都有具体的实现参考。那么这些步骤的背后，往往就是ATB算子的设计逻辑。本文件结合文章

https://www.hiascend.com/document/detail/zh/canncommercial/80RC22/developmentguide/acce/ascendtb/ascendtb_0039.html

中的描述，来做一个梳理和分析。

3 ATB算子工作原理

3.1 单算子执行流程

首先，看看单算子执行流程，如下图所示：

单算子执行流程

单算子的执行流程主要有5个部分：

1、kernel图构建

kernel是在device上运行的基本代码单元，类似于C语言中的函数、GPU的.cu文件等，device上基本以kernel为单位执行各种计算任务。

kernel图的构建本质是kernel任务队列的构建，或者说是device任务流的构建。 ATB对外提供的Operation都是具有相对复杂功能的kernel组合体，以图的形式组织各个kernel之间的关系。由于同一个operation在使用不同特性或不同输入时会组合不同的kernel，因此单Operation对应的kernel图或device任务流，只有在运行时才能确定。

2、kernel运行的必要输入

要使得一个kernel任务流正常执行，需要在device侧为每个kernel都准备好对应的输入。 对于单个kernel而言，有三种必要的输入类型：输入输出张量、tiling data、scratch memory。这三种输入都是以device侧地址的形式提供给kernel的。

• 输入输出张量一般由用户提供，ATB将会准备好该kernel所需的tiling data与scratch memory。
• tiling data是kernel的分片参数，用于决定kernel实际计算时的分片策略，通常以结构体的形式存储，由用户输入的参数与张量Shape计算而来。
• scratch memory则是kernel用于存放临时数据的空间。

3、device内存计算

对于一个kernel图而言，提供了整图的输入输出张量后，还需要分配图中间的张量。因此除了每个kernel的tiling data、scratch memory以外，ATB还将准备好kernel图的中间张量。 单kernel可当作中间张量大小为0的kernel图，ATB在计算单算子所需的device内存时，都是以kernel图来计算的。 ATB会计算kernel图所需的中间张量大小以及每个kernel的scratch memory大小，再将其作为WorkSpaceSize（如调用Setup(variantPack, workspaceSize, context)）返回给用户。

而tiling data则是在计算出大小后由ATB统一分配与管理。

4、tiling data计算与搬移

tiling data的计算通常放到host侧，tiling data在host侧计算完毕后，ATB再将其拷贝到device侧，作为kernel的输入提供给kernel。

tiling可以参考文章：

https://zhuanlan.zhihu.com/p/720175682

5、计算任务下发

在任务执行队列已确定且准备好输入，最后一步就是下发任务给device侧。 在这一步骤中，ATB会根据前面构造好的kernel队列，将准备好的输入作为入参依序给到device任务下发接口，最后等待device侧完成任务执行即可。

3.2 图算子执行流程

ATB内部通过一个有序的Operation流来处理复杂的图结构，用户可通过调用ATB接口来完成计算图的构建与优化。

在使用ATB时需要先梳理整个计算图及其计算流，将计算图的拓扑结构转为FIFO结构表示，再根据梳理好的FIFO队列构造ATB的图参数，计算图与FIFO队列转换请参考下图。

在获取到有序的Operation流后，ATB就会遍历队列中每一个Operation，并执行对应的单算子执行流程。

计算图与FIFO队列的转换

计算图与FIFO队列转换

4 ATB中的host性能优化

ATB提供了四种对host侧性能进行优化的机制：Tiling Cache、Setup复用、InferShape复用与Runner Pool。

4.1 Tiling Cache

Tiling Cache的作用是缓存kernel的tiling data（参考：ascend 算子tiling - 知乎 (zhihu.com)）。

根据transfomer结构模型的特点，推理过程中大量kernel的tiling data实际可以进行复用，因此ATB会对已计算的tiling data进行缓存，当检测到可复用tiling data时，将直接通过缓存中获取而不是重复计算，从而节省了大量tiling data计算的时间。

如下图所示，ATB内包含两种Tiling Cache，本地Tiling Cache与全局Tiling Cache（以下简称为本地Cache与全局Cache）。

ATB内的两种Tiling Cache

• 本地Cache（Local Cache）存储在Operation对象内，只能被当前的Operation对象读取或写入，其包含的Cache槽位数由环境变量控制。
• 全局Cache（Global Cache）是线程级的对象，在同一个线程内的所有Operation对象都可读取或写入。与本地Cache不同，全局Cache是一组Cache的总和。

一个全局Cache中包含多少个Cache由当前ATB支持多少种Operation决定，其中每个Cache的Cache槽位数都由环境变量来控制。

4.2 Setup复用与InferShape复用

对于单个特定Operation而言，若两次输入的Shape与参数相同，则该Operation两次构造的kernel图也是相同的。

基于这个结论，ATB提供了一种跳过kernel图构造这一步骤的机制，即Setup复用，每个Operation对象会存储自己上一次的输入张量并记录参数是否有被修改，每次Operation对象进行kernel图构造前，都会检查当前输入张量的个数与Shape与上次输入是否相同、参数是否有被修改，若输入相同且参数未被修改则会跳过kernel图构造这一步骤，直接使用上次构造好的kernel图。

InferShape复用与Setup复用类似，当同一个Operation对象两次输入的shape与参数相同时，就会跳过该Operation的InferShape步骤。对于图算子来说，图算子的InferShape是图内的单算子通过链式推导得来的。当整个计算图特别复杂庞大时，InferShape就成为host侧最主要的性能开销，此时可使用InferShape复用机制显著优化性能。

该优化手段当前只针对图算子生效，由于单算子的InferShape逻辑复杂度较低，此时使用InferShape复用性能优化效果不明显。

4.3 Runner Pool

Runner是Operation的执行单元，可以理解为Operation是面向用户的前端，而Runner则是真正处理逻辑的后端。

在重复多次创建与释放Operation对象的场景下，Runner的创建耗时占据了host侧耗时较大部分。为减少Runner的创建开销，ATB新增了Runner Pool这一特性。在使用Runner Pool的情况下，ATB每次在创建Runner时，需要先从Runner Pool中寻找是否有可以使用的Runner，有则直接使用Runner Pool中的Runner，否则就创建新Runner并放到Runner Pool中。Runner Pool存放于Context中，每个Runner类型都有一个自己的Runner Pool，每个Runner Pool中存放有多个Runner槽位，该槽位数量可通过环境变量ATB_RUNNER_POOL_SIZE控制。

Runner Pool

5 ATB中内存优化与管理机制

内存优化与管理机制主要涉及ATB对device内存的计算与管理机制。

一个算子下发所需要的device内存空间分为三部分：中间张量内存、kernel的scratch memory和tiling data内存，下面将分别讲述这三部分内存在ATB中是如何进行计算及管理。

• 中间张量由于不作为整个kernel图的输入或输出，它的生命周期只有在执行到第一个以其作为输出张量的kernel时才开始，在所有依赖该中间张量作为输入的kernel运行结束后，其生命周期就立即结束。因此在整个kernel执行队列中，任一时刻存在哪些中间张量是可知的，整个kernel执行流程实际是不断对中间张量进行释放与分配的过程，在该过程中占用device内存的最大值可作为中间张量所需内存。
• kernel的scratch memory用于存放kernel运行时的一些临时数据。由于ATB在同一时刻仅运行一个kernel，因此不同kernel之间可以复用同一块scratch memory，其大小可取所有kernel所需的容量中的最大值。
• tiling data内存是三部分内存中占用相对较少的部分。ATB对tiling data内存采取的策略是一次性分配所有kernel的tiling data总和的内存空间。

kernel的scratch memory与tiling data的计算与管理

上述三部分内存中，中间张量内存与kernel的scratch memory是作为workspace由用户进行分配的，kernel的tiling data由ATB进行管理和分配。

• ATB在Setup接口内会计算好整图所需要的中间张量内存与scratch memory，两者相加后将其作为WorkspaceSize返回给用户。
• 用户根据Setup接口返回的WorkspaceSize申请device内存，再将申请的device内存通过Execute接口传递到ATB，ATB会根据之前的计算结果来使用该内存。
• kernel的tiling data存储在ATB的context类中，用户在构造context类时默认会生成一个32 * 3 mb大小的device内存池，每当一个Operation需要对tiling data进行搬移时就会从context中取出一块大小为3mb的device内存用于存放tiling data。内存池中的内存块数用户可通过环境变量进行配置，每块内存的大小当前暂不支持配置。

ATB内存来源

6 Context介绍

Context类是用于存放与管理ATB内各种公共资源的类，其包含了以下资源：两条stream、控制时序的事件、host内存池、device内存池、Runner池、溢出检测张量。

• 两条stream分别用于kernel执行与tiling data的拷贝，kernel执行的stream由用户设置，tiling data拷贝的stream则由ATB本身来创建。当不开启多stream功能时，用于tiling data拷贝的stream将不会创建。
• 控制时序的事件用于保证多stream功能中tiling data拷贝与kernel执行的顺序正确。
• host内存池是一个环状的host内存缓冲区，用于存放host侧计算出来的tiling data，其块数可通过环境变量控制。
• device内存池的结构和用法都与host内存池类似，区别在于其中的内存块都为device内存块，其块数可通过环境变量控制。
• Runner池的详细介绍见性能优化章节中的Runner Pool小节（4.3 Runner Pool）。Runner池存放在context中，作为公共资源供Operation对象使用。
• 溢出检测张量用于存放溢出检测算子的输出结果，其中包含的device内存会跟随Context统一申请或释放。