如何對(duì)spmv算法進(jìn)行優(yōu)化

主要是介紹如何對(duì)spmv算法進(jìn)行優(yōu)化。Spmv，即稀疏化的矩陣向量乘操作，關(guān)于稠密的矩陣向量乘操作，已經(jīng)在上一篇文章中介紹過(guò)了。關(guān)于稀疏kernel的優(yōu)化，是CUDA優(yōu)化中最難的一部分，其難度在于稀疏特性千差萬(wàn)別，需要針對(duì)不同的應(yīng)用、不同的數(shù)據(jù)選擇不同的數(shù)據(jù)存儲(chǔ)格式，然后再根據(jù)不同的數(shù)據(jù)特點(diǎn)進(jìn)行特定的并行算法設(shè)計(jì)。而現(xiàn)實(shí)生活中，尤其是科學(xué)計(jì)算里面，基本上都是稀疏問(wèn)題。在深度學(xué)習(xí)領(lǐng)域中，也一直有針對(duì)稀疏模型的研究，主要是針對(duì)推理方向，將模型進(jìn)行剪枝之后，直接減少了計(jì)算量來(lái)達(dá)到對(duì)模型的加速目的。但實(shí)際上，為了保證模型的精度，稀疏度有限，且稀疏問(wèn)題很難充分地利用硬件性能，導(dǎo)致了這一條路線其實(shí)并不好走。嘮嘮叨叨說(shuō)了挺多，總之，稀疏kernel的優(yōu)化是一個(gè)非常難的話題。本文會(huì)詳細(xì)地介紹一下spmv。

一、前言

在說(shuō)spmv之前，說(shuō)一下稀疏格式。當(dāng)矩陣中的絕大多數(shù)元素都是0時(shí)，需要一些特殊的格式來(lái)存儲(chǔ)非零元素。這些格式也就是稀疏格式，常用的稀疏格式有：COO、CSR、DIA、ELL、HYB。深度學(xué)習(xí)領(lǐng)域還有blocked CSR、blocked ELL等。具體的稀疏格式總結(jié)見以下鏈接：

稀疏矩陣存儲(chǔ)格式總結(jié)+存儲(chǔ)效率對(duì)比:COO,CSR,DIA,ELL,HYB - Bin的專欄 - 博客園www.cnblogs.com/xbinworld/p/4273506.html

在本文中，使用CSR格式存儲(chǔ)稀疏矩陣，后續(xù)所說(shuō)的一系列優(yōu)化也是針對(duì)CSR格式而言。說(shuō)完了稀疏格式，現(xiàn)在再來(lái)說(shuō)一下spmv，即稀疏矩陣向量乘。稠密的矩陣向量乘，即gemv已經(jīng)在之前說(shuō)過(guò)了。具體的操作即給定稀疏矩陣A和向量x，需要計(jì)算兩者的乘積y。示意圖如下。

spmv介紹

二、并行算法設(shè)計(jì)

并行算法設(shè)計(jì)，主要是block和thread的設(shè)計(jì)，在這里主要參考了cusp的實(shí)現(xiàn)。有一個(gè)很重要的考慮是workload的分配。我們需要使用多少個(gè)線程來(lái)負(fù)責(zé)A矩陣中一行的計(jì)算？需要說(shuō)明的是，在不同的數(shù)據(jù)特性下，需要采用不同的取值。如果這一行的元素非常多，那使用一個(gè)warp或者一個(gè)block，如果這一行只有一個(gè)元素，只需要一次乘加指令，那顯然只能使用一個(gè)線程，畢竟用兩個(gè)線程處理一個(gè)元素，怎么都不像正常人能干出的事。因而，我們假定一個(gè)參數(shù)，即THREADS_PER_VECTOR，來(lái)代表這個(gè)值。每THREADS_PER_VECTOR個(gè)線程為一組，他們需要負(fù)責(zé)A矩陣中一行元素的計(jì)算。

說(shuō)完了這個(gè)核心思路，接下來(lái)看看每個(gè)線程需要干的工作。每個(gè)線程都要單獨(dú)地對(duì)A矩陣的offset和index等進(jìn)行讀取，然后計(jì)算當(dāng)前行的結(jié)果。如果每一行的元素特別少，比如這一行元素有4個(gè)，THREADS_PER_VECTOR就設(shè)為4，有8個(gè)元素就設(shè)為8，平均多少個(gè)元素，THREADS_PER_VECTOR就設(shè)為幾，但上限是32。元素比32多的話，就多進(jìn)行幾個(gè)迭代即可?？傊疃嗍褂靡粋€(gè)warp來(lái)處理一行元素。

至于如何得到一行元素有多少，row_offset數(shù)組長(zhǎng)度除以y數(shù)組長(zhǎng)度即可得。有必要在這里再提一句的是，這些參數(shù)的選擇，以及對(duì)于不均衡的A矩陣元素如何處理，這些都是比較棘手的問(wèn)題。有一大堆的論文在談負(fù)載均衡和自動(dòng)調(diào)參這兩個(gè)事情，大家可以搜一下相關(guān)的論文瞅瞅。

講完了思路，下面說(shuō)一下具體的代碼，如下：

