当前位置： 首页 > news >正文

【CUDA】Sgemm单精度矩阵乘法（上）

news 来源：原创 2025/5/15 9:58:13

目录

    • 前言
    • 1. 简述
    • 2. 框架搭建和CPU串行实现
    • 3. baseline算法：global memory
    • 4. 优化技巧1：shared memory
    • 5. 优化技巧2：shared memory + sliding window
    • 6. 优化技巧3：增加每个线程的工作量
    • 7. 优化技巧4：使用float4提高读取效率
    • 结语
    • 下载链接
    • 参考

前言

学习 UP 主 比飞鸟贵重的多_HKL 的 【CUDA】Sgemm单精度矩阵乘法（已完结~） 视频，记录下个人学习笔记，仅供自己参考😄

refer 1：【CUDA】Sgemm单精度矩阵乘法（已完结~）

refer 2：https://github.com/tpoisonooo/how-to-optimize-gemm/cuda

refer 3：https://chatgpt.com/

1. 简述

我们先来通俗的解释下 gemm、sgemm、gemv、spmv 等一些耳熟能详的名词缩写

GEMM（Genera Matrix-Matrix Multiplication）

含义：

  • GEMM 即 “通用矩阵-矩阵乘法”，英文全称 General Matrix-Matrix Multiplication
  • 它描述的是两个稠密（Dense）矩阵之间的乘法，通用形式为：

C = A × B C = A \times B C=A×B

其中：

  • A A A 是一个 ( m × k ) \left(m \times k \right) (m×k) 的矩阵
  • B B B 是一个 ( k × n ) \left(k \times n \right) (k×n) 的矩阵
  • C C C 是一个 ( m × n ) \left(m \times n \right) (m×n) 的矩阵
  • Note：这里每个矩阵的维度是大家约定俗成的，在后续的代码实现以及文章中我们也会按照这个维度来讲解

特点：

  • GEMM 一般由 BLAS（Basic Linear Algebra Subprograms）库提供
  • 常用于深度学习（如卷积、全连接层、attention 等都可以转化为 GEMM）和高性能科学计算领域
  • 它是性能优化的重点，GPU 库（如 cuBLAS）都会对 GEMM 做高度优化

SGEMM（Single Precision GEMM）

含义：

  • SGEMM 是 GEMM 的一种特定实现，表示单精度浮点数（Single precision，float）下的通用矩阵乘法
  • 完整名称为 Single Precision General Matrix-Matrix Multiplication
  • SGEMM 在 GPU 中应用非常广泛，因为 GPU 对单精度计算的优化程度很高

区别与特点：

  • 与 GEMM 最大的区别是 SGEMM 明确指出矩阵元素采用的是 32 位浮点数（float）
  • 如果使用的是双精度浮点数（double），则对应的称为 DGEMM（Double precision GEMM）
  • SGEMM 通常比 DGEMM 快，内存占用更少，在深度学习序列和推理等场景中使用较多

GEMV（General Matrix-Vector Multiplication）

含义：

  • GEMV 即 “通用矩阵-向量乘法”，英文全称 General Matrix-Vector Multiplication
  • 通用数学形式为：

y = A × x y = A \times x y=A×x

其中：

  • A A A 是一个 ( m × n ) \left(m \times n \right) (m×n) 的矩阵
  • x x x 是一个 ( n × 1 ) \left(n \times 1 \right) (n×1) 的向量
  • y y y 是一个 ( m × 1 ) \left(m \times 1 \right) (m×1) 的向量

区别与特点：

  • 与 GEMM 不同的是，GEMV 只处理一个矩阵和一个向量，而不是两个矩阵
  • 由于数据量少于 GEMM，因此通常速度更快，但相应地并行化利用程度不如 GEMM

SPMV（Sparse Matrix-Vector Multiplication）

含义：

  • SPMV 即 “稀疏矩阵-向量乘法”，英文全称为 Sparse Matrix-Vector Multiplication
  • 通用数学形式与 GEMV 相似：

y = A × x y = A \times x y=A×x

其中：

  • A A A 是稀疏矩阵（Sparse matrix），也就是说矩阵元素绝大多数是 0
  • x x x 和 y y y 是稠密向量（Dense vector）

区别与特点：

  • 与 GEMV 最大的区别是矩阵 A A A 是 稀疏的，绝大部分元素为 0
  • 计算效率严重依赖于稀疏矩阵的存储方式和数据结构设计，访问内存的模式与 GEMV 有很大的差别
  • 常用于图计算、有限元分析、推荐系统、科学计算等领域

下面的表格对比了它们之间的区别：

名称全称类型使用领域
GEMMGeneral Matrix-Matrix Multiplication稠密矩阵-稠密矩阵乘法科学计算、深度学习训练
SGEMMSingle-precision General Matrix-Matrix Multiplication单精度稠密矩阵-矩阵乘法深度学习训练推理、GPU 计算
GEMVGeneral Matrix-Vector Multiplication稠密矩阵-稠密向量乘法线性代数基础操作
SPMVSparse Matrix-Vector Multiplication稀疏矩阵-稠密向量乘法图计算、推荐系统、科学计算

还有一些别的名词缩写，例如 SGEMV（Single Precision General Matrix-Vector Multiplication）单精度的矩阵-向量乘法，HGEMM（Half Precision General Matrix-Matrix Multiplication）半精度的矩阵-矩阵乘法，HGEMV（Half Precision General Matrix-Vector Multiplication）半精度的矩阵-向量乘法，

我们了解完这些缩写名词的全称后，很容易知道每个字符的含义，其中：

  • ge 代表 general 通用的
  • m 代表 matrix 矩阵
  • mm 代表 matrix-matrix 矩阵与矩阵相乘
  • s 代表 single precision 单精度
  • v 代表 vector 向量
  • mv 代表 matrix-vector 矩阵与向量相乘
  • sp 代表 sparse 稀疏的
  • h 代表 half precision 半精度

知道每个字符的含义后现在我们也能自己推导出来了，例如 sgemv 的全称，其中 s 代表 single precision 单精度，ge 代表 general 通用的，mv 代表 matrix-vector 矩阵与向量乘法，因此它代表的含义就是 “半精度通用矩阵-向量乘法”

OK，下面我们就来看看 SGEMM 单精度通用矩阵乘法在 GPU 上是如何优化的

Note：这里我们的重点不是讨论哪种实现方法更优，而是把每种方法都实现一遍，来更好的理解矩阵乘法。由于我们以学习为目的，因此会做最大程度的简化，例如矩阵 A、B、C 的宽高都保持一致，grid 和 block 的设置都能被整除，不会有过多的边界条件判断。并且我们约定矩阵 A 维度是 M x K，矩阵 B 维度是 K x N，矩阵 C 维度是 M x N

2. 框架搭建和CPU串行实现

我们先来看看 CPU 上的实现，代码如下：

static void cpu_sgemm(float* A, float* B, float* C, const int M, const int N, const int K){for(int m = 0; m < M; m++){for(int n = 0; n < N; n++){float sum = 0;for(int k = 0; k < K; k++){sum += A[m * K + k] * B[k * N + n];}C[m * N + n] = sum;}}
}

CPU 上的串行实现就是通过三个 for 循环来完成的，那这个矩阵乘效率是比较慢的，特别是在一些大矩阵相乘时

