[发明专利]一种深度学习框架适配AI加速芯片的方法和装置

说明书

本发明的一种深度学习框架适配AI加速芯片的方法和装置，具体分为三个阶段：芯片类型定义、芯片类型注册、芯片内存支持，所述芯片类型定义就是把需要支持的芯片的类型以一个枚举值的形式定义到proto文件中去，从而使得框架内部可以正确识别该芯片类型；芯片类型注册是把该芯片所需的基础设施注册到哈希表中去，便于框架在需要的时候可以根据芯片类型方便地找到相应的内容；芯片内存支持是把该芯片内存相关的操作放到框架中去，使得框架可以对芯片的存储空间进行统一管理。本发明简化了深度学习框架适配AI加速芯片的工作。

技术领域

本发明属于深度学习基础软件领域，涉及一种深度学习框架适配AI加速芯片的方法和装置。

背景技术

深度学习框架是人工智能领域的操作系统，通过张量、基于张量的操作（Op）、计算图、自动微分工具、硬件拓展包（如cublas, cudnn）等五大核心组件帮助用户方便地实现各种深度学习算法，充分释放底层硬件的算力资源。

AI加速芯片也被称为AI加速器或计算卡，即专门用于处理人工智能应用中大量的计算任务的硬件。不同于传统芯片，AI芯片规模更大、结构更复杂、运算能力更强，为算力提供了有力的支持。

当前AI加速芯片种类日渐增多，呈现一种百花齐放的状态。深度学习框架底层兼容更多种类的AI加速器可以增加框架的兼容性，针对不同应用场景选择最合适的硬件，充分释放硬件算力资源。然而由于每一种AI加速芯片硬件结构不同，操作方式不同，导致在深度学习框架中每支持一种硬件就要从头开始把整个流程进行一遍，进行大量的重复性工作。

发明内容

为了解决现有技术中存在的上述技术问题，本发明提出了一种深度学习框架适配AI加速芯片的方法和装置，通过芯片类型定义、芯片类型注册、芯片内存支持三个主要的步骤，来简化深度学习框架适配AI加速芯片的工作，其具体技术方案如下：

一种深度学习框架适配AI加速芯片的方法，主要分为三个阶段：

芯片类型定义阶段，基于某种数据传输格式，如Protobuf，以自定义或编写的文件定义要支持的AI加速芯片类型，所述AI加速芯片类型包括枚举类型，用于在深度学习框架内部区分不同种类的芯片，使得深度学习框架根据不同的枚举类型值进行相应的处理，深度学习框架中的上下文管理器、设备线程、流索引生成器、计算核Kernel等基础设施是与芯片类型强绑定的，因为不同的芯片操作方式不同，故以上几种基础设施的实现方式也是不同的，以计算核Kernel的实现为例，在CPU上可能会用到OpenBLAS库，在GPU上会用到cuBLAS库，在寒武纪MLU上会用到cnrt和cnnl库，经过芯片类型注册阶段，深度学习框架就可以根据芯片类型自动选择相应的方式实现计算核Kernel；

芯片类型注册阶段，将AI加速芯片相关的芯片类型、上下文管理器、设备线程以及流索引生成器、计算核Kernel注册到各自的哈希表中，并采用基于单例模式的注册机制使得芯片类型到上下文管理器、设备线程、流索引生成器、计算核Kernel的一一映射，便于框架在需要的时候可以根据芯片类型方便地找到相应的内容，相应的内容包括相应的芯片类型对应的上下文管理器、设备线程以及流索引生成器和计算核Kernel；

芯片内存支持阶段，把AI加速芯片内存相关的操作放到深度学习框架中，使得框架可以对芯片的存储空间进行统一管理。

优选的，所述芯片类型定义阶段，在芯片类型定义相关的数据结构中增加该AI加速芯片的类型。

优选的，所述哈希表的 key 值是要注册的芯片类型，value 值是该芯片所对应的处理函数，完成各种句柄的创建、片上计算、内存管理的操作；