其他
MindSpore自定义算子:思考,挑战与实践
点击蓝字MindSpore关注我们引言陈天奇在《新一代深度学习编译技术变革和展望》一文中曾经提到,当前阻碍深度学习的全栈解决方案发展的主要是两个隔阂,包括:竖向隔阂,也就是手工优化的方案和自动编译优化的方案的隔阂。当前的大部分深度学习框架面临以手工算子优化为主的算子库和以自动优化为主的编译方案二选一的困境,如何让手工优化,机器学习优化专家的知识和自动优化做有机整合,是目前行业面临一个大的问题。横向隔阂,也就是图算软件分层引起的跨层隔阂。当前的大部分深度学习框架对不同层次的抽象分开设计导致,例如图层和算子层的分离,但是深度学习编译和优化本身不是一个一个层级可以全部完成优化的问题。解决相关问题需要各个层级抽象之间的联动。在网络开发实践中,我们也深深的感受到了这两种隔阂对计算表达的严重限制。一方面,我们注意到竖向隔阂不仅仅存在于手工优化的方案和自动编译优化的方案之间,同样存在于不同的手工优化的方案之间。深度学习框架在整合这些优化方案加速网络时面临比较大的阻碍:有的手工优化的算子以第三方算子库的形式呈现,全部加入框架会使框架过于厚重;有的手动优化方案只能针对某些特定场景有效,无法作为一个普遍的优化应用在全部场景。特别是在AI-HPC融合场景中,HPC的应用中的算子很多是针对特殊硬件做的专门优化,如何有机的整合这类算子,是深度学习框架面对AI-HPC融合场景的一个重要挑战。另一方面,在业界的大部分深度学习框架中,算子的表达与注册独立于模型(也就是图层),图层视算子为黑盒。当用户在框架之外独立添加算子时,图层无法了解算子内部的具体计算逻辑,因此无法进行例如算子融合和算子拆分等图层上的优化。如何在图层上直接定义算子,让图层感知算子的具体逻辑,是打破图层和算子层的横向隔阂,实现各个层级抽象之间的联动的重要议题。总的来说,一个深度学习的全栈解决方案总是会遇到不同算子优化方案的竖向隔阂和图算抽象分离的横向隔阂。为了解决如上隔阂引入的表达问题,MindSpore给出了自己的解决方案,即统一自定义算子表达。1新接口:MindSpore自定义算子统一接口Custom随着MindSpone在科学计算等新型网络场景的大量使用,对算子的灵活性表达提出更高的要求,针对传统深度学习网络设计的算子库越来越无法满足需求。因此,MindSpore自1.6起推出自定义算子统一接口Custom,并在1.8版本迎来全新升级,有机的结合了手动优化和自动编译算子,并让图层感知算子的定义,助力用户方便高效的添加自定义算子,可以满足包括快速验证,实时编译和第三方算子接入等不同场景下的用户需求。当前自定义算子支持的算子开发方式包括:ms_kernel、tbe、aicpu、aot、pyfunc、julia。不同的算子开发方式支持不同的场景和不同的平台。自定义算子模式和平台支持在设计自定义算子的多种模式时,我们充分考虑到了阻碍深度学习的全栈解决方案发展的两个隔阂,并用统一的接口去实现框架对算子的多种需求。1.1手工优化算子的一键接入在前面的讨论中我们曾经提到,深度学习全栈解决方案发展的纵向隔阂是手工优化的方案和自动编译优化的方案的隔阂。为了方便用户快捷的接入手工优化的算子,MindSpore自定义算子统一接口Custom提供aot模式,灵活封装手工算子。用户可以手工优化算子实现,并把算子以动态库的形式接入MindSpore加速网络。特别的,当第三方库提供CPP或者CUDA函数的API时,我们可以自定义算子中调用第三方库的API,然后通过编译链接将第三方库接入MindSpore加速网络,从而实现手工优化算子的一键接入。以PyTorch的Aten库为例进行介绍,我们在网络迁移的时候会遇到基于PyTorch的网络中的部分算子尚未被MindSpore支持的情况。为了快速支持网络,我们可以利用Custom算子的
2022年8月26日