我们在 CPU 上的实现主要是为了对比后续在 GPU 上的实现，看看二者矩阵乘结果是否保持一致，如果不一致则说明我们在 GPU 上实现的矩阵乘肯定存在着某些问题，因此我们需要有个对比函数。此外还需要为这些矩阵分配内存空间，包括 host 和 device 上的，整个框架代码如下所示：

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cuda.h>
#include <cuda_runtime.h>// error check#define checkCudaKernel(...)                                                                         \__VA_ARGS__;                                                                                     \do{cudaError_t cudaStatus = cudaPeekAtLastError();                                               \if (cudaStatus != cudaSuccess){                                                                  \printf("launch failed: %s", cudaGetErrorString(cudaStatus));                                 \}} while(0);#define checkCudaRuntime(call) check_runtime(call, #call, __LINE__, __FILE__)bool check_runtime(cudaError_t e, const char* call, int line, const char *file){if (e != cudaSuccess) {printf("CUDA Runtime error %s # %s, code = %s [ %d ] in file %s:%d", call, cudaGetErrorString(e), cudaGetErrorName(e), e, file, line);return false;}return true;
}// utils functionstatic void random_matrix(const int rows, const int cols, float* matrix){for(int i = 0; i < rows; ++i){for(int j = 0; j < cols; ++j){// [-1, 1]matrix[i * cols + j] = 2.0f * (float)drand48() - 1.0f;}}
}static float compare_matrices(const int rows, const int cols, float* matrix_A, float* matrix_B){float max_diff = 0.0f;for(int i = 0; i < rows; ++i){for(int j = 0; j < cols; ++j){float a = matrix_A[i * cols + j];float b = matrix_B[i * cols + j];float diff = fabsf(a - b);if(diff > max_diff){max_diff = diff;}if(diff > 0.5f){printf("Error at (%d, %d): diff = %f, a = %f, b = %f\n", i, j, diff, a, b);}}}return max_diff;
}// cpu implementstatic void cpu_sgemm(float* A, float* B, float* C, const int M, const int N, const int K){for(int m = 0; m < M; m++){for(int n = 0; n < N; n++){float sum = 0;for(int k = 0; k < K; k++){sum += A[m * K + k] * B[k * N + n];}C[m * N + n] = sum;}}
}// gpu implement__global__ void cuda_sgemm_v0_global_memory(float* A, float* B, float* C, const int M, const int N, const int K){// TODO:}int main(){// random seedsrand48(42);const int m = 512;const int n = 512;const int k = 512;const size_t mem_size_A = m * k * sizeof(float);const size_t mem_size_B = k * n * sizeof(float);const size_t mem_size_C = m * n * sizeof(float);float* matrix_A_host = (float*)malloc(mem_size_A);float* matrix_B_host = (float*)malloc(mem_size_B);float* matrix_C_host_cpu_calc = (float*)malloc(mem_size_C);float* matrix_C_host_gpu_calc = (float*)malloc(mem_size_C);random_matrix(m, k, matrix_A_host);random_matrix(k, n, matrix_B_host);memset(matrix_C_host_cpu_calc, 0, mem_size_C);memset(matrix_C_host_gpu_calc, 0, mem_size_C);float* matrix_A_device;float* matrix_B_device;float* matrix_C_device;checkCudaRuntime(cudaMalloc((void**)&matrix_A_device, mem_size_A));checkCudaRuntime(cudaMalloc((void**)&matrix_B_device, mem_size_B));checkCudaRuntime(cudaMalloc((void**)&matrix_C_device, mem_size_C));checkCudaRuntime(cudaMemcpy(matrix_A_device, matrix_A_host, mem_size_A, cudaMemcpyHostToDevice));checkCudaRuntime(cudaMemcpy(matrix_B_device, matrix_B_host, mem_size_B, cudaMemcpyHostToDevice));cpu_sgemm(matrix_A_host, matrix_B_host, matrix_C_host_cpu_calc, m, n, k);constexpr int BLOCK = 16;dim3 block(BLOCK, BLOCK);dim3 grid((m + BLOCK - 1) / BLOCK, (n + BLOCK - 1) / BLOCK);checkCudaKernel(cuda_sgemm_v0_global_memory<<<grid, block>>>(matrix_A_device, matrix_B_device, matrix_C_device, m, n, k));checkCudaRuntime(cudaDeviceSynchronize());checkCudaRuntime(cudaMemcpy(matrix_C_host_gpu_calc, matrix_C_device, mem_size_C, cudaMemcpyDeviceToHost));float max_diff = compare_matrices(m, n, matrix_C_host_cpu_calc, matrix_C_host_gpu_calc);printf("max_diff = %.2f\n", max_diff);free(matrix_A_host);free(matrix_B_host);free(matrix_C_host_cpu_calc);free(matrix_C_host_gpu_calc);checkCudaRuntime(cudaFree(matrix_A_device));checkCudaRuntime(cudaFree(matrix_B_device));checkCudaRuntime(cudaFree(matrix_C_device));return 0;
}

Note：博主自己有一些编码习惯，所以代码实现上与 UP 主的并非完全相同（大部分实现都来自于 ChatGPT），不过优化方式以及原理都一样

整个项目的目录结构如下：

cuda_learn/
├── CMakeLists.txt
└── cuda_sgemm_study├── CMakeLists.txt├── sgemm_v0_global_memory.cu└── ...2 directories, 3 files

cuda_learn 文件夹下的 CMakeLists.txt 内容如下：

cmake_minimum_required(VERSION 3.20.0)
project(cuda_practice VERSION 0.1.0 LANGUAGES CUDA CXX C)
find_package(CUDAToolkit)
add_subdirectory(cuda_sgemm_study)

cuda_sgemm_study 文件夹下的 CMakeLists.txt 内容如下：

