博客来源：https://leimao.github.io/article/CUDA-Matrix-Multiplication-Optimization/ ，来自Lei Mao，已获得作者转载授权。后续会转载一些Lei Mao的CUDA相关Blog，也是一个完整的专栏，Blog会从稍早一些的CUDA架构到当前最新的CUDA架构，也会包含实用工程技巧，底层指令分析，Cutlass分析等等多个课题，是一个时间线十分明确的专栏。

CUDA 矩阵乘法优化

介绍

通用矩阵乘法（GEMM）是线性代数中的基本运算。它也是许多科学计算应用中非常重要的运算，如机器学习和深度学习。

在本文中，我们将讨论如何使用 CUDA 优化 NVIDIA GPU 上 FP32 GEMM 的性能，以及如何使用 NVIDIA Tensor Core 将 FP32 GEMM 优化扩展到 FP16 GEMM。

通用矩阵乘法

GEMM 运算计算 ，其中 ，，，。在计算机程序中，通常  和  是常量输入矩阵， 将被输出矩阵  覆盖。

在我们的实现中，我们假设所有矩阵 、、 和  都以行主序存储在内存中，对于 FP32 矩阵，主维度填充为 64 字节，对于 FP16 矩阵，主维度填充为 32 字节。

具有非合并内存访问的朴素实现

朴素实现是使用 2D 块，其中每个线程负责计算输出矩阵的一个元素。具体来说，对于全局线程索引为  的每个线程，其中  和 ，它计算

以下代码片段显示了朴素实现。

template <typename T>
__global__ void gemm_v00(size_t m, size_t n, size_t k, T alpha, T const* A,
                         size_t lda, T const* B, size_t ldb, T beta, T* C,
                         size_t ldc)
{
    // Compute the row and column of C that this thread is responsible for.
    size_tconst C_row_idx{blockIdx.x * blockDim.x + threadIdx.x};
    size_tconst C_col_idx{blockIdx.y * blockDim.y + threadIdx.y};

    // Each thread compute
    // C[C_row_idx, C_col_idx] = alpha * A[C_row_idx, :] * B[:, C_col_idx] +
    // beta * C[C_row_idx, C_col_idx].
    if (C_row_idx < m && C_col_idx < n)
    {
        T sum{static_cast<T>(0)};
        for (size_t k_idx{0U}; k_idx < k; ++k_idx)
        {
            sum += A[C_row_idx * lda + k_idx] * B[k_idx * ldb + C_col_idx];
        }
        C[C_row_idx * ldc + C_col_idx] =
            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
    }
}

template <typename T>
void launch_gemm_kernel_v00(size_t m, size_t n, size_t k, T const* alpha,
                            T const* A, size_t lda, T const* B, size_t ldb,
                            T const* beta, T* C, size_t ldc,
                            cudaStream_t stream)
{
    dim3 const block_dim{32U, 32U, 1U};
    dim3 const grid_dim{
        (static_cast<unsignedint>(m) + block_dim.x - 1U) / block_dim.x,
        (static_cast<unsignedint>(n) + block_dim.y - 1U) / block_dim.y, 1U};
    gemm_v00<T><<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B,
                                                     ldb, *beta, C, ldc);
    CHECK_LAST_CUDA_ERROR();
}

然而，除了朴素算法的其他缺点之外，这个实现还有一个主要问题，即读取和写入全局内存时的非合并内存访问。在我们的实现中，由于疏忽使用快速线程索引来索引矩阵  和  的行，同一 warp 中的线程从以行主序存储在内存中的矩阵  的同一列读取元素，由于读取完全不连续，导致非合并内存访问。当 warp 覆写矩阵  的元素时也会发生同样的问题。同一 warp 中的线程读取矩阵  的同一元素，导致广播内存访问，这不受疏忽的影响。

这个 FP32 GEMM 实现的性能在 NVIDIA GeForce RTX 3090 GPU 上仅为 0.27 TFLOPS，性能非常差。

具有合并内存访问的朴素实现

解决非合并内存访问的方法是使用快速线程索引来索引以行主序存储在内存中的矩阵的行，这样同一 warp 中的线程读取或覆写矩阵同一行的元素时是合并的。在我们的实现中，我们只需要在内核函数中交换快速线程索引和慢速线程索引。

以下代码片段显示了具有合并内存访问的朴素实现。

