1 ATB 介绍

Ascend Transformer Boost 加速库（下文简称为 ATB 加速库）是一款高效、可靠的加速库，基于华为 Ascend AI 处理器，专门为 Transformer 类模型的训练和推理而设计。

ATB 加速库采用了一系列优化策略，包括算法优化、硬件优化和软件优化，能够显著提升 Transformer 模型的训练和推理速度，同时降低能耗和成本。具体来说，ATB 加速库通过优化矩阵乘法等核心算子和注意力机制的实现方式，实现了对 Transformer 模型的高效加速。此外，ATB 加速库还充分利用了 Ascend AI 处理器的硬件特性，如算力、存储带宽和内存带宽，通过硬件加速和数据重用等技术，进一步提升了性能和效率。ATB 加速库目前提供了底层基础的高性能算子以及高效的算子组合技术（Graph 图算子），同时上层支持对接多种模型框架如 PyTorch、MindSpore、Paddle。

总而言之，ATB 加速库中包含了各类 Transformer 类模型的高度优化模块，在各种应用场景中发挥重要作用，为模型的训练和推理提供了强有力的支持。

2 ATB 的软件架构

ATB 架构图从上述架构看，ATB 优化的核心内容是在算子计算方面，通过优化算子计算的方式，比如：增加算计算并行机会，优化内存排布等。涉及到 3 种算子，如下所述。

2.1 基础 Operation（原生算子）

用户可以根据需求使用对应的算子完成计算功能。这类算子为一系列基础算子，提供了如矩阵乘、转置等功能。详细信息请参考 atb/infer_op_params.h 和 atb/train_op_params.h。

2.2 插件（Plugin）机制（插件算子）

自定义插件算子（PluginOperation）是一种为用户实现特定功能提供的机制。如果一些功能通过单算子或图算子无法实现，用户可以通过开发自定义插件算子实现对应的功能。

2.3 Graph Frame（图算子）

提供图算子（Graph）机制，用户根据模型设计对应的图算子，使用加速库提供的原生算子和创建的自定义算子创建图算子，完成相应的计算。

graph 也就是图，其实就是将多个算子组合成一个图的形式进行调用。那么组合成图的形式有什么好处呢？

1 算子调用

说明这个问题前，首先来了解下算子的下发过程。举例来说，如下图：

算子 host->device 下发过程一个算子的下发，过程大概要经过 host（CPU）-> Device(Ascend, GPU...)的过程。

在这个过程中，Host 做的事情大概又可以做如下划分：

1、python 侧。模型一般是用 python 编写，算子首先发起调用是从 python 侧的接口；

2、c++侧。python 侧发起后，一般是通过 pybind 技术，路由到 C++侧，这个过程一般由 AI 框架完成，如 torch 或、mindspore 等。在 c++侧做的事情，一般有算子的 infer（mindspore），算子输入输出、workspace 内存申请，调用 Device 侧的接口（驱动层）。

那么明显的，从 host->devcie 是有个调用过程，这个过程肯定是存在开销的。

那如何能减少这个开销呢？ATB 中 graph frame 的实现，就是将算子组成图的形式，将其整体一次性下发到 device 上 ，那么 host->device 的开销是可以极大减少的。这样端到端下的性能也就提升了。

2 ATB 图算子构建

用户需要自行设计并定义图的结构，即图结构中节点（Node）的组合与依赖关系，包括节点对应的单算子、节点的输入 Tensor 与输出 Tensor，并识别这些 Tensor 为图的输入 Tensor、输出 Tensor 和中间 Tensor。

图输入 Tensor，为用户使用图算子时需要从外部输入的所有 Tensor。

图输出 Tensor，为用户使用图算子时不会再进行下一步运算操作的所有 Tensor。

图中间 Tensor，为图算子运算中产生的临时 Tensor。例如同一图算子中有算子 A 和算子 X，算子 A 的输出 Tensor 为算子 X 所用，且该 tensor 对于图外部模块不可见，则该 Tensor 为中间 tensor。

如下图所示，该示例图算子由两个节点组成，这两个节点均为 Elewise_Add 算子。a、b、c 为图算子的三个输入 Tensor，output 为图算子的输出 Tensor，a_add_b_output 为图算子的中间 Tensor。其中，节点 0 的输入为 a 和 b，输出为 a_add_b_output，节点 1 的输入为 a_add_b_output 和 c，输出为 output。

用户需要自行设计并定义图的结构，即图结构中节点（Node）的组合与依赖关系，包括节点对应的单算子、节点的输入 Tensor 与输出 Tensor，并识别这些 Tensor 为图的输入 Tensor、输出 Tensor 和中间 Tensor。图输入 Tensor 为用户使用图算子时需要从外部输入的所有 Tensor。图输出 Tensor 为用户使用图算子时不会再进行下一步运算操作的所有 Tensor。图中间 Tensor 为图算子运算中产生的临时 Tensor。例如同一图算子中有算子 A 和算子 X，算子 A 的输出 Tensor 为算子 X 所用，且该 tensor 对于图外部模块不可见，则该 Tensor 为中间 tensor。

如下图所示，该示例图算子由两个节点组成，这两个节点均为 Elewise_Add 算子。a、b、c 为图算子的三个输入 Tensor，output 为图算子的输出 Tensor，a_add_b_output 为图算子的中间 Tensor。其中，节点 0 的输入为 a 和 b，输出为 a_add_b_output，节点 1 的输入为 a_add_b_output 和 c，输出为 output。

3. ATB 的算子执行流程

如上所述，ATB 存在三种算子，那么每种算子的设计思考是什么，与 Ascend 上，或者说 CANN（CANN：CANN-昇腾社区 (hiascend.com)）中的算子区别是什么。由于篇幅问题，请移驾到下述文章。ATB三种算子的执行区别

本文主要参考：ATB加速库

加速库使用指导