set(CMAKE_CUDA_STANDARD 11)add_executable(sgemm_v0_global_memory sgemm_v0_global_memory.cu)
target_link_libraries(sgemm_v0_global_memory PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")target_compile_options(sgemm_v0_global_memory PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
endif()
set_target_properties(sgemm_v0_global_memory PROPERTIES CUDA_SEPARABLE_COMPILATION ON)/* add new */
// add_executable(sgemm_v1_shared_memory sgemm_v1_shared_memory.cu)
// target_link_libraries(sgemm_v1_shared_memory PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
// if(CMAKE_BUILD_TYPE STREQUAL "Debug")
//     target_compile_options(sgemm_v1_shared_memory PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
// endif()
// set_target_properties(sgemm_v1_shared_memory PROPERTIES CUDA_SEPARABLE_COMPILATION ON)

编译运行指令如下：

cd cuda_learn
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Debug && make -j24
./cuda_sgemm_study/sgemm_v0_global_memory

OK，整个框架搭建完成后我们就来看看 GPU 上的 sgemm 是如何实现的

3. baseline算法：global memory

我们先看最基本的利用 GPU 本身的并行计算能力来对 sgemm 加速：

在这里插入图片描述

在 reduce 归约求和中我们使用的是 1D-layout 的布局，但是在 sgemm 矩阵乘法中我们需要使用 2D-layout 的布局，因为矩阵本身就是二维的

我们启动的总线程数是 M x N，恰好是矩阵 C 元素的总和，这意味着每个线程将处理矩阵 C 中的一个元素，也就是完成矩阵 A 中的一行与矩阵 B 中的一列的乘加

图中灰色区域代表着一个 block 中需要完成的计算，在我们的示例中 block 设置的大小是 (16, 16)，也就是它需要完成的是矩阵 A 中的 16 x K 大小的小块矩阵与矩阵 B 中的 K x 16 大小的小块矩阵的乘法

block 中每个 thread 要处理的是 A 小块矩阵中 1 x K 大小的矩阵与 B 小块矩阵中 K x 1 大小的矩阵的乘法，也就是图中淡红色区域展示的部分

实现代码如下：

__global__ void cuda_sgemm_v0_global_memory(float* A, float* B, float* C, const int M, const int N, const int K){// A->MxK B->KxN C->MxNint row = blockIdx.y * blockDim.y + threadIdx.y;int col = blockIdx.x * blockDim.x + threadIdx.x;float sum = 0.0f;for(int k = 0; k < K; k++){sum += A[row * K + k] * B[k * N + col];}C[row * N + col] = sum;
}

我们同时启动 MxN 个线程，每个线程负责计算输出矩阵 C 中的一个元素 C[row][col]。每个线程需要从 A 中读取一行的数据，从 B 中读取一列的数据，计算完之后再写回 C 中，整体计算过程还是比较清晰的

博主绘制了一个草图，简单描述矩阵 C 中某个元素的计算过程：

在这里插入图片描述

在上图中，通过 y 方向和 x 方向的线程索引 row 和 col 可以定位到矩阵 C 中待计算的元素 C[row][col]，接着需要定位到矩阵 A 中的某行元素和矩阵 B 中的某列元素，图中可以很清晰的看到矩阵 A 中 row 行元素的索引是 row * K + k，矩阵 B 中 col 列元素的索引是 k * N + col

其中线程索引 row 和 col 的计算博主绘制了一个草图来说明：

在这里插入图片描述

对于 Block(1, 31) 这个 block 中的线程 Thread(1, 1) 的全局索引 row 和 col 的计算如图中所示，其中 blockIdx.y * blockDim.y 和 blockIdx.x * blockDim.x 定位到具体是哪个 block，threadIdx.y 和 threadIdx.x 定位到具体是哪个 thread

我们可以来简单计算下该核函数对 global memory 的访问次数：

  • 读取次数：
    • 每个线程读取 K 次 A 和 K 次 B，总共有 MxN 个线程
    • 总读取次数 = MxNxK（来自 A）+ MxNxK（来自 B）= 2MNK
  • 写入次数：
    • 每个线程写入 1 次 C，总共有 MxN 个线程
    • 总写入次数 = MxN

可以看到尽管该核函数通过并行计算加速了 sgemm，但其内存访问次数非常高，会产生极高的内存带宽压力

nsight compute 的性能和带宽测试结果如下：

性能和带宽测试结果如下：

优化手段矩阵维度GridBlock耗时(us)Memory Throughout(%)DRAM Throughout(%)
v0_global_memory512x512(32,32)(16,16)471.7896.941.56

Note：测试设备 NVIDIA RTX3060，CUDA-11.6，launch 次数 2000

关于 nsight compute 的简单使用可以参考：【CUDA】Reduce归约求和（下）

4. 优化技巧1：shared memory

要进一步提升 cuda_sgemm_v0_global_memory 的性能，需要优化全局内存访问模式，我们可以考虑利用共享内存：

在这里插入图片描述

相比于之前 global memory 的实现，这里我们会先将 block 中需要处理的矩阵 A 和 B 的 global memory 数据（图中灰色区域部分）移动到 shared memory 上（图中深绿色部分）再做乘加

和 UP 主视频中使用的静态共享内存不同，这里博主在实现时使用的是动态共享内存，定义如下：

bool checkSharedMemorySize(size_t requested_bytes){int device_id = 0;cudaError_t err = cudaGetDevice(&device_id);if(err != cudaSuccess){printf("cudaGetDevice failed: %s\n", cudaGetErrorString(err));return false;}cudaDeviceProp prop{};err = cudaGetDeviceProperties(&prop, device_id);if(err != cudaSuccess){printf("cudaGetDeviceProperties failed: %s\n", cudaGetErrorString(err));return false;        }size_t max_shared_memory = prop.sharedMemPerBlock;if(requested_bytes > max_shared_memory){printf("Error: requested %zu bytes of dynamic shared memory, but device only supports %zu bytes per block\n", requested_bytes, max_shared_memory);return false;        }return true;
}int main(){const int m = 256;  const int n = 256;const int k = 256;// ...constexpr int BLOCK = 16;dim3 block(BLOCK, BLOCK);dim3 grid((m + BLOCK - 1) / BLOCK, (n + BLOCK - 1) / BLOCK);constexpr int BM = BLOCK;constexpr int BN = BLOCK;constexpr int BK = k;size_t shared_mem_size = (BM * BK + BK * BN) * sizeof(float);if(!checkSharedMemorySize(shared_mem_size)){return;} checkCudaKernel(cuda_sgemm_v1_shared_memory<<<grid, block, shared_mem_size>>>(matrix_A_device, matrix_B_device, matrix_C_device, m, n, k));// ...
}

为了简便起见，这里我们会将一个 block 中要处理的元素全部放到 shared memory 上，也就是矩阵 A 的 BM * BK 大小的元素，矩阵 B 的 BN * BK 大小的元素。其中 BM 和 BN 设置为 BLOCK 大小，BK 设置为 k 的大小

那在执行过程中我们其实还需要将矩阵 A、B、C 的维度改小些，前面我们设置的是 M = N = K = 512，这里我们可能就只能使用 256 的大小，不然会出现 nvlink error : Entry function '_Z27cuda_sgemm_v1_shared_memoryPfS_S_iii' uses too much shared data 的错误

这是因为 shared memory 在 GPU 上的大小是有限的，当矩阵的维度较大时，如果还是像上面一样将一个 block 要处理的元素全部放到 shared memory 上去，会超过它的大小的限制。因此，在实际利用 shared memory 实现 sgemm 时我们一般会分块（Tiling）来计算，每次只放一小部分到 shared memory 上，这个是我们下一小节需要实现的

这个小节我们就简单认为 shared memory 足够大，可以将一个 block 要处理的矩阵元素全部放上去

核函数实现代码如下：

__global__ void cuda_sgemm_v1_shared_memory(float* A, float* B, float* C, const int M, const int N, const int K){int row = blockIdx.y * blockDim.y + threadIdx.y;int col = blockIdx.x * blockDim.x + threadIdx.x;// BM = BN = BLOCK = blockDim.x = blockDim.y = 16// BK = K = 256extern __shared__ float shared_mem[];float* A_shared = shared_mem;float* B_shared = shared_mem + blockDim.y * K;for(int k = 0; k < K; k += blockDim.x){// load A_shared[BM, BK]int a_col = k + threadIdx.x;A_shared[threadIdx.y * K + a_col] = A[row * K + a_col];// load B_shared[BK, BN]int b_row = k + threadIdx.y;B_shared[b_row * blockDim.x + threadIdx.x] = B[b_row * N + col];}__syncthreads();float sum = 0.0f;for(int k =0; k < K; k++){sum += A_shared[threadIdx.y * K + k] * B_shared[k * blockDim.x + threadIdx.x];}C[row * N + col] = sum;
}

共享内存被分为两部分：

  • A_shared：缓存 A 的子矩阵，大小为 BM * BK
  • B_shared：缓存 B 的子矩阵，大小为 BK * BN
  • 本示例中 BM = BN = BLOCK = 16，BK = K = 256，一次性将 block 所需的全部行和列读入

接着就是把数据从 global memory 中加载到 shared memory 中，先明确下具体的矩阵大小：

  • 待加载的矩阵 A 的子矩阵大小为 BM * BK 即 16x256
  • 待加载的矩阵 B 的子矩阵大小为 BK * BN 即 256x16
  • 共享内存大小为 BM * BK + BK * BN
  • 一个 block 大小为 16x16 即 256 个线程

其中的索引计算有些繁琐，这里我们简单分析下

首先一个 block 有 256 个线程，而具体子矩阵大小是 16x256 和 256x16，按照一个线程一次加载一个元素来看，可以知道我们需要循环 16 次，也就是一个线程总共要从矩阵 A 和矩阵 B 中分别找到 16 个元素加载到 A_shared 和 B_shared 中

我们先看矩阵 A 子矩阵的加载代码：

for(int k = 0; k < K; k += blockDim.x){int a_col = k + threadIdx.x;A_shared[threadIdx.y * K + a_col] = A[row * K + a_col];
}

首先是有一个循环，循环次数是 K / blockDim.x = 256 / 16 = 16，stride 步长是 blockDim.x = 16

接着需要找到当前线程在矩阵 A 中的全局索引 row * K + a_col，我们可以将它理解为二维数组 A[row][a_col]，那它其实跟我们前面的 cuda_sgemm_v0_global_memory 核函数的索引 A[row * K + k] 是一致的，row * K 定位到具体某一行，a_col 定位到具体某一列，由于我们这里的步长是 blockDim.x，因此列的索引 a_col 是 k + threadIdx.x

博主绘制了一个草图来帮助理解：

在这里插入图片描述

假设我们当前的 block 设置为 blockDim.x = blockDim.y = 2，需要加载的矩阵的维度是 2x8，因此需要循环 4 次，每次加载矩阵 A 中的 4 个元素数据

加载完矩阵 A 的数据后就要写入到 shared memory 中，对应的索引是 threadIdx.y * K + a_col，同理 threadIdx.y * K 定位到 A_shared 中具体某一行，a_col 定位到具体某一列，如下图所示：

在这里插入图片描述

我们再来看矩阵 B 子矩阵的加载代码：

for(int k = 0; k < K; k += blockDim.y){int b_row = k + threadIdx.y;B_shared[b_row * blockDim.x + threadIdx.x] = B[b_row * N + col];
}

首先也是有一个循环，循环次数是 K / blockDim.y = 256 / 16 = 16，stride 步长是 blockDim.y = 16

接着需要找到当前线程在矩阵 B 中的全局索引，b_row * N + col，b_row * N 定位到具体某一行，col 定位到具体某一列，由于我们这里的步长是 blockDim.y，因此行的索引 b_row 是 k * N + col

同样可以根据下面的草图来帮助理解：

在这里插入图片描述

假设我们当前的 block 设置为 blockDim.x = blockDim.y = 2，需要加载的矩阵的维度是 8x2，因此需要循环 4 次，每次加载矩阵 B 中的 4 个元素数据

加载完矩阵 B 的数据后就要写入到 shared memory 中，对应的索引是 b_row * blockDim.x + threadIdx.x，同理 b_row * blockDim.x 定位到 B_shared 中具体某一行，threadIdx.x 定位到具体某一列，如下图所示：

在这里插入图片描述

由于矩阵 A 和矩阵 B 都是方阵，同时 blockDim.x = blockDim.y，因此我们可以把矩阵 A 子矩阵和矩阵 B 子矩阵的加载放在一个循环中，代码如下：

for(int k = 0; k < K; k += blockDim.x){// load A_shared[BM, BK]int a_col = k + threadIdx.x;A_shared[threadIdx.y * K + a_col] = A[row * K + a_col];// load B_shared[BK, BN]int b_row = k + threadIdx.y;B_shared[b_row * blockDim.x + threadIdx.x] = B[b_row * N + col];
}

矩阵 A 和矩阵 B 的子矩阵从 global memory 加载到 shared memory 之后，要做一个线程同步（__syncthreads()），确保所有线程都完成了数据加载操作

数据加载到 shared memory 之后就是做一个乘加最后写回矩阵 C 的对应位置即可：

float sum = 0.0f;
for(int k =0; k < K; k++){sum += A_shared[threadIdx.y * K + k] * B_shared[k * blockDim.x + threadIdx.x];
}C[row * N + col] = sum;

由于我们事先将数据从 global memory 加载到 shared memory 中，因此对 global memory 的访问次数大大减小了：

  • 每个线程块每次需要加载 BM * K 个 A 的元素和 K * BN 个 B 的元素到共享内存
  • 总线程块数为 (M / BM) * (N / BN)
  • 总全局内存读取次数为 (M / BM) * (N / BN) * (BM * BK + BK * BN) = MNK / BM + MNK / BN

性能和带宽测试结果如下：

优化手段矩阵维度GridBlock耗时(us)Memory Throughout(%)DRAM Throughout(%)
v0_global_memory512x512(32,32)(16,16)471.7896.941.56
v1_shared_memory256x256(16,16)(16,16)82.1178.921.84

5. 优化技巧2：shared memory + sliding window

当矩阵维度较大时很容易超过 shared memory 内存大小的限制，这时就需要考虑采用分块（tiling）的方式加载矩阵 A 和 B 的子块：

在这里插入图片描述

在 v1 版本的实现中我们是将 blockDim.x * K 的子矩阵以及 K * blockDim.y 的子矩阵都放在 shared memory 中，但在这里我们只放一小块上去（图中深绿色部分），之后二者相乘将结果存放在寄存器中（图中深蓝色部分），接着不断的滑动窗口直至完成灰色部分的乘加，最后写回到全局内存中

代码如下：

template<unsigned int BLOCK_SIZE>
__global__ void cuda_sgemm_v2_shared_memory_sliding_window(float* A, float* B, float* C, const int M, const int N, const int K){int row = blockIdx.y * blockDim.y + threadIdx.y;int col = blockIdx.x * blockDim.x + threadIdx.x;extern __shared__ float shared_mem[];float* A_tile = shared_mem;float* B_tile = shared_mem + BLOCK_SIZE * BLOCK_SIZE;float sum = 0.0f;for(int k_base = 0; k_base < K; k_base += BLOCK_SIZE){// load A_tile from global memory to shared memoryint a_col = k_base + threadIdx.x;A_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x] = A[row * K + a_col];// load B_tile from global memory to shared memoryint b_row = k_base + threadIdx.y;B_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x] = B[b_row * N + col];__syncthreads();// compute the sum of A_tile * B_tilefor(int i = 0; i < BLOCK_SIZE; ++i){sum += A_tile[threadIdx.y * BLOCK_SIZE + i] * B_tile[i * BLOCK_SIZE + threadIdx.x];}__syncthreads();}C[row * N + col] = sum;
}

Note：其中 BLOCK_SIZE 的设置和 BLOCK 的大小保持一致，统一设置为 16

具体代码分析如下：(from ChatGPT)

1. 全局线程索引的计算

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = blockIdx.x * blockDim.x + threadIdx.x;

当前 grid 和 block 都是 2D-layout 的布局，启动的线程总数恰好是矩阵 C 的元素总数即 512x512。Grid 的大小是 32x32，Block 的大小是 16x16

其中：

  • blockDim.x = blocDim.y = BLOCK_SIZE = 16
  • blockIdx.y ∈ [0, 31]
  • threadIdx.y ∈ [0, 15]

因此 row = blockIdx.x * blockDim.y + threadIdx.x 代表着矩阵 C 全局的第 row 行，同理 col 代表着矩阵 C 的第 col 列，每个线程负责计算出元素 C[row][col]

2. 共享内存布局

extern __shared__ float shared_mem[];
float* A_tile = shared_mem;
float* B_tile = shared_mem + BLOCK_SIZE * BLOCK_SIZE;

这里 A_tile 和 B_tile 各自是一个 BLOCK_SIZE * BLOCK_SIZE（16x16）的行主序数组，如下图所示：

在这里插入图片描述

shared_mem 的总大小是 2 * BLOCK_SIZE * BLOCK_SIZE，一部分存储 A_tile 即矩阵 A 分块的元素，另一部分存储 B_tile 即矩阵 B 分块的元素

3. 沿 K 维度的 “分块累加” 主循环

float sum = 0.0f;
for(int k_base = 0; k_base < K; k_base += BLOCK_SIZE){// load A_tile // load B_tile// __syncthreads();// compute the sum of A_tile * B_tile// __syncthreads();
}

我们把完整的点积

C [ r o w ] [ c o l ] = ∑ k = 0 511 A [ r o w ] [ k ] ∗ B [ k ] [ c o l ] C[row][col] = \sum_{k=0}^{511}A[row][k] * B[k][col] C[row][col]=k=0∑511​A[row][k]∗B[k][col]

划分为 K / BLOCK_SIZE = 32 轮，每轮处理 BLOCK_SIZE = 16 项

3.1 加载 A_tile

int a_col = k_base + threadIdx.x;
A_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x] = A[row * K + a_col];
  • 全局索引 A[row * K + a_col]
    • row * K 跳到矩阵 A 的第 row 行起点
    • + a_col 取该行的第 a_col 列的元素
  • 分块矩阵索引 A_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x]
    • 将 A_tile 行内平铺，代表着第 threadIdx.y 行第 threadIdx.x 列