template <typename T>
__global__ void gemm_v01(size_t m, size_t n, size_t k, T alpha, T const* A,
                         size_t lda, T const* B, size_t ldb, T beta, T* C,
                         size_t ldc)
{
    // Compute the row and column of C that this thread is responsible for.
    size_tconst C_col_idx{blockIdx.x * blockDim.x + threadIdx.x};
    size_tconst C_row_idx{blockIdx.y * blockDim.y + threadIdx.y};

    // Each thread compute
    // C[C_row_idx, C_col_idx] = alpha * A[C_row_idx, :] * B[:, C_col_idx] +
    // beta * C[C_row_idx, C_col_idx].
    if (C_row_idx < m && C_col_idx < n)
    {
        T sum{static_cast<T>(0)};
        for (size_t k_idx{0U}; k_idx < k; ++k_idx)
        {
            sum += A[C_row_idx * lda + k_idx] * B[k_idx * ldb + C_col_idx];
        }
        C[C_row_idx * ldc + C_col_idx] =
            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
    }
}

template <typename T>
void launch_gemm_kernel_v01(size_t m, size_t n, size_t k, T const* alpha,
                            T const* A, size_t lda, T const* B, size_t ldb,
                            T const* beta, T* C, size_t ldc,
                            cudaStream_t stream)
{
    dim3 const block_dim{32U, 32U, 1U};
    dim3 const grid_dim{
        (static_cast<unsignedint>(n) + block_dim.x - 1U) / block_dim.x,
        (static_cast<unsignedint>(m) + block_dim.y - 1U) / block_dim.y, 1U};
    gemm_v01<T><<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B,
                                                     ldb, *beta, C, ldc);
    CHECK_LAST_CUDA_ERROR();
}

现在，由于这个修复，同一warp中的线程从以行主序存储在内存中的矩阵的同一行读取元素，导致合并内存访问。当warp覆写矩阵的元素时也会发生同样的情况。同一warp中的线程读取矩阵的同一元素，导致广播内存访问。因此，这个实现应该比非合并内存访问的实现性能好得多。

这个FP32 GEMM实现的性能在NVIDIA GeForce RTX 3090 GPU上变成了1.72 TFLOPS，比之前的实现好得多。然而，考虑到GPU的理论峰值性能是35.58 TFLOPS，这个实现的性能仍然很差。

具有2D块分块的实现

由于前面的实现频繁访问全局内存，GEMM实现变成了内存受限的。由于访问共享内存比访问全局内存快得多，为了提高性能，我们可以使用共享内存来缓存输入矩阵和以实现数据重用。

然而，由于共享内存大小有限，我们无法在共享内存中缓存整个输入矩阵和。相反，我们可以在共享内存中缓存和的2D分块，并使用2D分块来计算输出矩阵的2D分块。然后，我们可以将和的下一个2D分块加载到共享内存中，并计算的下一个2D分块。

数学上，给定一个GEMM运算，其中，，，，矩阵可以被分成更小的矩阵。

每个中的小矩阵都是通过多个小矩阵乘法和累加计算的。

在这个实现中，每个具有块索引的2D块，其中和，负责计算一个小矩阵。共享内存用于缓存大小分别为和的和的2D分块。的2D分块由索引，其中和。的2D分块由索引，其中和。缓存和小矩阵乘法计算过程重复次，直到整个小矩阵被累加完成。

与之前的实现类似，每个块需要个线程来计算小矩阵，每个具有块线程索引的线程，其中和，负责计算小矩阵的一个元素。

下面的代码片段展示了使用2D块分块的实现。

template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
          size_t BLOCK_TILE_SIZE_K, size_t NUM_THREADS, size_t BLOCK_TILE_SKEW_SIZE_X = 0U, size_t BLOCK_TILE_SKEW_SIZE_K = 0U>