template 
__device__ __forceinline__ float warpReduceSum(float sum) {
    if (WarpSize >= 32)sum += __shfl_down_sync(0xffffffff, sum, 16); // 0-16, 1-17, 2-18, etc.
    if (WarpSize >= 16)sum += __shfl_down_sync(0xffffffff, sum, 8);// 0-8, 1-9, 2-10, etc.
    if (WarpSize >= 8)sum += __shfl_down_sync(0xffffffff, sum, 4);// 0-4, 1-5, 2-6, etc.
    if (WarpSize >= 4)sum += __shfl_down_sync(0xffffffff, sum, 2);// 0-2, 1-3, 4-6, 5-7, etc.
    if (WarpSize >= 2)sum += __shfl_down_sync(0xffffffff, sum, 1);// 0-1, 2-3, 4-5, etc.
    return sum;
}

template 
__global__ void My_spmv_csr_kernel(const IndexType row_num,
                       const IndexType * A_row_offset,
                       const IndexType * A_col_index,
                       const ValueType * A_value,
                       const ValueType * x,
                       ValueType * y)
{
    const IndexType THREADS_PER_BLOCK = VECTORS_PER_BLOCK * THREADS_PER_VECTOR;
    const IndexType thread_id   = THREADS_PER_BLOCK * blockIdx.x + threadIdx.x;    // global thread index
    const IndexType thread_lane = threadIdx.x & (THREADS_PER_VECTOR - 1);          // thread index within the vector
    const IndexType row_id   = thread_id   /  THREADS_PER_VECTOR;               // global vector index

    if(row_id < row_num){
        const IndexType row_start = A_row_offset[row_id];                  //same as: row_start = Ap[row];
        const IndexType row_end   = A_row_offset[row_id+1];

        // initialize local sum
        ValueType sum = 0;

        // accumulate local sums
        for(IndexType jj = row_start + thread_lane; jj < row_end; jj += THREADS_PER_VECTOR)
            sum += A_value[jj] * x[ A_col_index[jj] ];

        sum = warpReduceSum(sum);
        if (thread_lane == 0){
            y[row_id] = sum;
        }   
    }
}

首先說(shuō)一下傳入的幾個(gè)模板參數(shù)，IndexType代表索引類型，一般用int，如果矩陣十分巨大，可以考慮long long，ValueType代表數(shù)值存儲(chǔ)的格式，科學(xué)計(jì)算一般都雙精度，深度學(xué)習(xí)一般用單精度，或者fp16。推理甚至有一些int8這樣的需求。THREADS_PER_VECTOR這個(gè)參數(shù)已經(jīng)在前面說(shuō)過(guò)了，VECTORS_PER_BLOCK則是代表一個(gè)block中有多少vector。我們盡可能地保證一個(gè)block有256個(gè)線程，所以VECTORS_PER_BLOCK = 256 / THREADS_PER_VECTOR。看完了模板參數(shù)，再看函數(shù)輸入?yún)?shù)，分別是A矩陣的行數(shù)，A矩陣的CSR表示，有3個(gè)數(shù)組，然后是向量x和向量y。

接下來(lái)到了具體的kernel邏輯，首先計(jì)算了四個(gè)參數(shù)，需要注意的是thread_lane和row_id參數(shù)，thread_lane代表當(dāng)前元素是當(dāng)前組里面的第幾個(gè)元素，row_id代表當(dāng)前元素負(fù)責(zé)A矩陣中第幾行的計(jì)算。接下來(lái)的邏輯也比較明了，先計(jì)算當(dāng)前行對(duì)應(yīng)的索引值，即row_start和row_end。定義sum變量來(lái)存儲(chǔ)該行的計(jì)算結(jié)果，而后進(jìn)行多次迭代，將每個(gè)線程對(duì)應(yīng)的sum取出來(lái)，最后將sum元素進(jìn)行warp的reduce_sum操作。最后將元素寫到y(tǒng)向量中。

三、優(yōu)化技巧

在上一節(jié)中已經(jīng)把代碼說(shuō)完了，接下來(lái)盤點(diǎn)一下具體的優(yōu)化技巧，以及優(yōu)化中需要考慮的方方面面。

1、合理的block和thread調(diào)整。我一直覺(jué)得這個(gè)點(diǎn)是優(yōu)化中最重要的一點(diǎn)。核心就是THREADS_PER_VECTOR需要根據(jù)實(shí)際的數(shù)據(jù)進(jìn)行調(diào)整。這一點(diǎn)主要是考慮到如果使用更多的線程處理的話，只有THREADS_PER_VECTOR個(gè)線程在工作，其他的線程都被浪費(fèi)了。

2、Shuffle指令減少訪存的latency。在不使用shuffle指令的話，只能通過(guò)shared memory完成最后的求和操作。從shared memory中取數(shù)比寄存器之間直接訪問(wèn)要花費(fèi)更多的latency。因而要盡可能地使用shuffle指令。

