\u200E
飞桨首创算子内核开发KP体系,快速适配硬件,高效优化性能
发布日期:2022-08-11T13:11:52.000+0000 浏览量:135次

随着深度学习的蓬勃发展,其应用领域也愈加丰富,对深度学习框架的计算算子需求也随之增多。一个成熟的框架,其算子数目通常达到几百甚至上千。面对同样繁荣发展的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,对开发者的硬件知识以及开发能力有着很高的要求。

图1 算子执行流程

在设计理念上,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实现即可。

图2 KP多平台算子执行流程
在具体实现上,KP内核开发接口是Block级别的函数封装,不会增加kernel launch和访存开销,内部逻辑使用寄存器作为缓冲,工具链编译器进行自动优化,同时不增加核函数调用次数及全局内存访问次数,因此使用内核开发接口构建算子核函数不会带来额外的性能开销。内核级API具有完备性和可复用等特点,API一次调用可完成一个Block内的数据读取和计算操作。例如reduce API通过ReduceMode标识规约模式,当ReduceMode为kGlobalMode时说明当前API对Block内的线程进行规约操作,reduceMode为kLocalMode时进行线程内规约操作,如下图3所示,kGlobalMode的返回结果为当前Block的最大值,注意此处使用到了批量处理,可一次产生多个最值,计算结束后每个线程具有相同的返回结果;kLocalMode是进行线程内的规约操作,返回当前线程内的最值,可一次返回多组最值,计算结束后每个线程的计算结果各不相同。

图3 reduce API功能展示图


使用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代码的复用,实现同类算子功能的快速支持。

图4 KP Functor定义及使用
如上图所示,当elementwise_add Kernel实现完成后,将AddFunctor替换为MulFunctor即可完成mul 算子的功能支持,Functor主要在计算类API中使用,通过函数模板的形式转入,相比于分支判断,模板传递实现的Kernel具有更高的性能优势。

简洁易维护

KP内封装了通用计算操作,将复杂的计算和数据搬运流程下沉到了简单的接口之下,开发者仅需要使用简单的API调用即可完成复杂的功能支持,使用KP实现的Kernel代码具有极度简洁可维护性高等特点以softmax为例,使用KP替换softmax原有Kernel实现,替换后的softmax与替换前相比,整体Kernel性能保持一致,而在某些数据规模下,性能还优于原始Kernel实现。使用KP实现的softmax Kernel代码量从之前的155行减少为30行,大幅降低了Kernel开发的工作量同时替换后的 softmax 代码逻辑更加简洁清晰,与softmax的计算公式高度一致,代码可读性和易维护性明显增强。

图5 KP softmax Kernel实现

 高性能实现

目前基于主流AI加速器实现的算子性能主要受访存和计算两个方面的约束,为保证Kernel性能,KP在设计之初就针对API特性进行了优化调整。对于IO类接口,KP设计了Block边界模板参数,除非指定,所有Block默认是非边界Block,仅在边界Block中添加边界判断和处理,通过减少不必要的分支判断,可以有效的提升API性能。此外为提升IO类API的访存效率,ReadData、WriteData等API提供了向量化数据读取参数VecSize,保证AI加速器的每个Core能够一次读取多个数据,以保证访存类API具有较高的访存效率,配合Block边界模板参数可进一步提升Kernel性能。

图6 KP 边界block设置及ReadData边界处理

在进行计算类API优化时,为充分发挥编译器优化效果,在数据计算展开时,增加了Program unroll操作,数据循环操作能够在编译阶段实现循环展开,降低判断操作产生的时间开销。使用KP实现的elementwise类、activation类、reduce类等算子在适配前后性能分别有5%~15%的提升。

图7 使用KP后,三类重点算子性能提升幅度


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%。

图8 KP昆仑芯2代芯片上三类重点算子适配代码量

一处优化多处受益

当前飞桨中使用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的Shared Memory,每个XPU Core有8KB的local memory,支持scalar运算并且具备512bit的SIMD计算能力。CLUSTER具有非常好的通用性和可编程性,用户可以根据需求来灵活实现各种函数,昆仑芯产品上大部分的AI操作都是由这个计算单元来实现,类似于NVIDIA GPU的CUDA Core。

图10 昆仑芯XPU-R架构图

昆仑芯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可以做到对数据进行批量读取,减少访存耗时。

图11 调整VecSize,合理划分任务示例
另外XPU CLUSTER包含丰富的SIMD指令集,XPU编译器将这些指令封装成了内建函数供开发者调用,使编程方式更加灵活并且可以提高计算效率。以加法操作AddFunctor为例,相比优化前的“a+b"操作,使用SIMD指令一次可以并行完成16个数据的计算,对计算效率有很大提升。

图12-使用SIMD指令对AddFunctor优化

基于XPU的硬件特性和编程模型,通过对XPU Core进行合理的任务划分,以及采用SIMD指令等方式来对KP进行优化后,调用该API的算子性能都同步得到优化,实现了一处优化多处收益的优化效果。通过实际测试,优化后算子和模型性能相比原KP都有大幅提高,且能接近甚至达到XDNN(XPU Deep Neural Network Library,昆仑芯深度神经网络高性能优化加速算子库)的性能。

图13 优化前后的算子性能对比

图14 优化前后的模型性能对比


总结

深度学习框架在高性能算子的具体实现上,目前主流的方式是针对不同硬件单独开发及优化,同一个算子提供多个硬件的实现版本,在编译时或运行时根据实际运行设备选择执行。这会使得各个算子间的代码无法系统性复用,并且类似的性能优化方法需要在不同硬件上重复实现,导致算子开发成本高,性能优化难度大,硬件迁移成本高等问题。

飞桨KP提供了一套高度可复用的算子内核开发接口,屏蔽不同硬件底层架构细节,起到性能优化统一进行等作用,极大减轻框架与硬件对接的成本。
在飞桨框架v2.3版本适配昆仑芯2代AI芯片的工作中,KP体系的优越性得到了充分验证。不仅适配的代码量相比传统方式减少了90%以上,而且通过针对7个KP的优化,57个算子性能均达到了持平昆仑芯XDNN的表现,大大减轻了性能优化成本。

关注【飞桨PaddlePaddle】公众号

获取更多技术内容~