cuda编程笔记(9)--使用 Shared Memory 实现 tiled GEMM
tiled GEMM
是在 GPU 上使用 共享内存(Shared Memory)优化通用矩阵乘法(GEMM, General Matrix Multiply) 的一种经典方法,其核心思想是 将大矩阵拆分为更小的 tile(子块),再通过共享内存提高缓存命中率和并行计算效率。
任务:给定矩阵 A(M×K),B(K×N),计算 C = A × B,得到矩阵 C(M×N)
使用 tile(块) 的思想:
-
把矩阵分成小块(如 16×16),一块块加载进共享内存;
-
让每个线程处理一个 C[i][j] 元素;
-
每次只加载 A 和 B 的一块(tile),乘积累加进 C 的对应值;
-
这样能显著减少对 global memory 的慢速访问。
先给出代码,再详细讲解
#ifndef __CUDACC__
#define __CUDACC__
#endif
#include <cuda_runtime.h>
#include <device_launch_parameters.h>#include <iostream>
#include<cstdio>#define TILE_SIZE 16
void error_handling(cudaError_t res) {if (res !=cudaSuccess) {std::cout << "error!" << std::endl;}
}
__global__ void tiledGEMM(float* A, float* B, float* C, int M, int N, int K) {//一个block里有TILE_SIZE*TILE_SIZE个线程__shared__ float tileA[TILE_SIZE][TILE_SIZE];__shared__ float tileB[TILE_SIZE][TILE_SIZE];//当前线程负责的C矩阵中的行和列//为什么这么算,TILE_SIZE说成blcokDim.y更直观,因为两者是相等的int row = blockIdx.y * TILE_SIZE + threadIdx.y;int col = blockIdx.x * TILE_SIZE + threadIdx.x;float temp = 0;//每个block内,把对应区域的A和B矩阵填入tileA和tileBfor (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {if (row < M && t * TILE_SIZE + threadIdx.x < K)tileA[threadIdx.y][threadIdx.x] = A[row * K + t * TILE_SIZE + threadIdx.x];elsetileA[threadIdx.y][threadIdx.x] = 0;// 从全局内存拷贝 B 的 tile 到共享内存(小心边界)if (t * TILE_SIZE + threadIdx.y < K && col < N)tileB[threadIdx.y][threadIdx.x] = B[(t * TILE_SIZE + threadIdx.y) * N + col];elsetileB[threadIdx.y][threadIdx.x] = 0;// 等待所有线程把 tile 加载进来__syncthreads();// 每个线程负责 tile 的一行和一列相乘求和for (int k = 0; k < TILE_SIZE; ++k)temp += tileA[threadIdx.y][k] * tileB[k][threadIdx.x];// 所有线程同步,准备下一轮 tile 加载__syncthreads();}// 将结果写回 C(小心越界)if (row < M && col < N)C[row * N + col] = temp;
}int main() {const int M = 64;const int K = 64;const int N = 64;size_t bytesA = M * K * sizeof(float);size_t bytesB = K * N * sizeof(float);size_t bytesC = M * N * sizeof(float);float* h_A = new float[M * K];float* h_B = new float[K * N];float* h_C = new float[M * N];// 初始化 A 和 Bfor (int i = 0; i < M * K; ++i) h_A[i] = 1.0f;for (int i = 0; i < K * N; ++i) h_B[i] = 1.0f;float* d_A, * d_B, * d_C;cudaMalloc(&d_A, bytesA);cudaMalloc(&d_B, bytesB);cudaMalloc(&d_C, bytesC);cudaMemcpy(d_A, h_A, bytesA, cudaMemcpyHostToDevice);cudaMemcpy(d_B, h_B, bytesB, cudaMemcpyHostToDevice);// 设置 Grid 和 Blockdim3 blockDim(TILE_SIZE, TILE_SIZE);dim3 gridDim((N + TILE_SIZE - 1) / TILE_SIZE,(M + TILE_SIZE - 1) / TILE_SIZE);// 调用 kerneltiledGEMM << <gridDim, blockDim >> > (d_A, d_B, d_C, M, N, K);cudaDeviceSynchronize();// 拷贝结果回主机cudaMemcpy(h_C, d_C, bytesC, cudaMemcpyDeviceToHost);// 简单输出前 10 个值检查for (int i = 0; i < 10; ++i)std::cout << h_C[i] << " ";std::cout << std::endl;// 清理内存delete[] h_A;delete[] h_B;delete[] h_C;cudaFree(d_A);cudaFree(d_B);cudaFree(d_C);return 0;
}
在本例子中,BlockDim.x=BlockDim.y=TILE_SIZE
- 因为矩阵C的规模是(M×N)所以总共的线程是也是(M×N)个;
- 每个Block负责矩阵C中一块(BlockDim.x × BlockDim.y)也即(TILE_SIZE × TILE_SIZE)大小的区域;
- 而Block中的每个线程负责C[i][j],也即矩阵C的单独一个位置的值的计算
矩阵C中C[i][j]的计算公式为:
由于每次只能将(TILE_SIZE × TILE_SIZE)大小的矩阵加载进共享内存,而一块(TILE_SIZE × TILE_SIZE)大小的矩阵A区域和一块(TILE_SIZE × TILE_SIZE)大小的矩阵B区域,不一定能够完成C[i][j]的计算
见下图:
红色区域是本Block负责的(TILE_SIZE × TILE_SIZE)的区域,而计算这些区域C矩阵的值需要用到如下A,B矩阵的区域
可是Block内的共享矩阵大小也是 (TILE_SIZE × TILE_SIZE)的,我们只能每次框住对应大小,并且不停迭代来进行移动
我们来逐句解释tiledGEMM核函数:
__shared__ float tileA[TILE_SIZE][TILE_SIZE];__shared__ float tileB[TILE_SIZE][TILE_SIZE];
这是每个Block内共享的小块矩阵,对应的即使图3中蓝色框圈住的A、B的区域
//TILE_SIZE说成blcokDim.y更直观,因为两者是相等的int row = blockIdx.y * TILE_SIZE + threadIdx.y;int col = blockIdx.x * TILE_SIZE + threadIdx.x;
这是计算本线程负责的C[row][col]的位置,计算原理在CUDA编程笔记(1)--最简单的核函数-CSDN博客中获取线程全局 ID小节有讲过,请自行学习
float temp = 0;
对于每个C[row][col],不能一次性算出它的值,得经过多次累加,所以需要声明一个变量进行存储
for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; t++) {}
整个for循环,参考图3的情况,即是蓝色框进行移动,每次迭代,蓝色框就向对应方向进行移动
for循环内部:
if (row < M && t * TILE_SIZE + threadIdx.x < K)tileA[threadIdx.y][threadIdx.x] = A[row * K + t * TILE_SIZE + threadIdx.x];elsetileA[threadIdx.y][threadIdx.x] = 0;// 从全局内存拷贝 B 的 tile 到共享内存(小心边界)if (t * TILE_SIZE + threadIdx.y < K && col < N)tileB[threadIdx.y][threadIdx.x] = B[(t * TILE_SIZE + threadIdx.y) * N + col];elsetileB[threadIdx.y][threadIdx.x] = 0;
这些代码的作用就是把蓝色框内的A、B矩阵是数据转移到共享内存tileA,tileB中。A,B是将逻辑上的二维矩阵转为实际存储的一维数组;其中A[],B[]的下标计算即是按照行/列优先的方式来计算的,这里不懂的去自行复习一下矩阵下标的计算;
__syncthreads();
第一个__syncthreads();是等待所有线程都把当前蓝色框的内容转移到tileA和tileB中
for (int k = 0; k < TILE_SIZE; ++k)temp += tileA[threadIdx.y][k] * tileB[k][threadIdx.x];
一个C[row][col]对应A的一长行和B的一长列,但是一次tileA和tileB只能加载一部分,只能先把tileA和tileB里面的对应部分相乘加起来存起来
__syncthreads();
所有线程同步,准备下一轮 tile 加载
dim3 blockDim(TILE_SIZE, TILE_SIZE);
dim3 gridDim((N + TILE_SIZE - 1) / TILE_SIZE,(M + TILE_SIZE - 1) / TILE_SIZE);
这里blockDim和gridDim总乘积为M*N(M、N为TILE_SIZE倍数的情况下),刚好每个线程对应矩阵C的一个元素