3.2 加载 B_tile

int b_row = k_base + threadIdx.y;
B_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x] = B[b_row * N + col];
  • 全局索引 B[b_row * N + col]
    • b_row * N 跳到矩阵 B 的第 b_row 行起点
    • + col 取该行的第 col 列元素
  • 分块矩阵索引 B_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x]
    • 将 B_tile 行内平铺，代表着第 threadIdx.y 行第 threadIdx.x 列

3.3 同步 __syncthreads()

__syncthreads();

第一次同步确保所有线程都已经把 A_tile 和 B_tile 从 global memory 中加载到 shared memory 中，确保所有数据加载完成

3.4 分块矩阵计算

for(int i = 0; i < BLOCK_SIZE; ++i){sum += A_tile[threadIdx.y * BLOCK_SIZE + i] * B_tile[i * BLOCK_SIZE + threadIdx.x];
}
  • A_tile 的访问：取 A_tile 的第 threadIdx.y 行，第 i 列
  • B_tile 的访问：取 B_tile 的第 i 行，第 threadIdx.x 列
  • 累加 BLOCK_SIZE 次，完成分块矩阵 A_tile 和 B_tile 的乘积，获得 C[row][col] 的部分点积和

3.5 同步 __syncthreads()

__syncthreads();

第二次同步确保所有线程都完成了当前子块的计算读取，才开始下一轮写新的子块到相同的 shared memory 区域

