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

权利要求书

1.一种深度学习框架适配AI加速芯片的方法，其特征在于，分为三个阶段：

芯片类型定义阶段，基于不同种类的数据传输格式，以自定义或编写的文件定义要支持的AI加速芯片类型，所述AI加速芯片类型包括枚举类型，用于在深度学习框架内部区分不同种类的芯片，使得深度学习框架根据不同的枚举类型值进行相应的处理；

芯片类型注册阶段，将AI加速芯片相关的芯片类型、设备上下文管理器、设备线程以及流索引生成器、计算核Kernel注册到各自的哈希表中，并采用基于单例模式的注册机制，使得芯片类型到上下文管理器、设备线程、流索引生成器、计算核Kernel一一映射；其中，所述设备上下文管理器里面是可以产生芯片操作句柄的方法，其注册过程中将设备上下文管理器注册到待注册入所述设备上下文管理器的哈希表中与芯片类型一一对应，具体为：创建上下文句柄，再通过调用上下文句柄来创建设备上下文管理器用于向外部调用者提供各种句柄并进行设备同步的操作，其中句柄包括：流句柄、芯片操作句柄；所述流索引生成器用于对不同的操作生成相应的流索引号，该类型的注册也是创建一个哈希表，key值是芯片类型，value值是相应的流索引生成器；所述设备线程，用于创建启动芯片上计算过程的线程；

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

2.如权利要求1所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述芯片类型定义阶段，在芯片类型定义相关的数据结构中增加该AI加速芯片的类型。

3.如权利要求1所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述哈希表的 key 值是要注册的芯片类型，value 值是该芯片所对应的处理函数，完成各种句柄的创建、片上计算、内存管理的操作。

4.如权利要求1所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述设备线程注册如下，创建芯片类型相关的设备线程，设备线程创建完成之后将其注册到自己的哈希表中完成线程与芯片类型的一一对应。

5.如权利要求1所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述计算核Kernel 注册，首先实现计算核Kernel内部的计算逻辑，然后将芯片类型和数据类型两项组成的二元组作为key值，计算核Kernel作为value值注册到计算核kernel相关的哈希表中去。

6.如权利要求1所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述芯片内存支持阶段，基于不同种类的数据传输格式，定义需要支持的AI加速芯片的内存类型，对片上内存进行分配、释放、空间切分、数据拷贝操作。

7.如权利要求6所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述定义需要支持的AI加速芯片的内存类型，用以在框架内部区分不同芯片的内存类型，在内存类型相关的数据结构中添加该AI加速芯片的内存类型。

8.如权利要求6所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述对芯片的存储空间进行统一管理包括内存分配，所述内存分配包括：将所述芯片的存储空间统一分配出第一存储空间，根据不同模块分别对应的所需空间大小将所述第一存储空间划分成空间大小不同的段，划分的方法采用初始地址加偏移量的方式或调用芯片提供的特定API接口的方式；内存释放用于内存统一释放。

9.如权利要求6所述的一种深度学习框架适配AI加速芯片的方法，其特征在于，所述对芯片的存储空间进行统一管理包括对片上内存进行数据拷贝，AI加速芯片在开始计算之前及完成计算之后均与主机之间进行数据的拷贝，具体为：在使用AI加速芯片的过程中先从主机内存将数据拷贝到片上内存，完成计算之后再将结果拷回主机内存，计算过程中对主机与芯片之间、芯片与芯片之间的数据进行拷贝，在内存拷贝时根据数据拷贝的源和目的地的不同分情况处理。

10.一种深度学习框架适配AI加速芯片的装置，其特征在于，包括一个或多个处理器，用于实现权利要求1-9中任一项所述的深度学习框架适配AI加速芯片的方法。