[發明專利]一種矩陣乘分片配置選擇方法及系統

在審

• 申請號: 202111506483.9
• 申請日: 2021-12-10
• 公開（公告）號: CN114238853A
• 公開（公告）日: 2022-03-25
• 發明（設計）人: 方寶輝;石加圣;單天逸
• 申請（專利權）人: 上海天數智芯半導體有限公司
• 主分類號: G06F17/16
• 分類號: G06F17/16;G06F17/15
• 代理公司: 南京鐘山專利代理有限公司 32252
• 代理人: 王磊

本發明涉及一種矩陣乘分片配置選擇方法及系統，首先，根據硬件特性羅列出線程塊所有可能的分片配置方案；然后，將各分片配置方案中的線程塊分成M方向和N方向的多個線程，并根據數據復用性能對各分片配置方案進行初步篩選；接著，從保留方案中選擇出能夠充分利用硬件計算資源的分片配置方案；然后，進一步選擇出具有最大計算訪存比的分片配置方案；最后，找出具有最小核函數發射次數的分片配置方案作為最佳分片配置方案。相比于現有技術，本發明根據芯片的硬件資源使用情況和輸入矩陣的尺寸決定分片方案，平臺可移植性強；且能夠實時動態的選擇分片方案，對調用矩陣乘算法的上層用戶更友好、便捷。

技術領域

本發明屬于高性能計算技術領域，具體涉及一種矩陣乘分片配置選擇方法及系統。

背景技術

矩陣乘法運算應用非常廣泛，如在深度學習、科學計算、圖像處理中扮演重要角色。矩陣乘法的性能受計算規模、計算用例的尺寸、硬件計算資源和數據復用等諸多因素影響，因此也有許多的參數配置方案。

現有的矩陣乘分片解決方案有些在方案設計的時候使用了經驗值，但由于軟硬件平臺的變化(如編譯器升級、更換新一代芯片)，該方法會出現明顯的失效或者分片不準確的情況。此外，對于卷積推理和訓練中的變體矩陣乘，由于使用的共享內存和寄存器等硬件資源明顯不同，經驗閾值也會明顯不同，因此這樣的方案不便于推廣。還有些矩陣乘分片方案是通過性能測試將每種分片方案的時間測試出來，再由人工選擇最優的配置方案，然而實際應用中，矩陣乘的尺寸會不斷變化，通過性能測試人工獲得最佳配置的方法嚴重浪費了時間和人力。

發明內容

本發明針對現有技術中的不足，提供一種矩陣乘分片配置選擇方法及系統，采用以下技術方案：

一種矩陣乘分片配置選擇方法，包括以下步驟：

步驟1：根據硬件特性，羅列出線程塊所有可能的分片配置方案；

步驟2：將各分片配置方案中的線程塊分成M方向和N方向的多個線程，并根據數據復用性能對各分片配置方案進行初步篩選；

步驟3：從初步篩選保留下來的分片配置方案中，選擇出能夠充分利用硬件計算資源的分片配置方案并保留；

步驟4：計算步驟3中保留的各分片配置方案的計算訪存比，選擇出具有最大計算訪存比的分片配置方案并保留；

步驟5：計算步驟4中保留的各分片配置方案的核函數發射次數，找出具有最小核函數發射次數的分片配置方案作為最佳分片配置方案。

進一步地，所述步驟2中，若采用通用乘加指令做矩陣乘運算，則M方向和N方向的線程個數滿足以下條件：

16≤(block_m_threads,block_n_threads)≤max(block_m_threads)/16；

若采用矩陣乘硬件加速單元做矩陣乘運算，則M方向和N方向的線程個數滿足以下條件：

其中，block_m_threads和block_n_threads分別表示M方向的線程個數和N方向的線程個數，TCUmin shapesize表示矩陣乘加速單元TCU支持的M方向和N方向的最小線程個數。

進一步地，所述步驟2中，根據數據復用性能對各分片配置方案進行初步篩選的方法為：

通過寄存器使用量將活躍線程數量少于設定閾值的分片配置方案排除；

或先羅列出分片配置方案的數據復用方案，然后通過多種尺寸的矩陣乘測試用例進行性能測試，排除在所有測試中性能表現低于設定閾值的數據復用方案。

進一步地，所述步驟3具體包括以下步驟：