整个外层循环过程有点像滑动窗口（sliding window），每次只处理一个小的 BLOCK，然后不断地向右滑动，直至处理完所有行或列的数据，如下图所示：

在这里插入图片描述

我们以 Block(0, 0) 为例，这个 Block 块要计算的是 A 和 B 实线部分包围的子矩阵相乘，即 BLOCK_SIZE * K 的矩阵 A 和 K * BLOCK_SIZE 的矩阵 B，由于 shared memory 大小有限，因此我们每次只加载 BLOCK_SIZE * BLOCK_SIZE 大小的矩阵 A_tile 和矩阵 B_tile 相乘计算部分和，即图中灰色小块和蓝色小块部分，接着依次循环所有 K 直至完成 BLOCK_SIZE * K 的矩阵 A 和 K * BLOCK_SIZE 的矩阵 B 相乘

由于启动的 block 大小刚好是 BLOCK_SIZE * BLOCK_SIZE，因此每个 thread 负责从矩阵 A 和矩阵 B 的 global memory 中加载一个元素并放在 shared memory 的 A_tile 和 B_tile 的对应位置上。这样就确保当前 block 从 global memory 中加载整个 A_tile 和 B_tile 放在 shared memory 中来计算，

4. 写回 C

C[row * N + col] = sum;

外层循环做完 32 轮后，sum 就是完整的 ∑ k = 0 K − 1 A [ r o w ] [ k ] ∗ B [ k ] [ c o l ] \sum_{k=0}^{K-1}A[row][k] * B[k][col] ∑k=0K−1​A[row][k]∗B[k][col]，接着写回到全局内存即可

v2 版本通过将 A、B 矩阵按子块加载到高速的 shared memory 中，并在子块内利用这些数据来完成乘加，相比于 v0 版本减少了对全局内存的访问次数、提高了访存合并效率和带宽利用率，在大矩阵运算中可以实现更高的吞吐和性能

性能和带宽测试结果如下：

优化手段矩阵维度GridBlock耗时(us)Memory Throughout(%)DRAM Throughout(%)
v0_global_memory512x512(32,32)(16,16)471.7896.941.56
v1_shared_memory256x256(16,16)(16,16)82.1178.921.84
v2_shared_memory_sliding_window512x512(32,32)(16,16)362.5094.457.05

6. 优化技巧3：增加每个线程的工作量

这个小节的优化策略是增加每个线程的工作量，前面我们学习 reduce 的时候也有提到过，让每个线程多做一些工作，可能会提高计算强度，有助于延迟隐藏

在这里插入图片描述

v2 版本中是一个线程处理矩阵 C 中的一个元素计算，现在我们将 blockDim.x 和 blockDim.y 都减少到原来的 1/2，也就是现在一个线程要处理矩阵 C 中的四个元素计算，如上图所示

代码如下：

