[發明專利]一種深度學習框架適配AI加速芯片的方法和裝置

有效

• 申請號: 202111497148.7
• 申請日: 2021-12-09
• 公開（公告）號: CN113918507B
• 公開（公告）日: 2022-04-08
• 發明（設計）人: 王拓;楊非;黃振華;鮑虎軍;華煒
• 申請（專利權）人: 之江實驗室
• 主分類號: G06F15/78
• 分類號: G06F15/78;G06N3/04;G06N3/063
• 代理公司: 杭州浙科專利事務所(普通合伙) 33213
• 代理人: 孫孟輝

本發明的一種深度學習框架適配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 值是該芯片所對應的處理函數，完成各種句柄的創建、片上計算、內存管理的操作；