【CUDA】矩陣乘法 matMatMul

矩陣乘法 matMatMul

矩陣乘法是基本線性代數子程序（BLAS）的重要組成部分，而且線性代數中許多其他操作以此為基礎。

圖1是兩個矩陣的乘法。

基礎方法，正方形tile和長方形tile

基礎方法

執行矩陣乘法的基礎方法是使用單個線程執行輸出矩陣的單個元素乘法。這意味著所需的線程數等于輸出矩陣中的元素數。線程排列在二維網格中，每個線程分配一個唯一的索引。線程的索引用于訪問輸入矩陣的相應行和列。執行行和列的乘法，結果存儲在輸出矩陣的相應元素中。

這種方法的主要缺點是多個線程加載相同的行和列來計算它們的輸出。圖2是一個示例。

圖2

正如上圖所示，為了計算P0,0和P0,1，兩個線程都需要加載整個M0。對于P0,0和P1,0也是如此。兩個線程都需要加載整個N1列。這意味著線程將多次訪問相同的內存位置。

正方形tile

為了避免這個問題，我們可以使用tile。tile是將輸入矩陣的一個小部分加載到共享內存中。線程將把tile加載到共享內存中，然后執行乘法運算。圖3描述了這種技術。

圖3

kernel將計算分為幾個階段。每個階段，線程將將M和N矩陣中的tile加載到共享內存中。然后，線程將執行tile的乘法，并將結果累積到輸出矩陣的相應元素中。

通過這種技術，我們可以看到全局內存訪問減少了TILE_WIDTH（圖示）倍。

長方形tile

為了進一步減少全局內存訪問，我們可以使用圖4所示的矩形塊。

圖4

通過這種技術，我們增加了每個線程的工作量，以進一步減少全局內存訪問的次數。kernel再次將計算分成多個階段。在每個階段中，線程將從中加載一個tile M 和兩個tile N 到共享內存中。然后線程將對這些tile進行乘法運算，并將結果累加到輸出矩陣的相應元素中。

Code

host代碼使用隨機值初始化輸入矩陣，并調用kernel執行乘法運算。輸入矩陣以線性格式存儲。