__device__ void load_data_to_shared_memory(T const* A, size_t lda,
                                           T const* B, size_t ldb,
                                           T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K + BLOCK_TILE_SKEW_SIZE_K],
                                           T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X + BLOCK_TILE_SKEW_SIZE_X],
                                           size_t thread_block_tile_idx,
                                           size_t thread_linear_idx,
                                           size_t m, size_t n,
                                           size_t k)
{
    // 将DRAM中的矩阵A数据加载到共享内存中的A_thread_block_tile
#pragma unroll
    for (size_t load_idx{0U};
         load_idx <
         (BLOCK_TILE_SIZE_Y * BLOCK_TILE_SIZE_K + NUM_THREADS - 1U) /
             NUM_THREADS;
         ++load_idx)
    {
        // 计算在共享内存分块中的行列索引
        size_tconst A_thread_block_tile_row_idx{
            (thread_linear_idx + load_idx * NUM_THREADS) /
            BLOCK_TILE_SIZE_K};
        size_tconst A_thread_block_tile_col_idx{
            (thread_linear_idx + load_idx * NUM_THREADS) %
            BLOCK_TILE_SIZE_K};
        // 计算在全局矩阵A中的行列索引
        size_tconst A_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
                               A_thread_block_tile_row_idx};
        size_tconst A_col_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
                               A_thread_block_tile_col_idx};

        // 边界检查可能会在一定程度上降低内核性能
        // 但是能够保证内核对所有不同GEMM配置的正确性
        T val{static_cast<T>(0)};
        if (A_row_idx < m && A_col_idx < k)
        {
            val = A[A_row_idx * lda + A_col_idx];
        }
        // 这个if语句会降低内核性能
        // 在主机代码中添加静态断言来保证这个if条件总是为true
        static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS ==
                      0U);
        // if (A_thread_block_tile_row_idx < BLOCK_TILE_SIZE_Y &&
        //     A_thread_block_tile_col_idx < BLOCK_TILE_SIZE_K)
        // {
        //     A_thread_block_tile[A_thread_block_tile_row_idx]
        //                        [A_thread_block_tile_col_idx] = val;
        // }
        A_thread_block_tile[A_thread_block_tile_row_idx]
                           [A_thread_block_tile_col_idx] = val;
    }
// 将DRAM中的矩阵B数据加载到共享内存中的B_thread_block_tile
#pragma unroll
    for (size_t load_idx{0U};
         load_idx <
         (BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_X + NUM_THREADS - 1U) /
             NUM_THREADS;
         ++load_idx)
    {
        // 计算在共享内存分块中的行列索引
        size_tconst B_thread_block_tile_row_idx{
            (thread_linear_idx + load_idx * NUM_THREADS) /
            BLOCK_TILE_SIZE_X};
        size_tconst B_thread_block_tile_col_idx{
            (thread_linear_idx + load_idx * NUM_THREADS) %
            BLOCK_TILE_SIZE_X};
        // 计算在全局矩阵B中的行列索引
        size_tconst B_row_idx{thread_block_tile_idx * BLOCK_TILE_SIZE_K +
                               B_thread_block_tile_row_idx};
        size_tconst B_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
                               B_thread_block_tile_col_idx};

        // 边界检查可能会在一定程度上降低内核性能
        // 但是能够保证内核对所有不同GEMM配置的正确性
        T val{static_cast<T>(0)};
        if (B_row_idx < k && B_col_idx < n)
        {
            val = B[B_row_idx * ldb + B_col_idx];
        }
        // 这个if语句会降低内核性能
        // 在主机代码中添加静态断言来保证这个if条件总是为true
        static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS ==
                      0U);
        // if (B_thread_block_tile_row_idx < BLOCK_TILE_SIZE_K &&
        //     B_thread_block_tile_col_idx < BLOCK_TILE_SIZE_X)
        // {
        //     B_thread_block_tile[B_thread_block_tile_row_idx]
        //                        [B_thread_block_tile_col_idx] = val;
        // }
        B_thread_block_tile[B_thread_block_tile_row_idx]
                           [B_thread_block_tile_col_idx] = val;
    }
}

template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
          size_t BLOCK_TILE_SIZE_K>
