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

在審

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

1.一種矩陣乘分片配置選擇方法，其特征在于，包括以下步驟：

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

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

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

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

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

2.如權利要求1所述的一種矩陣乘分片配置選擇方法，其特征在于，所述步驟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 shape size表示矩陣乘加速單元TCU支持的M方向和N方向的最小線程個數。

3.如權利要求2所述的一種矩陣乘分片配置選擇方法，其特征在于，所述步驟2中，根據數據復用性能對各分片配置方案進行初步篩選的方法為：

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

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

4.如權利要求2所述的一種矩陣乘分片配置選擇方法，其特征在于，所述步驟3具體包括以下步驟：

S1、分別獲取各分片配置方案將硬件計算資源占滿所需的最少線程塊數min_theory_block_num；

S2、對于m*k的左矩陣與k*n的右矩陣間的矩陣乘，輸入矩陣乘的尺寸(m，k，n)并計算各分片配置方案下的線程塊數Thread_Blocks，

其中，m_tiles和n_tiles分別表示M方向和N方向的數據復用次數；

S3、若Thread_Blocks不小于min_theory_block_num，則保留該分片配置方案；若所有分片配置方案的Thread_Blocks均小于min_theory_block_num，則說明輸入矩陣乘的尺寸太小，無法充分利用起計算資源，將所有分片配置方案中的數據復用次數設置為1后保留。

5.如權利要求4所述的一種矩陣乘分片配置選擇方法，其特征在于，步驟4中所述計算訪存比TLP的計算公式為：

其中，m*k的左矩陣的分片形狀是(block_m_threads*m_tiles,block_k_threads)，k*n的右矩陣的分片形狀是(block_k_threads,block_n_threads*m_tiles)，block_k_threads表示K方向的分片長度。

6.如權利要求4所述的一種矩陣乘分片配置選擇方法，其特征在于，步驟5中所述核函數發射次數Launch_kernel_times的計算公式為：

其中，block_per_cu表示每個計算單元上可以同時并行存在的線程塊數，且

max available register num of a thread表示，由于芯片硬件限制，每個線程可以使用的最大的寄存器數；num of registers actuaily used表示每個線程實際使用的寄存器數，max shared memory size of a SM表示每個流處理器簇可以使用的最大的共享內存的大小，shared memory size used by per thread block表示每個線程塊實際使用的共享內存的大小。

7.一種矩陣乘分片配置選擇系統，其特征在于，包括輸入模塊、編譯模塊和權利要求1至6任一權利要求所述矩陣乘分片配置選擇方法的實現模塊，輸入模塊用于輸入矩陣乘的尺寸，編譯模塊用于將各分片配置方案的配置參數作為矩陣乘kernel函數的模板參數并編譯生成所有分片配置方案的kernel函數；運行時，輸入矩陣乘的尺寸，通過矩陣乘分片配置選擇方法的實現模塊選擇出最佳的分片配置方案及對應的kernel函數，以執行矩陣乘運算。