#include <iostream>#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
#include <thrust/transform.h>#include "mat_mat_mul_kernels.cuh"int main(int argc, char *argv[])
{int rows1, cols1rows2, cols2, t_x;if(argc != 5){std::cout << "Usage: " << argv[0] << " <rows1> <cols1rows2> <cols2> <t_x>" << std::endl;return 1;}rows1 = atoi(argv[1]);cols1rows2 = atoi(argv[2]);cols2 = atoi(argv[3]);t_x = atoi(argv[4]);thrust::host_vector<float> h_in_mat1(rows1 * cols1rows2);thrust::host_vector<float> h_in_mat2(cols1rows2 * cols2);thrust::host_vector<float> h_out_host(rows1 * cols2);srand(time(NULL));for(int i = 0; i < rows1 * cols1rows2; i++)h_in_mat1[i] = rand() / (float)RAND_MAX;for(int i = 0; i < cols1rows2 * cols2; i++)h_in_mat2[i] = rand() / (float)RAND_MAX;for(int i = 0; i < rows1; ++i)for(int j = 0; j < cols2; ++j)for(int k = 0; k < cols1rows2; ++k)h_out_host[i * cols2 + j] += h_in_mat1[i * cols1rows2 + k] * h_in_mat2[k * cols2 + j];thrust::device_vector<float> d_in_mat1 = h_in_mat1;thrust::device_vector<float> d_in_mat2 = h_in_mat2;thrust::device_vector<float> d_out(rows1 * cols2);dim3 blockDim(t_x, t_x);dim3 gridDim((cols2 + blockDim.x - 1) / blockDim.x, (rows1 + blockDim.y - 1) / blockDim.y);mat_mat_mul<float><<<gridDim, blockDim>>>(thrust::raw_pointer_cast(d_in_mat1.data()),thrust::raw_pointer_cast(d_in_mat2.data()),thrust::raw_pointer_cast(d_out.data()),rows1, cols1rows2, cols1rows2, cols2);bool success = true;for(int i = 0; i < rows1 * cols2; ++i)if(abs(h_out_host[i] - d_out[i]) >= 0.001){std::cout << "Error at " << i << ": " << h_out_host[i] << " != " << d_out[i] << " (Mat Mat Mul)" << std::endl;success = false;break;}if(success)std::cout << "Success (Mat Mat Mul)" << std::endl;blockDim = dim3(t_x, t_x);gridDim = dim3((cols2 + blockDim.x - 1) / blockDim.x, (rows1 + blockDim.y - 1) / blockDim.y);mat_mat_mul_tiles<float><<<gridDim, blockDim, 2 * blockDim.x * blockDim.x * sizeof(float)>>>(thrust::raw_pointer_cast(d_in_mat1.data()),thrust::raw_pointer_cast(d_in_mat2.data()),thrust::raw_pointer_cast(d_out.data()),rows1, cols1rows2, cols1rows2, cols2);success = true;for(int i = 0; i < rows1 * cols2; ++i)if(abs(h_out_host[i] - d_out[i]) >= 0.001){std::cout << "Error at " << i << ": " << h_out_host[i] << " != " << d_out[i] << " (Mat Mat Mul Tiles)" << std::endl;success = false;break;}if(success)std::cout << "Success (Mat Mat Mul Tiles)" << std::endl;blockDim = dim3(t_x, t_x);gridDim = dim3((cols2 + blockDim.x - 1) / blockDim.x / 2, (rows1 + blockDim.y - 1) / blockDim.y);mat_mat_mul_rec_tiles<float><<<gridDim, blockDim, 3 * blockDim.x * blockDim.x * sizeof(float)>>>(thrust::raw_pointer_cast(d_in_mat1.data()),thrust::raw_pointer_cast(d_in_mat2.data()),thrust::raw_pointer_cast(d_out.data()),rows1, cols1rows2, cols1rows2, cols2);success = true;for(int i = 0; i < rows1 * cols2; ++i)if(abs(h_out_host[i] - d_out[i]) >= 0.001){std::cout << "Error at " << i << ": " << h_out_host[i] << " != " << d_out[i] << " (Mat Mat Mul Rec Tiles)" << std::endl;success = false;break;}if(success)std::cout << "Success (Mat Mat Mul Rec Tiles)" << std::endl;return 0;}

以下是kernel的展示。

基礎方法

template<typename T> __global__
void mat_mat_mul(T *in_mat1, T *in_mat2, T *out_mat, int mat1_rows, int mat1_cols, int mat2_rows, int mat2_cols) {int row = blockIdx.y * blockDim.y + threadIdx.y;int col = blockIdx.x * blockDim.x + threadIdx.x;if (row >= mat1_rows || col >= mat2_cols) return;T sum = 0;for (int k = 0; k < mat1_cols; ++k)sum += in_mat1[row * mat1_cols + k] * in_mat2[k * mat2_cols + col];out_mat[row * mat2_cols + col] = sum;
}

在這種基本方法中，kernel非常簡單。

線程首先計算它們的行和列索引。如果行或列索引大于輸入矩陣的行數或列數，則線程返回。

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;if (row >= mat1_rows || col >= mat2_cols) return;

然后，這些線程會對行和列進行相乘，并將結果存儲在輸出矩陣的相應元素中。

T sum = 0;
for (int k = 0; k < mat1_cols; ++k)sum += in_mat1[row * mat1_cols + k] * in_mat2[k * mat2_cols + col];out_mat[row * mat2_cols + col] = sum;

我們可以觀察到，線程不僅會多次加載相同的行和列，而且它們還會以非合并的方式加載M的元素。這意味著線程將無法充分利用內存帶寬。

正方形tile