__global__ void gemm_v02(size_t m, size_t n, size_t k, T alpha, T const* A,
                         size_t lda, T const* B, size_t ldb, T beta, T* C,
                         size_t ldc)
{
    // 避免使用blockDim.x * blockDim.y作为每个块的线程数
    // 因为它是一个运行时常数，编译器无法基于它优化循环展开
    // 改为使用编译时常数
    constexprsize_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
    size_tconst thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};

    // 计算此线程负责的矩阵C的行列索引
    size_tconst C_col_idx{blockIdx.x * blockDim.x + threadIdx.x};
    size_tconst C_row_idx{blockIdx.y * blockDim.y + threadIdx.y};

    // 在共享内存中缓存A和B的分块以实现数据重用
    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];

    // 计算需要处理的线程块分块数量
    size_tconst num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
                                        BLOCK_TILE_SIZE_K};

    // 累加和初始化为0
    T sum{static_cast<T>(0)};
    // 遍历每个线程块分块
    for (size_t thread_block_tile_idx{0U};
         thread_block_tile_idx < num_thread_block_tiles;
         ++thread_block_tile_idx)
    {
        // 加载当前分块的数据到共享内存
        load_data_to_shared_memory<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
                                   BLOCK_TILE_SIZE_K, NUM_THREADS>(
            A, lda, B, ldb, A_thread_block_tile, B_thread_block_tile,
            thread_block_tile_idx, thread_linear_idx, m, n, k);
        __syncthreads();

        // 计算当前分块的矩阵乘法
#pragma unroll
        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
        {
            // 这样做会导致性能达到2 TOPS
            // 假设blockDim.x = blockDim.y = 32
            // 实际上，对于一个warp，在一次迭代中，我们从A_thread_block_tile的
            // 共享内存同一位置读取值，导致广播，我们也从B_thread_block_tile
            // 读取32个没有bank冲突的值。即使如此，所有值都必须从共享内存读取，
            // 结果是共享内存指令运行得非常密集，只是为了使用简单的算术指令
            // 计算少量值，这是不高效的
            sum += A_thread_block_tile[threadIdx.y][k_i] *
                   B_thread_block_tile[k_i][threadIdx.x];
        }
        __syncthreads();
    }
    // 将最终结果写入输出矩阵C（边界检查）
    if (C_row_idx < m && C_col_idx < n)
    {
        C[C_row_idx * ldc + C_col_idx] =
            alpha * sum + beta * C[C_row_idx * ldc + C_col_idx];
    }
}

template <typename T>
void launch_gemm_kernel_v02(size_t m, size_t n, size_t k, T const* alpha,
                            T const* A, size_t lda, T const* B, size_t ldb,
                            T const* beta, T* C, size_t ldc,
                            cudaStream_t stream)
{
    // 可以随意调整块分块大小
    // 算法正确性应该始终得到保证
    constexprunsignedint BLOCK_TILE_SIZE_X{32U};
    constexprunsignedint BLOCK_TILE_SIZE_Y{32U};
    constexprunsignedint BLOCK_TILE_SIZE_K{32U};
    constexprunsignedint NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y};
    // 静态断言确保分块大小与线程数兼容
    static_assert(BLOCK_TILE_SIZE_K * BLOCK_TILE_SIZE_Y % NUM_THREADS == 0U);
    static_assert(BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_K % NUM_THREADS == 0U);
    // 设置块维度
    dim3 const block_dim{BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, 1U};
    // 设置网格维度
    dim3 const grid_dim{
        (static_cast<unsignedint>(n) + block_dim.x - 1U) / block_dim.x,
        (static_cast<unsignedint>(m) + block_dim.y - 1U) / block_dim.y, 1U};
    // 启动GEMM内核
    gemm_v02<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K>
        <<<grid_dim, block_dim, 0U, stream>>>(m, n, k, *alpha, A, lda, B, ldb,
                                              *beta, C, ldc);
    CHECK_LAST_CUDA_ERROR();
}

这个FP32 GEMM实现的性能在NVIDIA GeForce RTX 3090 GPU上变成了2.66 TFLOPS，比之前的实现好得多。然而，它仍然远远低于GPU的理论峰值性能。

这个实现的问题是共享内存被频繁访问。即使访问共享内存比访问全局内存快得多，共享内存指令运行得非常密集，只是为了使用简单的算术指令计算少量值，这是不高效的。因此，这个实现的性能仍然受限于内存带宽，这次来自共享内存。

具有2D块分块和1D线程分块的实现

为了进一步提高性能，我们可以通过进一步缓存输入矩阵和的更小的分块到线程的寄存器中来缓解共享内存带宽问题。这次，每个线程负责计算输出矩阵的一个小分块，而不是一个单个元素。因为寄存器是最快的访问方式，这个实现的性能应该比之前的实现好得多。

我们首先只缓存矩阵的数据从共享内存到寄存器。每个具有块线程索引的线程，其中和，现在负责计算小矩阵的个元素，其中是线程分块大小。

在我们之前的实现中没有线程分块，为了计算小矩阵的一个元素，我们需要从共享内存中缓存的矩阵中读取个值，从共享内存中缓存的矩阵中读取个值。总共需要从共享内存中读取个值。

现在，有了1D线程分块，为了计算小矩阵的个元素，我们只需要从共享内存中缓存的矩阵中读取个值，从共享内存中缓存的矩阵中读取个值。具体来说，在每个内层循环中，被缓存到寄存器中以被重复使用次。总共需要从共享内存中读取个值。平均来说，为了计算小矩阵的一个元素，我们需要从共享内存中读取个值。