3、對(duì)于global memory的合并訪存。對(duì)于稀疏問(wèn)題，由于CSR格式中的col數(shù)組和val數(shù)組不能保證地址對(duì)齊，因而針對(duì)global memory的合并訪存其實(shí)是有一定的困難。我們可以仔細(xì)地來(lái)進(jìn)行分析。當(dāng)A矩陣行數(shù)比較多的情況下，spmv主要的訪存有3部分，分別是A_value，A_col_index和x。其中，對(duì)于A_value和A_col_index的訪存是連續(xù)的，但是由于地址不能保證對(duì)齊，所以訪存效率大概率不會(huì)太高。而對(duì)于x的訪存本身就是不連續(xù)的，因而cache命中率會(huì)顯然易見地低。如何解決這些問(wèn)題呢？對(duì)于A_value和A_col_index的訪存問(wèn)題，尚可以嘗試對(duì)其進(jìn)行數(shù)據(jù)填充，強(qiáng)制其地址對(duì)齊。而對(duì)于x的非連續(xù)訪存問(wèn)題，如何提高訪存效率，這個(gè)問(wèn)題就非常困難了。

4、關(guān)于向量化指令的使用。之前在進(jìn)行g(shù)emm和gemv優(yōu)化中大量地使用了float4這樣的向量化訪存結(jié)構(gòu)。如何將向量化帶到spmv中，這也是一個(gè)非常困難的問(wèn)題。最大的根源是因?yàn)槊恳恍械脑夭淮_定，并且本身A中每行的元素就比較少，根本沒(méi)有那么多數(shù)據(jù)去喂到LDS128指令上。

5、關(guān)于負(fù)載均衡的思考。CUDA上的負(fù)載均衡問(wèn)題可以從兩個(gè)角度考慮，一個(gè)是block之間的負(fù)載均衡，另一個(gè)是block/warp內(nèi)，不同線程之間的負(fù)載均衡。關(guān)于spmv的負(fù)載均衡問(wèn)題，可以參考一下Speculative Segmented Sum for Sparse Matrix-Vector Multiplication on Heterogeneous Processors。

四、實(shí)驗(yàn)與總結(jié)

最后，我們來(lái)說(shuō)一下實(shí)驗(yàn)環(huán)節(jié)。實(shí)驗(yàn)主要用來(lái)說(shuō)明兩個(gè)問(wèn)題，第一個(gè)是THREADS_PER_VECTOR參數(shù)對(duì)性能的影響，第二個(gè)是與cusparse的對(duì)比，用以觀察不同數(shù)據(jù)下的性能表現(xiàn)。

實(shí)驗(yàn)一，從佛羅里達(dá)矩陣庫(kù)里面選了一個(gè)稀疏矩陣，shyy41。平均一行有4.2個(gè)元素。我們?cè)诓煌膮?shù)下進(jìn)行了實(shí)驗(yàn)，其結(jié)果如下：

這個(gè)結(jié)果和預(yù)期的差不多，因?yàn)槠骄恍性貍€(gè)數(shù)是4.2，所以THREADS_PER_VECTOR參數(shù)取4或8會(huì)有更好的性能表現(xiàn)。

實(shí)驗(yàn)二，從佛羅里達(dá)矩陣庫(kù)里面隨機(jī)選取了一些矩陣，其稀疏特性如下，矩陣旁邊有x-y-z標(biāo)識(shí)。x和y代表矩陣的行數(shù)和列數(shù)，z代表矩陣中的非零元個(gè)數(shù)。

稀疏矩陣

性能表現(xiàn)如下，

與cusparse的性能對(duì)比

結(jié)論：

1、先單獨(dú)看cusparse的表現(xiàn)，庫(kù)里面會(huì)調(diào)用兩個(gè)kernel，分別是binary_seach和load_balance。這個(gè)名稱簡(jiǎn)寫了?？傊?，就是cusparse不管來(lái)的數(shù)據(jù)是啥，都會(huì)進(jìn)行負(fù)載均衡，在數(shù)據(jù)量比較多的時(shí)候，額外的開銷比較少，能夠取到足夠的效益。

2、如果是結(jié)構(gòu)化的網(wǎng)格，即元素聚集在對(duì)角線附近，且每行的非零元都差不了太多的時(shí)候，我寫的spmv會(huì)比cusparse快一些。如果每行的非零元方差特別大，cusparse中的負(fù)載均衡工作就發(fā)揮了威力，在web網(wǎng)絡(luò)這種矩陣上能夠比我的spmv快2-3倍。總之，在sparse問(wèn)題中，負(fù)載均衡非常重要，我會(huì)在下一篇博文中說(shuō)明如何在spmm中進(jìn)行負(fù)載均衡。

總之，我們實(shí)現(xiàn)了spmv kernel，并對(duì)主要的優(yōu)化技巧進(jìn)行了解析和說(shuō)明，然后大概地說(shuō)了一下在spmv上需要注意的問(wèn)題。通過(guò)實(shí)驗(yàn)評(píng)估了不同參數(shù)對(duì)性能的影響以及在不同的稀疏矩陣下同cusparse進(jìn)行了比較，在部分矩陣上性能能夠超越cusparse。但由于沒(méi)有考慮負(fù)載均衡，在非均勻網(wǎng)格上，與cusparse有一定的差距。