工作原理
单算子执行流程
ATB内单算子的执行流程分为以下几个步骤：kernel图构建、device内存计算、tiling data计算与搬移、计算任务下发。
- kernel图构建
kernel是在device上运行的基本代码单元，类似于C语言中的函数，device上基本以kernel为单位执行各种计算任务，kernel图的构建本质是kernel任务队列的构建，或者说是device任务流的构建。
ATB对外提供的Operation都是具有相对复杂功能的kernel组合体，以图的形式组织各个kernel之间的关系。由于同一个operation在使用不同特性或不同输入时会组合不同的kernel，因此单Operation对应的kernel图或device任务流，只有在运行时才能确定。
- 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用于存放临时数据的空间。
- device内存计算
对于一个kernel图而言，提供了整图的输入输出张量后，还需要分配图中间的张量。因此除了每个kernel的tiling data、scratch memory以外，ATB还将准备好kernel图的中间张量。
单kernel可当作中间张量大小为0的kernel图，ATB在计算单算子所需的device内存时，都是以kernel图来计算的。
ATB会计算kernel图所需的中间张量大小以及每个kernel的scratch memory大小，再将其作为WorkSpaceSize返回给用户。而tiling data则是在计算出大小后由ATB统一分配与管理。
- tiling data计算与搬移
tiling data的计算通常放到host侧，tiling data在host侧计算完毕后，ATB再将其拷贝到device侧，作为kernel的输入提供给kernel。
- 计算任务下发
在任务执行队列已确定且准备好输入，最后一步就是下发任务给device侧。
在这一步骤中，ATB会根据前面构造好的kernel队列，将准备好的输入作为入参依序给到device任务下发接口，最后等待device侧完成任务执行即可。
Tiling data搬移策略
ATB存在三种不同的tiling data搬移策略：tiling整体搬移、多stream搬移、tiling随kernel下发搬移，当前默认使用的是tiling随kernel下发搬移这一方式。
- tiling整体搬移方式会把每次计算出的tiling data存放到一片连续的host内存中，待所有kernel的tiling data计算完成后，再一次性搬移到device侧。
- 多stream搬移方式则是针对整体搬移方式做出的改良，其核心思路是通过stream并行的方式来减少tiling data搬移所消耗的时间。
在该搬移策略中，ATB会准备好两个stream、一个环状的device缓冲区以及一系列同步信号量。
- 其中一条stream用于kernel的执行，另一条则单独用于tiling data的拷贝。由于kernel的执行依赖于tiling data拷贝的完成，此时需要用同步信号量来保证另一条stream上的kernel执行动作位于当前kernel的tiling data拷贝完成之后。
- 环状的device内存缓冲区则是用于处理在tiling data拷贝速度快于kernel执行速度时，提前拷贝到device侧的tiling data数据可以保存下来且不互相冲突。
图1 多stream搬移策略介绍
- tiling随kernel下发搬移方式是默认使用的tiling data搬移策略。
该策略对tiling整体搬移方式进行了性能优化，但优化方式与多stream搬移方式不同。
tiling随kernel下发搬移的核心思路是：
- 不再等待所有tiling data计算完毕后再一次性搬移到device侧，而是每次计算完一个kernel的tiling data就搬移至device侧。
- 在kernel任务下发至device侧时，同时启动tiling data的搬移任务。这样就可以使host侧与device侧的设备并行，在host侧准备下一个kernel的tiling data的时候，device侧同时执行当前kernel任务，从而显著提升了tiling data的搬移效率。
如图2所示，相较于多stream搬移方式，在tiling data搬移速度过快时会导致的device缓存区不足，从而导致tiling data被覆盖，tiling随kernel下发搬移的方式不受两者（tiling data拷贝速度与kernel执行速度）速度的限制，且性能优化也更进一步。
图算子执行流程
ATB内部通过一个有序的Operation流来处理复杂的图结构，用户可通过调用ATB接口来完成计算图的构建与优化。
在使用ATB时需要先梳理整个计算图及其计算流，将计算图的拓扑结构转为FIFO结构表示，再根据梳理好的FIFO队列构造ATB的图参数，计算图与FIFO队列转换请参考图3。
在获取到有序的Operation流后，ATB就会遍历队列中每一个Operation，并执行对应的单算子执行流程。
性能优化机制
从实际数据来看，整个构图流程的host耗时较长。为减少host侧耗时，进一步优化端到端的性能，ATB提供了四种对host侧性能进行优化的机制：Tiling Cache、Setup复用、InferShape复用与Runner Pool。
- Tiling Cache
Tiling Cache的作用是缓存kernel的tiling data。
根据transfomer结构模型的特点，推理过程中大量kernel的tiling data实际可以进行复用，因此ATB会对已计算的tiling data进行缓存，当检测到可复用tiling data时，将直接通过缓存中获取而不是重复计算，从而节省了大量tiling data计算的时间。
如图4所示，ATB内包含两种Tiling Cache，本地Tiling Cache与全局Tiling Cache（以下简称为本地Cache与全局Cache）。
- 本地Cache（Local Cache）存储在Operation对象内，只能被当前的Operation对象读取或写入，其包含的Cache槽位数由环境变量控制。
- 全局Cache（Global Cache）是线程级的对象，在同一个线程内的所有Operation对象都可读取或写入。与本地Cache不同，全局Cache是一组Cache的总和。
一个全局Cache中包含多少个Cache由当前ATB支持多少种Operation决定，其中每个Cache的Cache槽位数都由环境变量来控制。
- 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复用性能优化效果不明显。
- 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控制。
图5 Runner Pool结构示意图
内存优化与管理机制
内存优化与管理机制主要涉及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是作为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。内存池中的内存块数用户可通过环境变量进行配置，每块内存的大小当前暂不支持配置。
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小节。Runner池存放在context中，作为公共资源供Operation对象使用。
- 溢出检测张量用于存放溢出检测算子的输出结果，其中包含的device内存会跟随Context统一申请或释放。