template<unsigned int BLOCK_SIZE, unsigned int STRIDE>
__global__ void cuda_sgemm_v3_increase_work_of_per_thread(float* A, float* B, float* C, const int M, const int N, const int K) {constexpr int STEP = BLOCK_SIZE * STRIDE;const int block_row = STEP * blockIdx.y;const int block_col = STEP * blockIdx.x;extern __shared__ float shared_mem[];float* A_shared = shared_mem;float* B_shared = shared_mem + STEP * STEP;float Sum[STRIDE * STRIDE] = {0.0f};for(int k = 0; k < K; k += STEP){// load data from globale memory to shared memoryfor(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){const int tile_row = threadIdx.y + i * BLOCK_SIZE;const int tile_col = threadIdx.x + j * BLOCK_SIZE;A_shared[tile_row * STEP + tile_col] = A[(block_row + tile_row) * K + k + tile_col];B_shared[tile_row * STEP + tile_col] = B[(k + tile_row) * N + block_col + tile_col];}}__syncthreads();// computefor(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){float sum = 0.0f;for(int s = 0; s < STEP; ++s){const int tile_row = threadIdx.y + i * BLOCK_SIZE;const int tile_col = threadIdx.x + j * BLOCK_SIZE;                    sum += A_shared[tile_row * STEP + s] * B_shared[s * STEP + tile_col];}Sum[i * STRIDE + j] += sum;}}__syncthreads();}for(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){const int row = block_row + threadIdx.y + i * BLOCK_SIZE;const int col = block_col + threadIdx.x + j * BLOCK_SIZE;C[row * N + col] = Sum[i * STRIDE + j];}}
}

具体代码分析如下：(from ChatGPT)

1. 模板参数与关键变量解析

  • BLOCK_SIZE：子矩阵分块大小，每次将 BLOCK_SZIE * BLOCK_SIZE 大小的分块矩阵加载到 shared memory 中，当前代码中 BLOCK_SIZE 等于线程块 block 的大小，等于 16
  • STRIDE：每个线程在行和列方向上负责计算的元素数量，每个线程计算 STRIDE * STRIDE 个结果，当前代码中 STRIDE 等于 2
  • STEP：每次迭代中，从 global memory 加载到 shared memory 的矩阵块（Tile）大小，计算方式为 BLOCK_SIZE * STRIDE，等于 32

2. Tile 基址计算（block_row，block_col）

const int block_row = STEP * blockIdx.y;
const int block_col = STEP * blockIdx.x;
  • block_row：当前线程块处理的子矩阵 C 的起始行
  • block_col：当前线程块处理的子矩阵 C 的起始列

和 v2 版本不同的是，这里我们一个 block 要处理之前四个 block 的数据，因此，我们必须要知道当前 block 对应的索引，这样我们才能拿到对应位置的矩阵 A 和矩阵 B 的数据。其索引不再是之前的 blockIdx.y 和 blockIdx.x，而是这里的 STEP * blockIdx.y 以及 STEP * blockIdx.x，每个 block 处理的数据量是 STEP * STEP，也就是 stride 步长为 STEP

3. 分配共享内存

extern __shared__ float shared_mem[];
float* A_shared = shared_mem;
float* B_shared = shared_mem + STEP * STEP;

A_shared 和 B_shared 分别指向共享内存中 A 和 B 的分块数据，大小为 STEP * STEP，当前示例中为 32x32

4. 沿 K 维度的 “分块累加” 主循环

float Sum[STRIDE * STRIDE] = {0.0f};
for(int k = 0; k < K; k += STEP){// load data from global memory to shared memory// __syncthreads()// compute// __syncthreads()
}

我们把完整的点积

C [ r o w ] [ c o l ] = ∑ k = 0 511 A [ r o w ] [ k ] ∗ B [ k ] [ c o l ] C[row][col] = \sum_{k = 0}^{511}A[row][k] * B[k][col] C[row][col]=k=0∑511​A[row][k]∗B[k][col]

划分为 K / STEP = 16 轮，每轮处理 STEP = BLOCK_SIZE * STRIDE = 32 项

那 v3 版本的总体流程其实和前面 v2 版本没有什么区别，主要在内部还有循环，现在一个 block 要同时处理之前 4 个 block 的数据，我们重点就是来看具体的索引是如何计算的

4.1 数据加载

// load data from globale memory to shared memory
for(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){const int tile_row = threadIdx.y + i * BLOCK_SIZE;const int tile_col = threadIdx.x + j * BLOCK_SIZE;A_shared[tile_row * STEP + tile_col] = A[(block_row + tile_row) * K + k + tile_col];B_shared[tile_row * STEP + tile_col] = B[(k + tile_row) * N + block_col + tile_col];}
}
__syncthreads();

每个线程加载 STRIDE * STRIDE 个元素到共享内存，通过 i 和 j 循环覆盖整个 STEP * STEP 的块

博主绘制了一个草图来简单说明数据加载部分的流程，如下图所示：

在这里插入图片描述

在这里插入图片描述

我们启动的总线程数是 16x16，每次只能加载 16x16 大小的矩阵元素到共享内存。因此，如果我们想要每个线程同时处理 STRIDE * STRIDE 个元素，那么就要经历两个循环，分别对应 x 方向和 y 方向

当 i = 0, j = 0 时，我们需要加载矩阵 A 和矩阵 B 对应位置的 16x16 大小的元素到 shared memory 中（图中蓝色区域），每个线程加载一个元素，完成第一个分块的加载。接着依次循环 i 和 j，完成剩余分块的加载，加载完的 A_shared 和 B_shared 大小为 32x32，总共是四个分块。

再回到我们的线程索引的计算中来，首先来看共享内存索引 tile_row 和 tile_col

const int tile_row = threadIdx.y + i * BLOCK_SIZE;
const int tile_col = threadIdx.x + j * BLOCK_SIZE;

tile_row 和 tile_col 主要是用来计算当前线程在共享内存块（STEP * STEP） 中写入的位置

每个线程通过循环 i 和 j 扩展其工作范围，覆盖 STEP * STEP 的分块，i * BLOCK_SIZE 和 j * BLOCK_SIZE 将线程的工作范围从 BLOCK_SIZE * BLOCK_SIZE 扩展到 STEP * STEP

例如线程 (threadIdx.y=0, threadIdx.x=0) 在 i = 0, j = 0 时处理 tile_row = 0, tile_col = 0；在 i = 1, j = 0 时处理 tile_row = 16, tile_col = 0，以此类推

再来看全局内存索引：

矩阵 A 的索引：

A[(block_row + tile_row) * K + k + tile_col]
  • 可以将其理解为二维数组 A[block_row + tile_row][k + tile_col]
  • block_row + tile_row：当前 tile 分块在全局矩阵 A 中的行位置，其中 block_row 定位到具体哪个行 block，tile_row 定位到 block 中具体的哪个行 tile
  • k + tile_col：当前 tile 分块在全局矩阵 A 中的列位置，其中 k 定位到哪个 STEP，tile_col 定位到 STEP 中具体哪一列

矩阵 B 的索引：

B[(k + tile_row) * N + block_col + tile_col]
  • 同样可以将其理解为二维数组 B[k + tile_row][block_col + tile_col]
  • k + tile_row：当前 tile 分块在全局矩阵 B 中的行位置，其中 k 定位到哪个 STEP，tile_row 定位到 STEP 中具体哪一行
  • block_col + tile_col：当前 tile 分块在全局矩阵 B 中的列位置，其中 block_col 定位到具体哪个列 block，tile_col 定位到 block 中具体的哪个列 tile

4. 计算累加

// compute
for(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){float sum = 0.0f;for(int s = 0; s < STEP; ++s){const int tile_row = threadIdx.y + i * BLOCK_SIZE;const int tile_col = threadIdx.x + j * BLOCK_SIZE;                    sum += A_shared[tile_row * STEP + s] * B_shared[s * STEP + tile_col];}Sum[i * STRIDE + j] += sum;}
}
__syncthreads();

4.1 tile_row 和 tile_col 的计算

const int tile_row = threadIdx.y + i * BLOCK_SIZE;  // A_shared 的行索引
const int tile_col = threadIdx.x + j * BLOCK_SIZE;  // B_shared 的列索引

tile_row 和 tile_col 的作用和前面的一样，主要是用来计算当前线程在共享内存块（STEP * STEP） 中的位置

4.2 共享内存的访问：

A_shared[tile_row * STEP + s]  // 读取 A_shared 的 (tile_row, s) 元素
B_shared[s * STEP + tile_col]  // 读取 B_shared 的 (s, tile_col) 元素
  • A_shared 和 B_shared 均为 STEP * STEP 的二维数组
  • A_shared[tile_row][s]：固定 tile_row，遍历所有 s（即 A 分块的一行）
  • B_shared[s][tile_col]：固定 tile_col，遍历所有 s（即 B 分块的一列）