因为 ，共享内存访问频率降低，共享内存带宽问题得到缓解。

下面的代码片段展示了使用2D块分块和1D线程分块的实现。

// 使用2D块分块和1D线程分块的GEMM内核模板
// T: 数据类型
// BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K: 块分块大小
// THREAD_TILE_SIZE_Y: 线程分块大小（每个线程处理的行数）
template <typename T, size_t BLOCK_TILE_SIZE_X, size_t BLOCK_TILE_SIZE_Y,
          size_t BLOCK_TILE_SIZE_K, size_t THREAD_TILE_SIZE_Y>
__global__ void gemm_v03(size_t m, size_t n, size_t k, T alpha, T const* A,
                         size_t lda, T const* B, size_t ldb, T beta, T* C,
                         size_t ldc)
{
    // 避免使用blockDim.x * blockDim.y作为每个块的线程数
    // 因为它是运行时常量，编译器无法基于此优化循环展开
    // 使用编译时常量代替
    constexprsize_t NUM_THREADS{BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y /
                                 THREAD_TILE_SIZE_Y};
    // 计算线程的线性索引
    size_tconst thread_linear_idx{threadIdx.y * blockDim.x + threadIdx.x};

    // 在共享内存中缓存A和B的分块以实现数据重用
    __shared__ T A_thread_block_tile[BLOCK_TILE_SIZE_Y][BLOCK_TILE_SIZE_K];
    __shared__ T B_thread_block_tile[BLOCK_TILE_SIZE_K][BLOCK_TILE_SIZE_X];

    // 计算需要处理的线程块分块数量
    size_tconst num_thread_block_tiles{(k + BLOCK_TILE_SIZE_K - 1) /
                                        BLOCK_TILE_SIZE_K};

    // 块中的每个线程处理THREAD_TILE_SIZE_Y个输出值
    // 具体来说，这些值对应于
    // C[blockIdx.y * BLOCK_TILE_SIZE_Y + threadIdx.x / BLOCK_TILE_SIZE_X *
    // THREAD_TILE_SIZE_Y : blockIdx.y * BLOCK_TILE_SIZE_Y + (threadIdx.x /
    // BLOCK_TILE_SIZE_X + 1) * THREAD_TILE_SIZE_Y][blockIdx.x *
    // BLOCK_TILE_SIZE_X + threadIdx.x % BLOCK_TILE_SIZE_X]
    T C_thread_results[THREAD_TILE_SIZE_Y] = {static_cast<T>(0)};

    // 遍历所有线程块分块
    for (size_t thread_block_tile_idx{0U};
         thread_block_tile_idx < num_thread_block_tiles;
         ++thread_block_tile_idx)
    {
        // 将数据从全局内存加载到共享内存
        load_data_to_shared_memory<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y,
                                   BLOCK_TILE_SIZE_K, NUM_THREADS>(
            A, lda, B, ldb, A_thread_block_tile, B_thread_block_tile,
            thread_block_tile_idx, thread_linear_idx, m, n, k);
        // 同步所有线程，确保共享内存数据加载完成
        __syncthreads();

        // 遍历K维度进行矩阵乘法计算
#pragma unroll
        for (size_t k_i{0U}; k_i < BLOCK_TILE_SIZE_K; ++k_i)
        {
            size_tconst B_thread_block_tile_row_idx{k_i};
            // B_val被缓存在寄存器中以减轻共享内存访问压力
            T const B_val{
                B_thread_block_tile[B_thread_block_tile_row_idx]
                                   [thread_linear_idx % BLOCK_TILE_SIZE_X]};
            // 遍历线程分块的每一行
#pragma unroll
            for (size_t thread_tile_row_idx{0U};
                 thread_tile_row_idx < THREAD_TILE_SIZE_Y;
                 ++thread_tile_row_idx)
            {
                // 计算A矩阵在共享内存中的行索引
                size_tconst A_thread_block_tile_row_idx{
                    thread_linear_idx / BLOCK_TILE_SIZE_X * THREAD_TILE_SIZE_Y +
                    thread_tile_row_idx};
                // 计算A矩阵在共享内存中的列索引
                size_tconst A_thread_block_tile_col_idx{k_i};
                // 从共享内存中读取A矩阵的值
                T const A_val{A_thread_block_tile[A_thread_block_tile_row_idx]
                                                 [A_thread_block_tile_col_idx]};
                // 执行乘累加操作
                C_thread_results[thread_tile_row_idx] += A_val * B_val;
            }
        }
        // 同步所有线程，准备下一次迭代
        __syncthreads();
    }

    // 将结果写入DRAM
#pragma unroll
    for (size_t thread_tile_row_idx{0U};
         thread_tile_row_idx < THREAD_TILE_SIZE_Y; ++thread_tile_row_idx)
    {
        // 计算输出矩阵C的行索引
        size_tconst C_row_idx{blockIdx.y * BLOCK_TILE_SIZE_Y +
                               thread_linear_idx / BLOCK_TILE_SIZE_X *
                                   THREAD_TILE_SIZE_Y +
                               thread_tile_row_idx};
        // 计算输出矩阵C的列索引
        size_tconst C_col_idx{blockIdx.x * BLOCK_TILE_SIZE_X +
                               thread_linear_idx % BLOCK_TILE_SIZE_X};
        // 边界检查并写入结果
        if (C_row_idx < m && C_col_idx < n)
        {
            C[C_row_idx * ldc + C_col_idx] =
                alpha * C_thread_results[thread_tile_row_idx] +
                beta * C[C_row_idx * ldc + C_col_idx];
        }
    }
}

