飞桨首创算子内核开发KP体系，快速适配硬件，高效优化性能
随着深度学习的蓬勃发展，其应用领域也愈加丰富，对深度学习框架的计算算子需求也随之增多。一个成熟的框架，其算子数目通常达到几百甚至上千。面对同样繁荣发展的 AI 硬件（加速器）市场，在算子适配层面，深度学习框架与硬件的通常做法是，对每个算子在各个硬件平台上全部单独实现。这导致了各硬件的算子代码无法系统性复用，甚至相同的优化策略也要在不同硬件平台上重复实现。无论是框架还是硬件，其开发成本都是巨大的。
针对这一系列问题，飞桨在业内首创提出一套高效的算子内核开发体系——Kernel Primitive API（简称 KP）。通过对通用功能的封装及优化，实现跨硬件平台的算子 Kernel 代码复用。不同硬件间，不仅可以共用相同的算子计算逻辑，而且可以共用部分优化逻辑，这极大提升了框架与新硬件适配的效率。
例如，通过 KP 体系，飞桨与昆仑芯2代 AI 芯片（基于自研2代架构昆仑芯 XPU-R，以 CR20 为例）的适配，相比传统方式减少了90%的代码量，同时也极大减少了性能优化的工作量，提升了双方的工作效率。
KP 设计思想
当前，神经网络算子加速的主流方法是异构编程法，即在 CPU 端（host 端）进行数据准备和资源分配，在AI加速器端（device 端）进行数据计算。以 reduce 算子实现为例（如图1），首先根据算子计算规则在 CPU 端进行数据空间申请以及 Kernel 启动信息配置，然后在 AI 加速器端进行实际的规约操作。由于 AI 加速器的架构及开发语言各不相同，每新增一种硬件结构就需要针对该硬件重新开发上述算子执行流程。深度学习网络模型中使用到的算子众多，若每种新硬件架构的接入都新增一遍上述流程则工程量浩大。另外当开发者需要新增一个算子，并且扩充到已适配的所有新硬件时，就需要根据相关硬件特性实现 Kernel，对开发者的硬件知识以及开发能力有着很高的要求。
在设计理念上，KP 采用分层结构，将算子在硬件上的实现分成2层结构。
首先是 Kernel 层，调用各类 KP 实现算子计算逻辑，在 Kernel 内统一使用寄存器进行临时数据存储，不依赖于任何 AI 加速器硬件特性。
之后是 API 层，实现 Block 级别的数据读取和数据计算，各类 AI 处理器的 API 层对外接口完全一致，从而保证 Kernel 层代码的可复用性。
不同平台的 API 内部实现略有差异，主要依赖于当前 AI 处理器的硬件特性，将同一硬件适配的 API 放置到固定文件，通过编译宏进行头文件控制，保证不同硬件间的独立性，从而达到屏蔽硬件细节的效果。
根据数据操作规则，KP 可分为数据读写，数据计算两大部分。
数据读写类 API，用于完成全局内存与寄存器间的数据搬运工作。 数据计算类 API 进行通用数据计算，例如加法、求和、排序等操作。
API 接口调用简单，在进行 API 调用时仅需传递相应的数据指针及计算规则即可完成相应的功能支持。使用 KP 开发的算子在多 AI 加速器上的执行流程如图2所示。所有硬件平台下的算子完全共用 CPU 数据准备、Kernel 计算、CPU 数据返回等代码，差异仅在于 KP 具体实现，在进行新硬件接入时仅需增加相应的 API 实现即可。
在具体实现上，KP 内核开发接口是 Block 级别的函数封装，不会增加 kernel launch 和访存开销，内部逻辑使用寄存器作为缓冲，工具链编译器进行自动优化，同时不增加核函数调用次数及全局内存访问次数，因此使用内核开发接口构建算子核函数不会带来额外的性能开销。内核级 API 具有完备性和可复用等特点，API 一次调用可完成一个 Block 内的数据读取和计算操作。例如 reduce API 通过 ReduceMode 标识规约模式，当 ReduceMode 为 kGlobalMode 时说明当前 API 对 Block 内的线程进行规约操作，reduceMode 为 kLocalMode 时进行线程内规约操作，如下图3所示，kGlobalMode 的返回结果为当前 Block 的最大值，注意此处使用到了批量处理，可一次产生多个最值，计算结束后每个线程具有相同的返回结果；kLocalMode 是进行线程内的规约操作，返回当前线程内的最值，可一次返回多组最值，计算结束后每个线程的计算结果各不相同。
使用 KP 加速算子开发
KP 在进行 API 封装时总结了当前主流的数据读写和计算操作，设计并实现了一组通用访存类、计算类函数及计算 Functor。使用6个 API 即可完成大部分elementwise 类、activation 类、reduce 类、以及如 softmax 等复杂算子的功能支持， 能够有效的提升算子开发效率，降低算子开发成本。其中，计算 Functor 是对计算操作的抽象， 用于表示当前 OP的计算规则，例如减法操作，使用 SubFunctor 配合数据读写 API 即可完成 sub 算子实现。
▌代码复用，降低开发量
以当前深度学习模型中经常大量使用的 elementwise 类、activation 类和 reduce 类算子为例，这些类别的算子在某些模型中甚至占50%左右，且具体实现时用到的基础计算操作往往是相同的，比如 reduce、softmax 均会使用到数据规约操作，elementwise、activation类均需连续数据读取以达到高性能实现，因此结合算子功能需求进行通用API封装能够极大的减少相似功能实现，通过进行 API 代码复用，降低开发量。
此外 compare 类与 elementwise 类算子的 Kernel 实现流程基本一致，首先将数据存全局内存读取到寄存器，然后根据特定计算规则对读取的数据进行计算，最后将计算结果写回到全局内存，具体 Kernel 实现差别只有计算规则不同。若每个算子都单独实现一份 Kernel 代码，将会导致大量的冗余代码存在，占用大量的人力，后续的 Kernel 的性能优化依旧需要逐个 Kernel 进行优化，工程量巨大。KP 通过抽象操作规则，将数据操作定义为 Functor 并以函数模板的的形式传入，能够实现通用 Kernel 代码的复用，实现同类算子功能的快速支持。
如上图所示，当 elementwise_add Kernel 实现完成后，将 AddFunctor 替换为 MulFunctor 即可完成 mul 算子的功能支持，Functor 主要在计算类 API 中使用，通过函数模板的形式转入，相比于分支判断，模板传递实现的 Kernel 具有更高的性能优势。
▌简洁易维护
KP 内封装了通用计算操作，将复杂的计算和数据搬运流程下沉到了简单的接口之下，开发者仅需要使用简单的 API 调用即可完成复杂的功能支持，使用 KP 实现的 Kernel 代码具有极度简洁，可维护性高等特点。以 softmax 为例，使用 KP 替换 softmax 原有 Kernel 实现，替换后的s oftmax 与替换前相比，整体 Kernel 性能保持一致，而在某些数据规模下，性能还优于原始 Kernel 实现。使用 KP 实现的 softmax Kernel 代码量从之前的155行减少为30行，大幅降低了 Kernel 开发的工作量，同时替换后的 softmax 代码逻辑更加简洁清晰，与 softmax 的计算公式高度一致，代码可读性和易维护性明显增强。
▌高性能实现
目前基于主流 AI 加速器实现的算子性能主要受访存和计算两个方面的约束，为保证 Kernel 性能，KP 在设计之初就针对 API 特性进行了优化调整。对于 IO 类接口，KP 设计了 Block 边界模板参数，除非指定，所有 Block 默认是非边界Block，仅在边界 Block 中添加边界判断和处理，通过减少不必要的分支判断，可以有效的提升 API 性能。此外为提升 IO 类 API 的访存效率，ReadData、WriteData 等 API 提供了向量化数据读取参数 VecSize，保证 AI 加速器的每个 Core 能够一次读取多个数据，以保证访存类 API 具有较高的访存效率，配合 Block 边界模板参数可进一步提升 Kernel 性能。
在进行计算类 API 优化时，为充分发挥编译器优化效果，在数据计算展开时，增加了 Program unroll 操作，数据循环操作能够在编译阶段实现循环展开，降低判断操作产生的时间开销。使用 KP 实现的 elementwise 类、activation 类、reduce 类等算子在适配前后性能分别有5%~15%的提升。
KP 加速硬件适配实战
昆仑芯2代产品适配
飞桨框架 v2.3 版本正式适配了昆仑芯2代 AI 芯片。由于在算子对接中充分使用了 KP 进行开发，相比传统方式减少了90%以上的代码开发量，并且在昆仑芯2代 AI 芯片上的算子性能优化工作也大大减轻。下面我们就来具体看看 KP 是如何大幅减少硬件适配工作量的。
▌降低硬件感知度，减少硬件适配开发成本
KP 将硬件处理细节封装在 API 内部，各 AI 加速器的 KP 接口完全一致，使得算子开发者在进行 Kernel 开发时无需区分硬件平台，从而实现不同硬件平台复用同一套 Kernel 代码的编程效果。该种方式不仅能降低硬件感知度，保证开发者在进行新硬件支持和适配时无需学习复杂的硬件特性，同时能够实现一套 Kernel 代码实现多个硬件平台的算子支持，从而快速完成新硬件支持。
以昆仑芯2代 AI 芯片 CR20 为例，如果采用传统算子开发方式接入飞桨框架，需根据新硬件特性在飞桨中添加相应的 host 端代码，完成 elementwise, reduce, activation 三类约70个算子的功能支持约需添加约10,800行代码；而通过 KP 完成这三类算子的功能支持仅需添加707行代码，代码适配量可减少93.4%。
▌一处优化多处受益
当前飞桨中使用 KP 支持的算子约70个，完成该70个算子 Kernel 实现仅使用到6个 IO 类 API， 7个计算类 API，13个通用 Functor。不同算子间存在大量共用的 API，因此通过对少量KP的性能优化，可以直接提升这70个算子的性能，实现一处优化多处受益的效果。
该方式将原始算子性能优化转移到低层 KP 性能优化，能够在收获性能收益的同时，减少性能优化的工程量。
例如 elementwise 类与 activation 类算子均使用到了 ReadData、WriteData，通过对这两个 API 进行性能优化，其优化效果可以直接体现在多个 elementwise 类与 activation 类算子上，如下图：
图9 API 与 OP 一对多展示图
为充分展示特定平台下的 KP 内部性能优化效果及优化细节，此处以 CR20 为例进行阐述。CR20 包含8个 CLUSTER，CLUSTER 是细粒度可编程处理器，用于加速深度学习中非计算密集型的算子。每个 CLUSTER 包含64个 XPU Core 和 256KB 的S hared Memory，每个 XPU Core 有 8KB 的 local memory，支持 scalar 运算并且具备 512bit 的 SIMD 计算能力。CLUSTER 具有非常好的通用性和可编程性，用户可以根据需求来灵活实现各种函数，昆仑芯产品上大部分的 AI 操作都是由这个计算单元来实现，类似于 NVIDIA GPU 的 CUDA Core。
昆仑芯2代芯片的编程是采用SIMT的并行计算模型，将计算任务下发到多个 CLUSTER 和 Core 上进行并行计算。为提升数据访存效率，结合 XPU 编程模型和 KP 访存类 API 功能需求，我们对输入数据进行分块处理，使得输入数据尽可能的均分到每个 Core 上去处理，从而最大限度保证资源均衡。以 broadcast_add 为例：
输入规模： {x_input[1]，y_input[65536]} 设置向量化数据读取参数： VecSize=input_len/(CLUSTER_num*core_num)=128
每个 Core 每次可以通过 GM2LM 读取128个数据，这样 XPU 可以做到对数据进行批量读取，减少访存耗时。
另外 XPU CLUSTER 包含丰富的 SIMD 指令集，XPU 编译器将这些指令封装成了内建函数供开发者调用，使编程方式更加灵活并且可以提高计算效率。以加法操作 AddFunctor 为例，相比优化前的“a+b”操作，使用 SIMD 指令一次可以并行完成16个数据的计算，对计算效率有很大提升。
基于 XPU 的硬件特性和编程模型，通过对 XPU Core 进行合理的任务划分，以及采用 SIMD 指令等方式来对 KP 进行优化后，调用该 API 的算子性能都同步得到优化，实现了一处优化多处收益的优化效果。通过实际测试，优化后算子和模型性能相比原 KP 都有大幅提高，且能接近甚至达到 XDNN（XPU Deep Neural Network Library，昆仑芯深度神经网络高性能优化加速算子库）的性能。
总结
深度学习框架在高性能算子的具体实现上，目前主流的方式是针对不同硬件单独开发及优化，同一个算子提供多个硬件的实现版本，在编译时或运行时根据实际运行设备选择执行。这会使得各个算子间的代码无法系统性复用，并且类似的性能优化方法需要在不同硬件上重复实现，导致算子开发成本高，性能优化难度大，硬件迁移成本高等问题。
飞桨 KP 提供了一套高度可复用的算子内核开发接口，屏蔽不同硬件底层架构细节，起到性能优化统一进行等作用，极大减轻框架与硬件对接的成本。
在飞桨框架 v2.3 版本适配昆仑芯2代 AI 芯片的工作中，KP 体系的优越性得到了充分验证。不仅适配的代码量相比传统方式减少了90%以上，而且通过针对7个 KP 的优化，57个算子性能均达到了持平昆仑芯 XDNN 的表现，大大减轻了性能优化成本。