4.3 内积计算

sum += A_shared[tile_row * STEP + s] * B_shared[s * STEP + tile_col];

计算 A 分块的第 tile_row 行与 B 分块的第 tile_col 列的点积，结果累加到 sum

4.4 结果累加

Sum[i * STRIDE + j] += sum;  // 保存到线程局部的累加数组

每个线程维护 STRIDE * STRIDE 个累加结果

5. 写回全局内存

for(int i = 0; i < STRIDE; ++i){for(int j = 0; j < STRIDE; ++j){const int tile_row = threadIdx.y + i * BLOCK_SIZE;const int tile_col = threadIdx.x + j * BLOCK_SIZE;const int row = block_row + tile_row;const int col = block_col + tile_col;C[row * N + col] = Sum[i * STRIDE + j];}
}

全局坐标 row 和 col 的计算：

  • block_row 和 block_col：当前线程块处理的子矩阵的起始行列
  • tile_row 和 tile_col：线程在分块内的局部偏移

最后每个线程将自己计算出来的 STRIDE * STRIDE 个元素写入 global memory

cuda_sgemm_v3_increase_work_of_per_thread 版本相比于 cuda_sgemm_v2_shared_memory_sliding_window 版本引入了一个重要优化：通过 STRIDE 参数让每个线程计算多个输出元素（一个 STRIDE * STRIDE 的小块），从而显著提升了每个线程的计算强度，减少了线程调度的开销

性能和带宽测试结果如下：

优化手段矩阵维度GridBlock耗时(us)Memory Throughout(%)DRAM Throughout(%)
v0_global_memory512x512(32,32)(16,16)471.7896.941.56
v1_shared_memory256x256(16,16)(16,16)82.1178.921.84
v2_shared_memory_sliding_window512x512(32,32)(16,16)362.5094.457.05
v3_increase_work_of_per_thread512x512(16,16)(16,16)204.2684.013.64

7. 优化技巧4：使用float4提高读取效率

这个小节的优化策略是使用 float4 来读取数据，一次性读取连续的四个数据，减少访存，如下图所示：

在这里插入图片描述

在 v3 版本中我们是通过增加每个线程的工作量来同时处理 4 个元素，左图所示，这里我们可以直接使用内置变量 float4 来一次性读取 4 个元素，与 v3 版本不同的是，这里的 4 个元素在内存中是连续的，右图所示

Note：由于我们使用 float4 内置变量读取数据，因此，之前的 block 的布局有所改变，在 v2 版本中每个 block 大小是 (16, 16)，也就是总共 256 个线程，每个 grid 大小是 (32, 32)。在 v4 版本中每个 block 大小是 (4, 16)，也就是 x 方向 blockDim.x = 16 / 4 = 4，不再是原来的 16，这是因为 x 方向我们每个线程使用 float4 一次性读取 4 个数据，grid 大小保持不变

代码如下：

