一种深度学习框架适配AI加速芯片的方法和装置与流程

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


背景技术：

2.深度学习框架是人工智能领域的操作系统，通过张量、基于张量的操作（op）、计算图、自动微分工具、硬件拓展包（如cublas, cudnn）等五大核心组件帮助用户方便地实现各种深度学习算法，充分释放底层硬件的算力资源。
3.ai加速芯片也被称为ai加速器或计算卡，即专门用于处理人工智能应用中大量的计算任务的硬件。不同于传统芯片，ai芯片规模更大、结构更复杂、运算能力更强，为算力提供了有力的支持。
4.当前ai加速芯片种类日渐增多，呈现一种百花齐放的状态。深度学习框架底层兼容更多种类的ai加速器可以增加框架的兼容性，针对不同应用场景选择最合适的硬件，充分释放硬件算力资源。然而由于每一种ai加速芯片硬件结构不同，操作方式不同，导致在深度学习框架中每支持一种硬件就要从头开始把整个流程进行一遍，进行大量的重复性工作。


技术实现要素：

5.为了解决现有技术中存在的上述技术问题，本发明提出了一种深度学习框架适配ai加速芯片的方法和装置，通过芯片类型定义、芯片类型注册、芯片内存支持三个主要的步骤，来简化深度学习框架适配ai加速芯片的工作，其具体技术方案如下：一种深度学习框架适配ai加速芯片的方法，主要分为三个阶段：芯片类型定义阶段，基于某种数据传输格式，如protobuf，以自定义或编写的文件定义要支持的ai加速芯片类型，所述ai加速芯片类型包括枚举类型，用于在深度学习框架内部区分不同种类的芯片，使得深度学习框架根据不同的枚举类型值进行相应的处理，深度学习框架中的上下文管理器、设备线程、流索引生成器、计算核kernel等基础设施是与芯片类型强绑定的，因为不同的芯片操作方式不同，故以上几种基础设施的实现方式也是不同的，以计算核kernel的实现为例，在cpu上可能会用到openblas库，在gpu上会用到cublas库，在寒武纪mlu上会用到cnrt和cnnl库，经过芯片类型注册阶段，深度学习框架就可以根据芯片类型自动选择相应的方式实现计算核kernel；芯片类型注册阶段，将ai加速芯片相关的芯片类型、上下文管理器、设备线程以及流索引生成器、计算核kernel注册到各自的哈希表中，并采用基于单例模式的注册机制使得芯片类型到上下文管理器、设备线程、流索引生成器、计算核kernel的一一映射，便于框架在需要的时候可以根据芯片类型方便地找到相应的内容，相应的内容包括相应的芯片类型对应的上下文管理器、设备线程以及流索引生成器和计算核kernel；芯片内存支持阶段，把ai加速芯片内存相关的操作放到深度学习框架中，使得框
架可以对芯片的存储空间进行统一管理。
6.优选的，所述芯片类型定义阶段，在芯片类型定义相关的数据结构中增加该ai加速芯片的类型。
7.优选的，所述哈希表的 key 值是要注册的芯片类型，value 值是该芯片所对应的处理函数，完成各种句柄的创建、片上计算、内存管理的操作；所述设备上下文管理器里面是可以产生芯片操作句柄的方法，其注册过程中将设备上下文管理器注册到哈希表中与芯片类型一一对应，具体为：创建上下文句柄，用于生成芯片计算过程中的各种句柄，再通过调用上下文句柄来创建设备上下文管理器用于向外部调用者提供各种句柄，并进行设备同步的操作，其中句柄包括：流句柄、芯片操作句柄；所述流索引生成器用于对不同的操作生成相应的流索引号，该类型的注册也是创建一个哈希表，key值是芯片类型，value值是相应的流索引生成器。
8.优选的，所述设备线程注册，该过程创建芯片类型相关的设备线程，用于创建启动芯片上计算过程的线程，设备线程创建完成之后将其注册到自己的哈希表中完成线程与芯片类型的一一对应。
9.优选的，所述计算核kernel 注册，首先实现计算核kernel内部的计算逻辑，然后将芯片类型和数据类型两项组成的二元组作为key值，计算核kernel作为value值注册到计算核kernel相关的哈希表中去。
10.优选的，所述芯片内存支持阶段，基于不同种类的数据传输格式，定义需要支持的ai加速芯片的内存类型，对片上内存进行分配、释放、空间切分、数据拷贝操作。
11.优选的，所述定义需要支持的ai加速芯片的内存类型，用以在框架内部区分不同芯片的内存类型，在内存类型相关的数据结构中添加该ai加速芯片的内存类型。
12.优选的，所述对芯片的存储空间进行统一管理包括内存分配，所述内存分配包括：将所述芯片的存储空间统一分配出第一存储空间，根据不同模块例如卷积、池化等，分别对应的所需空间大小将所述第一存储空间划分成空间大小不同的段，划分的方法采用初始地址加偏移量的方式或调用芯片提供的特定api接口的方式；内存释放用于内存统一释放。
13.优选的，所述对芯片的存储空间进行统一管理包括对片上内存进行数据拷贝，ai加速芯片在开始计算之前及完成计算之后均与主机之间进行数据的拷贝，具体为：在使用ai加速芯片的过程中先从主机内存将数据拷贝到片上内存，完成计算之后再将结果拷回主机内存，计算过程中对主机与芯片之间、芯片与芯片之间的数据进行拷贝，在内存拷贝时根据数据拷贝的源和目的地的不同分情况处理一种深度学习框架适配ai加速芯片的装置，包括一个或多个处理器，用于实现所述的深度学习框架适配ai加速芯片的方法。
14.本发明的有益效果：本发明简化了深度学习框架适配ai加速芯片的工作。
附图说明
15.图1所示的是根据本发明的深度学习框架适配ai加速芯片整体流程示意图；图2所示的是根据本发明的深度学习框架适配ai加速芯片设备管理器和设备线程注册示意图；
图3所示的是根据本发明的深度学习框架适配ai加速芯片注册步骤中需要注册的基础设施示意图；图4所示的是根据本发明的深度学习框架适配ai加速芯片内空间的分配与释放流程图；图5是本发明提供的一种深度学习框架适配ai加速芯片的装置的结构框图。
具体实施方式
16.为了使本发明的目的、技术方案和技术效果更加清楚明白，以下结合说明书附图和实施例，对本发明作进一步详细说明。
17.如图1所示，本发明的一种深度学习框架适配ai加速芯片的方法，主要分为三个阶段：芯片类型定义、芯片类型注册、芯片内存支持。所述芯片类型定义就是把需要支持的芯片的类型以一个枚举值的形式写入到proto文件的相关数据结构中去，从而使得框架内部可以正确识别该芯片类型，此处以protobuf作为序列化工具来说明整个过程，同类型的工具还有json、hessian等等，不过目前主流深度学习框架都是用的protobuf；；芯片类型注册是把该芯片所需的基础设施注册到哈希表中去，便于框架在需要的时候可以根据芯片类型方便地找到相应的内容；芯片内存支持是把该芯片内存相关的操作放到框架中去，使得框架可以对芯片的存储空间进行统一管理。
18.所述芯片类型定义阶段，基于某种数据传输格式，如protobuf，以自定义或编写的文件定义要支持的ai加速芯片类型，所述ai加速芯片类型包括枚举类型，用于在深度学习框架内部区分不同种类的芯片，使得深度学习框架根据不同的枚举类型值进行相应的处理，深度学习框架中的上下文管理器、设备线程、流索引生成器、计算核kernel等基础设施是与芯片类型强绑定的，因为不同的芯片操作方式不同，故以上几种基础设施的实现方式也是不同的，以计算核kernel的实现为例，在cpu上可能会用到openblas库，在gpu上会用到cublas库，在寒武纪mlu上会用到cnrt和cnnl库，经过芯片类型注册阶段，深度学习框架就可以根据芯片类型自动选择相应的方式实现计算核kernel。
19.如图3所示，芯片注册，包括芯片类型、上下文管理器、设备线程、流索引生成器、计算核kernel等几个基础设施的注册，利用基于单例模式的注册机制将芯片类型与该芯片对应的设备上下文管理器、设备线程、流索引生成器、计算核kernel注册到各自的哈希表中，使得芯片类型与该芯片对应的设备上下文管理器、设备线程、流索引生成器、计算核kernel一一映射。其中，哈希表的 key 值是要注册的芯片类型，本质上是枚举类型，value 值是该芯片所对应的处理函数，完成各种句柄的创建、片上计算、内存管理等多种操作，ai加速芯片的操作常常需要用到句柄，句柄本质上是一个指针，指向完成该项操作所需要的的资源，比如gpu上的多流会用到流句柄cudastream_t，cublas库要用到cublashandle_t，cudnn 库会用到 cudnnhandle_t；而在燧原dtu芯片上则会用到topscontext_t句柄等等。哈希表的value值是类类型，包含各种成员函数，其中句柄的创建是由成员函数实现的；注册完成后可以利用芯片类型查询哈希表从而方便找到相应的处理函数。具体的，这里的注册过程会用到四个哈希表，四个哈希表的key值都是芯片类型，如cpu、gpu、mlu等，哈希表的value值是以上几种基础设施，第一个表的value值是上下文管理器，第二个表的value值是设备线程，第三个表的value值是流索引生成器，第四个表的value值是计算核kernel，如下表1-4
所示；keyvaluecpucpu的上下文管理器gpugpu的上下文管理器mlumlu的上下文管理器
……
表1上下文管理器注册表keyvaluecpucpu的设备线程gpugpu的设备线程mlumlu的设备线程
……
表2设备线程注册表keyvaluecpucpu的流索引生成器gpugpu的流索引生成器mlumlu的流索引生成器
……
表3流索引生成器注册表keyvaluecpucpu上的计算核kernelgpugpu上的计算核kernelmlumlu上的计算核kernel
……
表4计算核kernel注册表如图2所示，展示了设备上下文管理器和设备线程两个主要基础设施的注册过程，设备上下文管理器里面是可以产生芯片操作句柄的方法。注册的过程就是把芯片类型作为key值，设备上下文管理器作为value值的一对key-value对放进一个哈希表里，后期使用的时候通过芯片类型和哈希表方便地查询到对应的方法，从而获取相应的句柄。设备线程里包含了用于创建该芯片类型相关线程的方法，其注册过程与设备上下文管理器的注册如出一辙，只是它需要的是另一个哈希表。
20.所述芯片内存支持步骤，训练过程中对片上内存进行统一管理，即统一分配和释放，统一分配出的大段存储空间经切割成小段分别用于不同的用途，ai加速芯片在开始计算之前及完成计算之后都需要与主机之间进行数据的拷贝，所以需要实现针对当前芯片的内存空间分配、释放、切分、数据拷贝等方法。
21.如图4所示的是片上内存分配、空间切分及释放的流程图，程序运行过程中会对存储空间进行统一分配和管理。具体的，程序首先会判断所需空间是否在该芯片上，如果不在则不进行处理，如果所需空间在该芯片上则按预先计算的结果分配所需的空间，通常是一段比较大的空间，然后根据不同部分的需求对该空间进行切分，切分的方式有两种：1）基地
址加偏移的方式；2）调用芯片特定api的方式，具体切分方式需要根据芯片的具体情况而定。存储空间分配好之后就可以进行后续的训练任务，训练结束之后程序退出前需要释放存储空间，程序首先会判断需要释放的空间是否在本芯片上，若是则释放该空间，否则直接退出。
22.与前述深度学习框架适配ai加速芯片的方法的实施例相对应，本发明还提供了深度学习框架适配ai加速芯片的装置的实施例。
23.参见图5，本发明实施例提供的一种深度学习框架适配ai加速芯片的装置，包括一个或多个处理器，用于实现上述实施例中的深度学习框架适配ai加速芯片的方法。
24.本发明深度学习框架适配ai加速芯片的装置的实施例可以应用在任意具备数据处理能力的设备上，该任意具备数据处理能力的设备可以为诸如计算机等设备或装置。装置实施例可以通过软件实现，也可以通过硬件或者软硬件结合的方式实现。以软件实现为例，作为一个逻辑意义上的装置，是通过其所在任意具备数据处理能力的设备的处理器将非易失性存储器中对应的计算机程序指令读取到内存中运行形成的。从硬件层面而言，如图5所示，为本发明深度学习框架适配ai加速芯片的装置所在任意具备数据处理能力的设备的一种硬件结构图，除了图5所示的处理器、内存、网络接口、以及非易失性存储器之外，实施例中装置所在的任意具备数据处理能力的设备通常根据该任意具备数据处理能力的设备的实际功能，还可以包括其他硬件，对此不再赘述。
25.上述装置中各个单元的功能和作用的实现过程具体详见上述方法中对应步骤的实现过程，在此不再赘述。
26.对于装置实施例而言，由于其基本对应于方法实施例，所以相关之处参见方法实施例的部分说明即可。以上所描述的装置实施例仅仅是示意性的，其中所述作为分离部件说明的单元可以是或者也可以不是物理上分开的，作为单元显示的部件可以是或者也可以不是物理单元，即可以位于一个地方，或者也可以分布到多个网络单元上。可以根据实际的需要选择其中的部分或者全部模块来实现本发明方案的目的。本领域普通技术人员在不付出创造性劳动的情况下，即可以理解并实施。
27.本发明实施例还提供一种计算机可读存储介质，其上存储有程序，该程序被处理器执行时，实现上述实施例中的深度学习框架适配ai加速芯片的方法。
28.所述计算机可读存储介质可以是前述任一实施例所述的任意具备数据处理能力的设备的内部存储单元，例如硬盘或内存。所述计算机可读存储介质也可以是风力发电机的外部存储设备，例如所述设备上配备的插接式硬盘、智能存储卡（smart media card，smc）、sd卡、闪存卡（flash card）等。进一步的，所述计算机可读存储介质还可以既包括任意具备数据处理能力的设备的内部存储单元也包括外部存储设备。所述计算机可读存储介质用于存储所述计算机程序以及所述任意具备数据处理能力的设备所需的其他程序和数据，还可以用于暂时地存储已经输出或者将要输出的数据。
29.以上结合具体实施例描述了本公开的基本原理，但是，需要指出的是，对本领域的普通技术人员而言，能够理解本公开的方法和装置的全部或者任何步骤或者部件，可以在任何计算装置，包括处理器、存储介质等，或者计算装置的网络中，以硬件、固件、软件或者它们的组合加以实现，这是本领域普通技术人员在阅读了本公开的说明的情况下运用他们的基本编程技能就能实现的。
30.因此，本公开的目的还可以通过在任何计算装置上运行一个程序或者一组程序来实现。所述计算装置可以是公知的通用装置。因此，本公开的目的也可以仅仅通过提供包含实现所述方法或者装置的程序代码的程序产品来实现。也就是说，这样的程序产品也构成本公开，并且存储有这样的程序产品的存储介质也构成本公开。显然，所述存储介质可以是任何公知的存储介质或者将来所开发出来的任何存储介质。
31.还需要指出的是，在本公开的装置和方法中，显然，各部件或各步骤是可以分解和/或重新组合的。这些分解和/或重新组合应视为本公开的等效方案。并且，执行上述系列处理的步骤可以自然地按照说明的顺序按时间顺序执行，但是并不需要一定按照时间顺序执行。某些步骤可以并行或彼此独立地执行。
32.上述具体实施方式，并不构成对本公开保护范围的限制。本领域技术人员应该明白的是，取决于设计要求和其他因素，可以发生各种各样的修改、组合、子组合和替代。任何在本公开的精神和原则之内所作的修改、等同替换和改进等，均应包含在本公开保护范围之内。