// 启动GEMM内核v03的模板函数
template <typename T>
void launch_gemm_kernel_v03(size_t m, size_t n, size_t k, T const* alpha,
                            T const* A, size_t lda, T const* B, size_t ldb,
                            T const* beta, T* C, size_t ldc,
                            cudaStream_t stream)
{
    // 可以随意调整块分块大小
    // 算法正确性应该始终得到保证
    constexprunsignedint BLOCK_TILE_SIZE_X{64U};      // 块分块X维度大小
    constexprunsignedint BLOCK_TILE_SIZE_Y{64U};      // 块分块Y维度大小
    constexprunsignedint BLOCK_TILE_SIZE_K{8U};       // 块分块K维度大小
    // 每个线程计算C矩阵的THREAD_TILE_SIZE_Y个值
    constexprunsignedint THREAD_TILE_SIZE_Y{8U};      // 线程分块大小
    // 计算每个块的线程数
    constexprunsignedint NUM_THREADS_PER_BLOCK{
        BLOCK_TILE_SIZE_X * BLOCK_TILE_SIZE_Y / THREAD_TILE_SIZE_Y};
    
    // 静态断言确保参数配置的正确性
    static_assert(BLOCK_TILE_SIZE_Y % THREAD_TILE_SIZE_Y == 0U);
    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_K == 0U);
    static_assert(NUM_THREADS_PER_BLOCK % BLOCK_TILE_SIZE_X == 0U);
    
    // 设置块维度（一维线程块）
    dim3 const block_dim{NUM_THREADS_PER_BLOCK, 1U, 1U};
    // 设置网格维度
    dim3 const grid_dim{
        (static_cast<unsignedint>(n) + BLOCK_TILE_SIZE_X - 1U) /
            BLOCK_TILE_SIZE_X,
        (static_cast<unsignedint>(m) + BLOCK_TILE_SIZE_Y - 1U) /
            BLOCK_TILE_SIZE_Y,
        1U};
    
    // 启动GEMM内核
    gemm_v03<T, BLOCK_TILE_SIZE_X, BLOCK_TILE_SIZE_Y, BLOCK_TILE_SIZE_K,
             THREAD_TILE_SIZE_Y><<<grid_dim, block_dim, 0U, stream>>>(
        m, n, k, *alpha, A, lda, B, ldb, *beta, C, ldc);
    // 检查CUDA错误
    CHECK_LAST_CUDA_ERROR();
}

这个FP32 GEMM实现的性能在NVIDIA GeForce RTX 3090 GPU上变成了8.91 TFLOPS。看起来我们一直在取得进展。

具有2D块分块和2D线程分块的实现

如果寄存器数量不是性能瓶颈，我们可以通过缓存矩阵和的数据从共享内存到寄存器来进一步提高性能。每个具有块线程索引的线程，其中和，现在负责计算小矩阵的个元素，其中和分别是行和列的线程分块大小。

在我们之前的实现中有了1D线程分块，为了计算小矩阵的一个元素，我们平均需要从共享内存中读取个值。

现在，有了2D线程分块，为了计算小矩阵的个元素，我们只需要从共享内存中读取个值。具体来说，在每个内层循环中，和被缓存到寄存器中以被重复使用来计算矩阵乘法。总共需要从共享内存中读取个值。平均来说，为了计算小矩阵的一个元素，我们需要从共享内存中读取个值。