template<typename T> __global__
void mat_mat_mul_tiles(T *in_mat1, T *in_mat2, T *out_mat,int mat1_rows, int mat1_cols,int mat2_rows, int mat2_cols) {// Initialize shared memoryint TILE_WIDTH = blockDim.x;extern __shared__ uint8_t shared_mem[];T *ds_mat1 = reinterpret_cast<T*>(shared_mem);T *ds_mat2 = reinterpret_cast<T*>(shared_mem + TILE_WIDTH * TILE_WIDTH * sizeof(T));int bx = blockIdx.x,  by = blockIdx.y;int tx = threadIdx.x, ty = threadIdx.y;int row = by * TILE_WIDTH + ty;int col = bx * TILE_WIDTH + tx;T out_value = 0;// Loop over the in_mat1 and in_mat2 tiles required to compute the out_mat elementfor (int ph = 0; ph < (mat1_cols + TILE_WIDTH - 1) / TILE_WIDTH; ++ph) {// Collaborative loading of in_mat1 and in_mat2 tiles into shared memoryds_mat1[ty * TILE_WIDTH + tx] = row < mat1_rows && ph * TILE_WIDTH + tx < mat1_cols ?in_mat1[row * mat1_cols + ph * TILE_WIDTH + tx] : 0;ds_mat2[ty * TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows && col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + col] : 0;// Synchronize to make sure the tiles are loaded__syncthreads();// Compute the out_mat elementfor (int k = 0; k < TILE_WIDTH; ++k)out_value += ds_mat1[ty * TILE_WIDTH + k] * ds_mat2[k * TILE_WIDTH + tx];// Synchronize to make sure the out_mat element is computed// before other threads load new tiles__syncthreads();}// Store the out_mat element in out_matif (row < mat1_rows && col < mat2_cols)out_mat[row * mat2_cols + col] = out_value;
}

在這種方法中，我們使用共享內存來存儲輸入矩陣的塊。線程首先計算它們的行和列索引。如果行或列索引大于輸入矩陣的行數或列數，則線程返回。

int row = by * TILE_WIDTH + ty;
int col = bx * TILE_WIDTH + tx;if (row >= mat1_rows || col >= mat2_cols) return;

在每個階段：

線程將tiles加載到共享內存中。這些tiles以合并訪存的方式加載。

// Collaborative loading of in_mat1 and in_mat2 tiles into shared memory
ds_mat1[ty * TILE_WIDTH + tx] = row < mat1_rows && ph * TILE_WIDTH + tx < mat1_cols ?in_mat1[row * mat1_cols + ph * TILE_WIDTH + tx] : 0;ds_mat2[ty * TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows && col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + col] : 0;// Synchronize to make sure the tiles are loaded
__syncthreads();

當線程加載完tile后，它們會計算輸出矩陣元素。線程通過將tiles的相應行和列相乘，并將結果相加來計算輸出矩陣元素。

// Compute the out_mat element
for (int k = 0; k < TILE_WIDTH; ++k)out_value += ds_mat1[ty * TILE_WIDTH + k] * ds_mat2[k * TILE_WIDTH + tx];// Synchronize to make sure the out_mat element is computed
// before other threads load new tiles
__syncthreads();

最后，線程將輸出矩陣元素存儲在輸出矩陣中。

// Store the out_mat element in out_mat
out_mat[row * mat2_cols + col] = out_value;

長方形tiles

