1 前言

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

2 实现一个 ATB 算子

阅读本文之前，可以先行看下前文如何使用Ascend的ATB加速库？

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

ATB 算子实现步骤每一步都有具体的实现参考。那么这些步骤的背后，往往就是 ATB 算子的设计逻辑。本文件结合文章工作原理-进阶专题-Ascend Transformer Boost加速库-领域加速库开发-CANN商用版8.0.RC2.2开发文档-昇腾社区

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

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 可以参考文章：Ascend算子tiling

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

3.2 图算子执行流程

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

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

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

计算图与 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 统一申请或释放。