cuda 加速矩陣乘法

拾墨、發表於2024-03-15

對於一個m * n的矩陣a和一個n * k的矩陣b

因為最後得到一個m * k的矩陣c,那麼我們可以分配m * k個執行緒。

線上程(i,j)裡矩陣a的第i行和矩陣b的第j列進行點積運算得到c[i][j]

#include<iostream>
#include "cuda_runtime.h"
#include "device_launch_parameters.h"

# define BLOCK_SIZE 2

__global__ void gpu_matrix_mult(int* a, int* b, int* c, int m, int n, int k)
{
        //row和col是該執行緒所在行數和列數
	int row = blockIdx.y * blockDim.y + threadIdx.y;
	int col = blockIdx.x * blockDim.x + threadIdx.x;

	int sum = 0;
	if (col < k && row < m)
	{
		for (int i = 0; i < n; i++)
		{
			sum += a[row * n + i] * b[i * k + col];
		}
		c[row * k + col] = sum;
	}
}
int main()
{
	int m = 100, n = 100, k = 100;
	
	int* h_a, * h_b, * h_c;
	cudaMallocHost((void**)&h_a, sizeof(int) * m * n);
	cudaMallocHost((void**)&h_b, sizeof(int) * n * k);
	cudaMallocHost((void**)&h_c, sizeof(int) * m * k);

	for (int i = 0; i < m; ++i)
	{
		for (int j = 0; j < n; ++j) 
			h_a[i * n + j] = rand() % 1024;
	}
	for (int i = 0; i < n; ++i)
	{
		for (int j = 0; j < k; ++j)
			h_b[i * k + j] = rand() % 1024;
	}

	int* d_a, * d_b, * d_c;
	cudaMalloc((void**)&d_a, sizeof(int) * m * n);
	cudaMalloc((void**)&d_b, sizeof(int) * n * k);
	cudaMalloc((void**)&d_c, sizeof(int) * m * k);

	cudaMemcpy(d_a, h_a, sizeof(int) * m * n, cudaMemcpyHostToDevice);
	cudaMemcpy(d_b, h_b, sizeof(int) * n * k, cudaMemcpyHostToDevice);
	
        //BLOCK_SIZE是一個block邊的大小
        //grid_rows是一個grid有幾行block
        //grid_cols是一個grid有幾列block
        //dimGrid是一個grid一行有幾個block,一列有幾個block
        //dimBlock是一個block一行有幾個thread,一列有幾個thread
	unsigned int grid_rows = (m + BLOCK_SIZE - 1) / BLOCK_SIZE; 
	unsigned int grid_cols = (k + BLOCK_SIZE - 1) / BLOCK_SIZE;
	dim3 dimGrid(grid_cols, grid_rows);
	dim3 dimBlock(BLOCK_SIZE, BLOCK_SIZE);

	gpu_matrix_mult<<<dimGrid , dimBlock>>>(d_a, d_b, d_c, m, n, k);
	cudaMemcpy(h_c, d_c, sizeof(int) * m * k, cudaMemcpyDeviceToHost);
	for (int i = 0; i < m*k; i++)
	{
		std::cout << h_c[i] << std::endl;
	}
	return 0;
}

相關文章