template<typename T> __global__
void mat_mat_mul_rec_tiles(T* in_mat1, T* in_mat2, T* out_mat,int mat1_rows, int mat1_cols,int mat2_rows, int mat2_cols) {// Initialize shared memoryint TILE_WIDTH = blockDim.x;extern __shared__ uint8_t shared_mem[];T* ds_mat1 = reinterpret_cast<T*>(shared_mem);T* ds_mat2 = reinterpret_cast<T*>(shared_mem + TILE_WIDTH * TILE_WIDTH * sizeof(T));T* ds_mat3 = reinterpret_cast<T*>(shared_mem + 2 * TILE_WIDTH * TILE_WIDTH * sizeof(T));int bx = blockIdx.x, by = blockIdx.y;int tx = threadIdx.x, ty = threadIdx.y;int row = by * TILE_WIDTH + ty;int col = bx * TILE_WIDTH * 2 + tx;T out_value1 = 0;T out_value2 = 0;// Loop over the in_mat1 and in_mat2 tiles required to compute the out_mat elementfor (int ph = 0; ph < (mat1_cols + TILE_WIDTH - 1) / TILE_WIDTH; ++ph) {// Collaborative loading of in_mat1 and in_mat2 tiles into shared memoryds_mat1[ty * TILE_WIDTH + tx] = row < mat1_rows&& ph* TILE_WIDTH + tx < mat1_cols ?in_mat1[row * mat1_cols + ph * TILE_WIDTH + tx] : 0;ds_mat2[ty * TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows&& col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + col] : 0;ds_mat3[ty * TILE_WIDTH + TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows&& TILE_WIDTH + col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + TILE_WIDTH + col] : 0;// Synchronize to make sure the tiles are loaded__syncthreads();// Compute the out_mat elementfor (int k = 0; k < TILE_WIDTH; k++) {out_value1 += ds_mat1[ty * TILE_WIDTH + k] * ds_mat2[k * TILE_WIDTH + tx];out_value2 += ds_mat1[ty * TILE_WIDTH + k] * ds_mat3[k * TILE_WIDTH + TILE_WIDTH + tx];}// Synchronize to make sure the Pvalues are computed// before other threads load new tiles__syncthreads();}// Store the Pvalues in out_matif (row < mat1_rows && col < mat2_cols)//out_mat[row][col];out_mat[row * mat2_cols + col] = out_value1;if (row < mat1_rows && TILE_WIDTH + col < mat2_cols)//out_mat[row][TILE_WIDTH + col];out_mat[row * mat2_cols + TILE_WIDTH + col] = out_value2;
}

這個kernel函數幾乎與正方形tile函數相同。

首先，線程計算出他們將計算的輸出矩陣元素的行和列。

int row = by * TILE_WIDTH     + ty;
int col = bx * TILE_WIDTH * 2 + tx;

然后線程將tile加載到共享內存中。唯一的區別是N加載2個tile。

// Collaborative loading of in_mat1 and in_mat2 tiles into shared memoryds_mat1[ty * TILE_WIDTH + tx] = row < mat1_rows&& ph* TILE_WIDTH + tx < mat1_cols ?in_mat1[row * mat1_cols + ph * TILE_WIDTH + tx] : 0;ds_mat2[ty * TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows&& col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + col] : 0;ds_mat3[ty * TILE_WIDTH + TILE_WIDTH + tx] = ph * TILE_WIDTH + ty < mat2_rows&& TILE_WIDTH + col < mat2_cols ?in_mat2[(ph * TILE_WIDTH + ty) * mat2_cols + TILE_WIDTH + col] : 0;

然后線程計算兩個輸出矩陣元素。

// Compute the out_mat elementfor (int k = 0; k < TILE_WIDTH; k++) {out_value1 += ds_mat1[ty * TILE_WIDTH + k] * ds_mat2[k * TILE_WIDTH + tx];out_value2 += ds_mat1[ty * TILE_WIDTH + k] * ds_mat3[k * TILE_WIDTH + TILE_WIDTH + tx];}

最后，存儲兩個輸出矩陣元素。

// Store the Pvalues in out_mat
if (row < mat1_rows && col < mat2_cols)//out_mat[row][col];out_mat[row * mat2_cols + col] = out_value1;if (row < mat1_rows && TILE_WIDTH + col < mat2_cols)//out_mat[row][TILE_WIDTH + col];out_mat[row * mat2_cols + TILE_WIDTH + col] = out_value2;

性能分析

運行時間：

矩陣維度：1024 × 1024

線程塊維度：32 × 32

可見使用共享內存可以有效降低運行速度。但長方形tile反而耗時更久。因為L1緩存與共享內存公用硬件空間。可能共享內存占據大部分空間，而L1緩存所剩無幾，從而導致長方形tile耗時更久。具體情況需要做性能分析，之后再補充。也許是長方形tile的復雜性增加。kernel引入了許多分支和同步點。這可能導致耗時更久。