#define FLOAT4(pointer) (reinterpret_cast<float4 *>(&(pointer))[0])template<unsigned int BLOCK_SIZE, unsigned int NUM_PER_THREAD>
__global__ void cuda_sgemm_v4_using_float4(float* A, float* B, float* C, const int M, const int N, const int K) {int row = blockIdx.y * blockDim.y + threadIdx.y;int col = (blockIdx.x * blockDim.x + threadIdx.x) * NUM_PER_THREAD;extern __shared__ float shared_mem[];float* A_tile = shared_mem;float* B_tile = shared_mem + BLOCK_SIZE * BLOCK_SIZE;float sum[NUM_PER_THREAD] = {0.0f};for(int k_base = 0; k_base < K; k_base += BLOCK_SIZE){// load A_tile from global memory to shared memoryint a_col = k_base + threadIdx.x * NUM_PER_THREAD;FLOAT4(A_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD]) = FLOAT4(A[row * K + a_col]);// load B_tile from global memory to shared memoryint b_row = k_base + threadIdx.y;FLOAT4(B_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD]) = FLOAT4(B[b_row * N + col]);__syncthreads();// compute the sum of A_tile * B_tilefor(int i = 0; i < NUM_PER_THREAD; ++i){for(int j = 0; j < BLOCK_SIZE; ++j){sum[i] += A_tile[threadIdx.y * BLOCK_SIZE + j] * B_tile[j * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD + i];}}__syncthreads();        }for(int i = 0; i < NUM_PER_THREAD; ++i){C[row * N + col + i] = sum[i];}
}

具体代码分析如下：(from ChatGPT)

1. 函数签名和宏

#define FLOAT4(pointer) (reinterpret_cast<float4 *>(&(pointer))[0])

这个宏把一个 float* 指针 reinterpret 成 float4*，然后解引用一次，得到一个 float4 类型变量，表示连续的 4 个 float

template<unsigned int BLOCK_SIZE, unsigned int NUM_PER_THREAD>
__global__ void cuda_sgemm_v4_using_float4(float* A, float* B, float* C, const int M, const int N, const int K)
  • 模板参数 BLOCK_SIZE：每个 tile 的大小（示例中是 16），也就是每次从矩阵 A 的子矩阵和矩阵 B 的子矩阵处理 BLOCK_SIZE * BLOCK_SIZE 大小的元素从 global memory 加载到 shared memory，和 v2 版本保持一致
  • 模板参数 NUM_PER_THREAD：每个线程负责多少个输出（默认是 4）

2. 索引计算：定位当前线程处理的矩阵 C 的全局位置

int row = blockIdx.y * blockDim.y + threadIdx.y;
int col = (blockIdx.x * blockDim.x + threadIdx.x) * NUM_PER_THREAD;
  • row 行的索引和 v2 版本的没有区别
  • 注意这里的 col 列索引的写法 (...) * NUM_PER_THREAD 非常关键
    • 每个线程现在需要处理 4 个列方向上的元素，因此要乘以 4，也就是列与列之间 stride 步长是 4，不再是原来的 1

3. 声明共享内存并划分 A_tile 和 B_tile

extern __shared__ float shared_mem[];
float* A_tile = shared_mem;
float* B_tile = shared_mem + BLOCK_SIZE * BLOCK_SIZE;
  • shared_mem 是传入 kernel 的动态共享内存，大小是 2 * BLOCK_SZIE * BLOCK_SIZE
  • 一半用来存 A_tile，另一个用来存 B_tile，和 v2 版本保持一致

4. 准备累加数组 sum

float sum[NUM_PER_THREAD] = {0.0f};
  • 每个线程要累加出 NUM_PER_THREAD 个输出元素

5. 核心主循环：tile-wise 乘法

for(int k_base = 0; k_base < K; k_base += BLOCK_SIZE){
  • 滑动窗口，从左到右（沿着 K 维方向）遍历 A 子矩阵和 B 子矩阵
  • 每轮计算一个 BLOCK_SIZE * BLOCK_SIZE 的子块（tile）

5.1 将 A 子块加载到共享内存（float4 加载）

int a_col = k_base + threadIdx.x * NUM_PER_THREAD;
FLOAT4(A_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD]) = FLOAT4(A[row * K + a_col]);
  • 矩阵 A 的索引可以理解为二维数组 A[row][a_col]，row * K 定位到具体哪一行，a_col 定位到具体哪一列，和之前 v2 版本保持一致，不过差别就是 a_col 中 threadIdx.x 需要乘上 NUM_PER_THREAD，代表着每个线程加载 A 中的 4 个元素数据
  • 共享内存中的 A_tile 也可以理解为二维数组 A[threadIdx.y][threadIdx.x * NUM_PER_THREAD]，threadIdx.y * BLOCK_SIZE 定位到具体哪一行，threadIdx.x * NUM_PER_THREAD 定位到具体哪一列，和之前有差异的就是 threadIdx.x 需要乘以 NUM_PER_THREAD
  • 每个线程负责从矩阵 A 拷贝 4 个连续的 float 到 shared memory 中

博主绘制了一个草图来简单说明和 v2 版本的差异，如下图所示：

在这里插入图片描述

在 v2 版本中我们启动的总线程数和 C 矩阵的元素总数是一样的，每个线程处理 C 中的一个元素，包括虚线部分。在当前版本中我们启动的总线程数是原来的 1/4，因为每个线程采用 float4 来同时处理 4 个元素

对于第一个线程（蓝色小块）而言，它会计算 C 矩阵前 4 个结果，主要是计算矩阵 A 中的一行和矩阵 B 中的四列，对于第二个线程（绿色小块）也是类似，它也是计算 C 矩阵的 4 个元素结果

Note：这里博主未展示逐 tile 块加载到 shared memory 中计算，不过你应该清楚这一点

5.2 将 B 子块加载到共享内存（float4 加载）

int b_row = k_base + threadIdx.y;
FLOAT4(B_tile[threadIdx.y * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD]) = FLOAT4(B[b_row * N + col]);
  • 矩阵 B 的索引计算同样可以理解为二维数组 B[b_row][col]，b_row * N 定位到具体哪一行，col 定位到具体哪一列
  • 共享内存中的 B_tile 索引和 A_tile 保持一致

5.3 同步线程，确保 tile 加载完成

__syncthreads();

5.4 实际乘法：A_tile * B_tile，累加进 sum

for(int i = 0; i < NUM_PER_THREAD; ++i){for(int j = 0; j < BLOCK_SIZE; ++j){sum[i] += A_tile[threadIdx.y * BLOCK_SIZE + j] *B_tile[j * BLOCK_SIZE + threadIdx.x * NUM_PER_THREAD + i];}
}
  • 外层 i 是当前线程正在计算的第 i 个输出元素（对应 C[row][col + i]）
  • 内层 j 遍历 k 方向的 BLOCK_SIZE
  • A_tile 的索引可以理解为二维数组 A_tile[threadIdx.y][j]，threadIdx.y * BLOCK_SIZE 定义到具体哪一行，j 则定位这一行的具体哪个元素
  • B_tile 的索引需要加上 NUM_PER_THREAD 的偏移量，它也可以理解为二维数组 B_tile[j][threadIdx.x * NUM_PER_THREAD + i]，j 定位到具体哪一行，threadIdx.x * NUM_PER_THREAD 定位到具体是哪个 threadIdx.x 负责的 float4 变量，+ i 定位到 float4 变量中的具体哪个 float

博主绘制了一个草图来理解这一过程，如下图所示：

在这里插入图片描述

在图中，BLOCK_SIZE 设置为 8，block 的大小为 (2, 8)，i 对应着外层循环，内层循环则是访问 A_tile 的一行和 B_tile 的一列，图中未展示。我们可以清晰的看到一个线程处理 B_tile 的四列，也就是一个线程处理 B_tile 中的 float4 即 4 个元素

5.5 再次同步线程，为下一轮 tile 做准备

__syncthreads();

6. 写回输出 C

for(int i = 0; i < NUM_PER_THREAD; ++i){C[row * N + col + i] = sum[i];
}
  • 每个线程写回自己负责的 4 个输出元素

v4 版本相比于 v2 版本访问效率提高了，主要体现在在：

  • 使用 float4 同时加载 4 个 float：提高 global memory 带宽利用率
  • 每个线程计算 NUM_PER_THREAD 个输出值：减少线程数量，提高 cache hit 率

性能和带宽测试结果如下：

优化手段矩阵维度GridBlock耗时(us)Memory Throughout(%)DRAM Throughout(%)
v0_global_memory512x512(32,32)(16,16)471.7896.941.56
v1_shared_memory256x256(16,16)(16,16)82.1178.921.84
v2_shared_memory_sliding_window512x512(32,32)(16,16)362.5094.457.05
v3_increase_work_of_per_thread512x512(16,16)(16,16)204.2684.013.64
v4_using_float4512x512(32,32)(4,16)209.6091.993.44

结语

这篇文章我们主要从三个方面出发优化 sgemm，首先 baseline 使用的 global memory 对内存的访问非常高，因此优化技巧 1 和 2 考虑使用 shared memory 减少对全局内存的访问次数，提高带宽利用率。优化技巧 3 通过增加每个线程的工作量来提高线程处理的强度，减少线程调度的开销。同样优化技巧 4 也是减少线程数量，提高线程利用率，不过它是通过内置变量 float4 来同时处理 4 个元素

OK，这就是这篇文章关于 sgemm 优化的全部内容了

篇幅原因，后续的 sgemm 优化我们放在下篇文章中讲解，敬请期待🤗

下载链接

  • Sgemm 矩阵乘法代码下载链接【提取码：1234】

参考

  • 【CUDA】Sgemm单精度矩阵乘法（已完结~）
  • https://github.com/tpoisonooo/how-to-optimize-gemm/cuda
  • 深入浅出GPU优化系列：GEMM优化（一）
  • [施工中] CUDA GEMM 理论性能分析与 kernel 优化
  • cuda 入门的正确姿势：how-to-optimize-gemm
  • CUDA 矩阵乘法终极优化指南
  • CUDA实现矩阵乘法的8种优化策略编程介绍
  • https://chatgpt.com/

相关文章：

  • 达梦数据库 【-6111: 字符串转换出错】问题处理
  • 【AI大模型】赋能【传统业务】
  • React构建组件
  • 微信小程序学习之轮播图swiper
  • 【unity游戏开发——编辑器扩展】EditorWindow自定义unity窗口拓展
  • 橙子、橘子相关（果实、叶片、疾病等）数据集大合集
  • SQL注入报错“Illegal mix of collations for operation ‘UNION‘”解决办法
  • 材料×工艺×AI：猎板PCB重构汽车电子四层板技术逻辑
  • [滑动窗口]越短越合法（可转化成越长越合法）
  • docker-compose的使用总结
  • Linux下的c/c++开发之操作Redis数据库
  • select、poll、epoll
  • MySQL库级管理：数据库管理与存储引擎剖析
  • kafka connect 大概了解
  • idea挂掉，会导致进程不结束，切换profile环境，导致token认证不通过
  • Linux Bash | Capture Output / Recall
  • Android Studio Meerkat与Gradle构建工具升级实战指南
  • 同设备访问php的多个接口会有先后等待问题
  • 电机的导程和脉冲之间的关系
  • Rust入门之高级Trait
  • 复原展出孙吴大墓，江苏首座考古博物馆将开放
  • 巴菲特谈卸任CEO：开始偶尔失去平衡，但仍然保持敏锐的头脑，仍打算继续工作
  • 绿景中国地产：洛杉矶酒店出售事项未能及时披露纯属疏忽，已采取补救措施
  • “异常”只停留在医院里，用艺术为“泡泡宝贝”加油
  • KPL“王朝”诞生背后：AG和联赛一起迈向成熟
  • 北斗专访｜特赞科技范凌：现在AI主要是“说话”，接下来要“干活”了