科研動(dòng)態(tài)
AmgT：Tensor Core加速的AMG解法器
中文題目：AmgT：Tensor Core加速的AMG解法器
論文題目：AmgT: Algebraic Multigrid Solver on Tensor Cores
錄用期刊/會(huì)議：37th ACM/IEEE International Conference for High Performance Computing, Networking, Storage, and Analysis. (CCF A) Best Paper Finalist. 最佳論文提名
原文DOI：10.1109/SC41406.2024.00058
會(huì)議時(shí)間：2024.11.17-22
作者列表：
1） 盧玥辰 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 先進(jìn)科學(xué)與工程計(jì)算 博22
2） 曾禮杰 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)科學(xué)與技術(shù) 碩24
3） 王騰程 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)科學(xué)與技術(shù) 碩21
4） 付 旭 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)科學(xué)與技術(shù) 碩21
5） 李文瑄 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)科學(xué)與技術(shù) 碩24
6） 程賀琳 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)科學(xué)與技術(shù) 碩24
7） 楊德闖 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)技術(shù) 碩21
8） 金 洲 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)系教師
9） Marc Casas 巴塞羅那超級(jí)計(jì)算中心 首席研究員
10） 劉偉峰 中國(guó)石油大學(xué)（北京）人工智能學(xué)院 計(jì)算機(jī)系教師
摘要:
代數(shù)多重網(wǎng)格（AMG）方法因其靈活性和適應(yīng)性，被廣泛用于求解稀疏線(xiàn)性系統(tǒng)。盡管現(xiàn)代GPU為AMG提供了大規(guī)模并行計(jì)算能力，但其最新硬件特性（即Tensor Core及低精度計(jì)算能力）尚未被很好地用于加速AMG。本文提出了一個(gè)新的AMG解法器AmgT，其能夠在A(yíng)MG算法的多個(gè)階段中利用最新GPU的Tensor Core和混合精度能力。我們首先提出了一種統(tǒng)一稀疏存儲(chǔ)格式，充分利用Tensor Core及混合精度來(lái)提高AMG算法中頻繁調(diào)用的稀疏矩陣-矩陣乘法（SpGEMM）和稀疏矩陣-向量乘法（SpMV）的性能，并盡量減少AMG整個(gè)數(shù)據(jù)流中格式轉(zhuǎn)換的成本。同時(shí)，為更好地利用現(xiàn)有庫(kù)中的算法組件，AmgT的數(shù)據(jù)格式和計(jì)算kernel被集成到HYPRE庫(kù)中。實(shí)驗(yàn)結(jié)果表明，在NVIDIA A100、H100 和AMD MI210 GPU上，AmgT比原GPU版本的HYPRE平均快1.46倍、1.32倍和2.24倍（最高達(dá)2.10倍、2.06倍和3.67倍）。
背景與動(dòng)機(jī):
代數(shù)多重網(wǎng)格（AMG）方法因其靈活性和適應(yīng)性，被廣泛用于求解稀疏線(xiàn)性系統(tǒng)。盡管現(xiàn)代GPU為AMG提供了大規(guī)模并行計(jì)算能力，但其最新硬件特性（即Tensor Core及其低精度計(jì)算能力）尚未被很好的用于加速AMG。然而利用Tensor Core加速AMG面臨著三個(gè)挑戰(zhàn)：（1）存儲(chǔ)方面，如何避免為SpGEMM和SpMV生成不同格式的矩陣；（2）計(jì)算方面，如何將一般稀疏結(jié)構(gòu)與Tensor Core的嚴(yán)格稠密GEMM模式相匹配；（3）精度方面，如何將可變精度的SpGEMM和SpMV集成到AMG的完整數(shù)據(jù)流中。
設(shè)計(jì)與實(shí)現(xiàn):
我們首先設(shè)計(jì)了一種統(tǒng)一稀疏存儲(chǔ)格式mBSR（圖1），這是經(jīng)典塊稀疏行（BSR）格式的一種變體。mBSR格式將稀疏矩陣存儲(chǔ)在一組大小為4×4的稠密塊中，并使用位圖存儲(chǔ)每個(gè)塊中的非零元位置。
圖1. mBSR格式
在這種數(shù)據(jù)格式的基礎(chǔ)上，我們提出了一種新的AMG求解器AmgT，它在A(yíng)MG算法的多個(gè)階段利用了最新GPU的Tensor Core和及其低精度能力。AmgT基于新的SpGEMM（圖2）和SpMV算法，能夠根據(jù)塊的稀疏程度同時(shí)使用Tensor Core和CUDA Core加速計(jì)算。AmgT中的SpGEMM算法會(huì)先分析矩陣數(shù)據(jù)，并將所有塊行歸入八個(gè)分區(qū)，然后執(zhí)行兩步哈希操作進(jìn)行符號(hào)計(jì)算，以獲取生成矩陣中塊的位置信息，最后同時(shí)使用Tensor Core和CUDA Core進(jìn)行數(shù)值計(jì)算。AmgT中的SpMV算法采用自適應(yīng)性選擇負(fù)載均衡和計(jì)算kernel策略，并同樣實(shí)現(xiàn)了Tensor Core和CUDA Core的混合使用，以提高整體性能。
圖2. mBSR格式的SpGEMM計(jì)算過(guò)程
最后，多種精度的新SpGEMM和SpMV算法分別在不同層的網(wǎng)格中調(diào)用，以有效利用Tensor Core的計(jì)算能力。
圖3. AmgT的完整數(shù)據(jù)流
實(shí)驗(yàn)結(jié)果及分析:
為更好地利用現(xiàn)有庫(kù)中的算法組件，AmgT的數(shù)據(jù)格式和計(jì)算kernel被集成到HYPRE庫(kù)中。我們的測(cè)試平臺(tái)為兩款NVIDIA GPU A100（Ampere）和H100（Hopper）以及一款A(yù)MG GPU MI210（CDNA2），并測(cè)試了SuiteSparse矩陣集的16個(gè)代表性矩陣。
實(shí)驗(yàn)結(jié)果表明，在NVIDIA A100、H100 和AMD MI210 GPU上，AmgT比原GPU版本的HYPRE平均快1.46倍、1.32倍和2.24倍（最高達(dá)2.10倍、2.06倍和3.67倍），如圖4所示。在A(yíng)100和H100上，我們的混合精度AmgT比雙精度AmgT分別快1.03倍和1.04倍（最高1.08倍和1.14倍）。此外，獨(dú)立kernel測(cè)試表明，我們的SpGEMM比cuSPARSE和rocSPARSE SpGEMM分別快3.09倍、2.40倍和4.67倍（最高達(dá)7.61倍、6.11倍和5.96倍），SpMV在三個(gè)GPU上的性能分別是cuSPARSE和rocSPARSE SpMV的1.34倍、1.19倍和2.92倍（最高達(dá)2.21倍、2.09倍和6.70倍）。
圖4. 不同方法的AMG在三個(gè)GPU上的性能比較
我們還比較了雙精度HYPRE、AmgT和混合精度AmgT在8個(gè)A100 GPU上的性能，如圖5所示。與雙精度HYPRE相比，我們的雙精度AmgT方法的速度快了1.35倍（最高1.84倍）。此外，與雙精度AmgT方法相比，我們的混合精度AmgT方法的速度快了1.06 倍（最高 1.11 倍）。雖然數(shù)據(jù)劃分會(huì)導(dǎo)致通信成本增加、每GPU的計(jì)算量減少，但我們的算法仍能保持相對(duì)于調(diào)用cuSPARSE 的HYPRE的穩(wěn)定優(yōu)勢(shì)。
圖5. HYPRE和AmgT在8個(gè)A100 GPU上的性能比較
作者簡(jiǎn)介:
劉偉峰，教授，博士生導(dǎo)師，歐盟瑪麗居里學(xué)者。2002年和2006年于中國(guó)石油大學(xué)（北京）計(jì)算機(jī)系獲學(xué)士與碩士學(xué)位。2006年至2012年在中國(guó)石化石油勘探開(kāi)發(fā)研究院歷任助理工程師、工程師和高級(jí)研究師，其間主要研究領(lǐng)域?yàn)槭偷厍蛭锢砜碧降母咝阅芩惴ā?016年于丹麥哥本哈根大學(xué)獲計(jì)算科學(xué)博士學(xué)位，主要研究方向?yàn)閿?shù)值線(xiàn)性代數(shù)和并行計(jì)算，其中尤其關(guān)注稀疏矩陣的數(shù)據(jù)結(jié)構(gòu)、并行算法和軟件。研究工作發(fā)表于SC、PPoPP、DAC、ASPLOS、ICS、IPDPS、ICPP、TPDS、JPDC、FGCS和Parco等重要國(guó)際會(huì)議和期刊。擔(dān)任TPDS、SISC和TKDE等多個(gè)重要國(guó)際期刊審稿人，以及SC、ICS、IPDPS和ICPP等多個(gè)重要國(guó)際會(huì)議的程序委員會(huì)委員。他是IEEE高級(jí)會(huì)員、CCF高級(jí)會(huì)員、ACM和SIAM會(huì)員。
聯(lián)系方式：[email protected]。