【CUDA进阶】Tensor Core实战教程（上）

前言

学习 UP 主 比飞鸟贵重的多_HKL 的 【CUDA进阶】Tensor Core实战教程（已完结） 视频，记录下个人学习笔记，仅供自己参考😄

refer 1：【CUDA进阶】Tensor Core实战教程（已完结）

refer 2：https://github.com/xlite-dev/LeetCUDA

refer 3：https://github.com/Bruce-Lee-LY/cuda_hgemm

refer 4：有关CUBLAS中的矩阵乘法函数

refer 5：https://chatgpt.com

1. 简述

这边 UP 推荐了两个相关的代码仓库可以用来学习 Tensor Core，大家感兴趣的可以看看

第一个仓库地址是：https://github.com/xlite-dev/LeetCUDA

LeetCUDA 是一个面向 CUDA 初学者及进阶开发者的「现代 CUDA 学习笔记」与代码示例集合，基于 PyTorch 环境但核心聚焦在底层 CUDA/Tensor Cores 编程和高性能计算

其核心内容包括以下几个部分：

1. HGEMM：自研高性能矩阵乘（GEMM）内核，覆盖 WMMA、MMA、CuTe API，能够在 NVIDIA L20、RTX 4090、RTX 3080 等设备上达到 98%～100% 的 cuBLAS Tensor Core 性能

2. FlashAttention-MMA：基于纯 MMA PTX 实现的 FlashAttention-2（FA2），支持多阶段流水、Warp/Thread 级切片、QKV 细粒度 Tiling、Shared QKV、Collective Store 等优化手段，在某些配置下性能超过标准 SDPA/FA2 实现

3. 200+ CUDA Kernels：从「Easy⭐️」到「Hard++⭐⭐⭐️⭐️⭐️」不同难度等级的示例，涵盖基础并行模式、共享内存优化、异步拷贝、Block/Thread Tiling、Warp Swizzle、分块调度等超过 200 个经典或创新内核

4. 100+ LLM/CUDA 博客：系统化讲解 Transformer Attention、矩阵乘加、数据类型（TF32/F16/BF16/F8）在 Tensor Core 上的实现要点与性能陷阱

5. 配套文档与幻灯片：docs/、slides/ 等目录下提供了可复用的讲义、流程图与注释，方便教学和分享

第二个仓库地址是：https://github.com/Bruce-Lee-LY/cuda_hgemm

cuda_hgemm 仓库是一个基于 CUDA 的半精度（FP16）矩阵乘加（GEMM）实现与性能基准测试库，旨在通过底层优化策略充分发挥 NVIDIA Tensor Core 的算力，达到或超越 cuBLAS 性能。核心优化技术如下：

• 分块与切片（Tiling）：使用如 Block 256x128、Warp 64x64 的多级切片策略
• 访存合并：使用宽字节（wide）内存指令保证 global memory 访存合并
• 内存复用：将 A、B 子块加载到 shared memory 中多次复用
• 异步拷贝（Async Copy）：利用非阻塞预取（copy.async）提前加载数据
• 消除访存冲突：对 WMMA 使用适当的内存 Padding，对 PTX MMA 进行数据重排
• L2 缓存重排（cache swizzle）：优化全局访问的缓存命中率
• 寄存器级流水线：多缓冲设计（阶段划分为 Pg2s、Ps2r、Stage）实现并行调度

我们后续的实现代码来自于 LeetCUDA 这个仓库，因为它的可读性较好，比较适合初学者，不过一些工具类和测试类的实现我们参考的是 cuda_hgemm 这个仓库

2. Tensor Core

在正式开始之前，我们有必要了解下 Tensor Core

以下内容均 Copy 自 ZOMI 酱的讲解，强烈建议大家阅读原文或观看视频

视频：NVIDIA英伟达Tensor Core基本原理(上)【AI芯片】GPU架构04

文章：https://github.com/Infrasys-AI/AISystem/tree/main/02Hardware/04NVIDIA

2.1 Tensor Core 工作原理

Tensor Core 是 NVIDIA 针对深度学习和 AI 工作负载而设计的专用核心，可以实现混合精度计算并加速矩阵运算，尤其擅长处理半精度（FP16）和全精度（FP32）的矩阵乘法和累加操作。Tensor Core 在加速深度学习训练和推理中发挥着重要作用。

随着 Volta 架构的推出，英伟达引入了 Tensor Core，在具体的运算过程中，Tensor Core 采用融合乘法加法（FMA）的方式来高效地处理计算任务。每个 Tensor Core 每周期能执行 4x4x4 GEMM，64 个浮点乘法累加（FMA）运算。

Tensor Core 4x4x4 matrix multiply and accumulate

如上图所示，在执行运算 D=A*B+C，其中 A、B、C 和 D 是 4x4 矩阵。矩阵乘法 输入 A 和 B 是 FP16 矩阵，而 累加矩阵 C 和 D 可以是 FP16 或 FP32 矩阵

具体来说，它首先接受两个 4x4 的 FP16 精度的输入矩阵 A 和 B，执行它们的矩阵乘法。然后，将这个乘法的结果与第三个 4x4 的矩阵 C 相加，其中矩阵 C 可以是 FP16 或 FP32 精度。最终，Tensor Core 输出一个新的 4x4 矩阵 D，该矩阵同样可以是 FP16 或 FP32 精度

这也就实现了底层硬件上的混合精度计算。通过将矩阵乘法的输入限定为 FP16 精度，可以大幅减少所需的计算资源和内存带宽，从而加速计算。同时，通过允许累加矩阵 C 和输出矩阵 D 使用 FP32 精度，可以保证运行结果的准确性和数值稳定性。这种灵活的精度策略，结合 Tensor Core 的高效计算能力，使得在保持高性能的同时，还能有效控制神经网络模型的训练和推理过程中的资源消耗

接下来我们再打开一层进一步探讨 Tensor Core 的运算能力。上文我们谈到在每个 Tensor Core 每个时钟执行 64 个 FP32 FMA 混合精度运算，一个 SM 中一共有 8 个 Tensor Core，所以每个时钟周期内总共执行 512 个浮点运算（8 个 Tensor Core $\times$ 64 个 FMA 操作/核）

因此在 AI 应用中，Volta V100 GPU 的吞吐量与 Pascal P100 GPU 相比，每个 SM 的 AI 吞吐量提高 8 倍。此外得益于 Volta 架构在 SM 数量和核心设计上的优化，总体上共提高 12 倍

2.2 Tensor Core 与 CUDA 编程

如下图所示，在 CUDA 编程体系中，我们并非直接对线程进行控制，也就是图中的弯弯的线，而是通过控制一个 Warp（线程束），一个 Warp 包含很多线程（通常为 32 个线程），这些线程同时并行执行，利用 GPU 的并行计算能力

在实际执行过程中，CUDA 会对 Warp 进行同步操作，确保其中的所有线程都达到同步点，并获取相同的数据。然后，这些线程将一起执行矩阵相乘和其他计算操作，通常以 16x16 的矩阵块为单位进行计算。最终，计算结果将被存储回不同的 Warp 中，以便后续处理或输出

我们可以把 Warp 理解为软件上的一个大的线程概念，它帮助简化了对 GPU 并行计算资源的管理和利用。通过有效地利用 Warp 的并行性，CUDA 程序可以实现高效、快速的并行计算

在 CUDA 程序执行过程中，我们可以通过线程的 Warp 来调度 Tensor Core 的执行。多个 Tensor Core 可以同时通过 Warp 内的线程来执行计算任务，利用 Tensor Core 提供的高性能矩阵运算能力，每个 Warp 内的线程可以利用 Tensor Core 执行 16x16x16 的矩阵运算，充分发挥 GPU 的计算潜能

下面是 CUDA 提供的 WMMA API（Warp-Matrix-Multiply-Accumulate），也就是给 Tensor Core 编程用的一套 C++ 接口：

template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment;

void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm);
void load_matrix_sync(fragment<...> &a, const T* mptr, unsigned ldm, layout_t layout);
void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);
void fill_fragment(fragment<...> &a, const T& v);
void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);

其中：

• fragment：Tensor Core 数据存储类，支持 matrix_a、matrix_b、accumulator
• load_matrix_sync：Tensor Core 数据加载 API，支持将矩阵数据从 global memory 或 shared memory 加载到 fragment
• store_matrix_sync：Tensor Core 结果存储 API，支持将计算结果从 fragment 存储到 global memory 或 shared memory
• fill_fragment：fragment 填充 API，支持常数值填充
• mma_sync：Tensor Core 矩阵乘计算 API，支持 D = AB + C 或者 C = AB +C

CUDA 通过 CUDA C++ WMMA API 向外提供了 Tensor Core 在 Warp 级别上的计算操作支持。这些 C++ 接口提供了专门用于矩阵加载、矩阵乘法和累加以及矩阵存储等操作的功能。例如上述代码中，其中的 mma_sync 就是执行具体计算的 API 接口。借助这些 API，开发者可以高效地利用 Tensor Core 进行深度学习中的矩阵计算，从而加速神经网络模型的训练和推理过程

那大家可能会有所困惑，明明一个 Tensor Core 每个周期只可以执行 4x4x4 的 GEMM 运算，为什么在 CUDA 的层面却提供了使用 16x16x16 的 GEMM 运算 API 呢？🤔

事实上，如果我们整体来看，如上图所示，一个 Tensor Core 是一个 4x4 的 Tensor Core 核心。但实际上，在一个 SM（Streaming Multiprocessor）中有多个 Tensor Core，我们无法对每个 Tensor Core 进行细粒度的控制，因为效率会很低。因此，一个 Warp 就扮演了重要角色，将多个 Tensor Core 打包在一起，以执行更大规模的计算任务

通过 Warp 层的计算指令，CUDA 向外提供了一个 16x16x16 的抽象层，使得开发者可以通过一条指令完成多个 Tensor Core 的协同工作，实现高效的并行计算。这条指令也即我们之前提到的 mma_sync API，它允许开发者利用 Warp 内的线程同时调度多个 Tensor Core 执行矩阵乘加操作，从而提高 GPU 计算的效率和性能

2.3 Tensor Core 实现大矩阵计算

在整体 CUDA 软件设计方面，其目的是与 GPU 计算和存储分层结构相匹配。英伟达 CUDA 对于 Tensor Core 的定义主要是通过 CUDA 提供一种通用的编程范式，即所谓的 General Programming。为了更高效地利用 Tensor Core，CUDA 允许开发者利用通用编程模型来调用 Tensor Core 的硬件功能，以实现高效的并行计算

假设我们 Tensor Core 中的 D = A * B + C 计算简化为 C = A * B。在实际应用中，由于 Tensor Core 只能处理 4x4 的简单计算，不可能直接将整个大矩阵载入 Tensor Core 中进行运算。因此，需要将矩阵切片，并将其分配到 Thread Block（线程块）中

接着，在软件层面会定义一个 Warp（线程束），将这些切片矩阵分配给不同的 Warp。最终，通过线程的执行实现对矩阵乘法的并行计算，充分利用 Tensor Core 的计算能力。这种分块、分配和并行执行的方式有效地利用了硬件资源，提高了计算效率，从而加速了深度学习和科学计算等领域的应用

下面我们将详细展开 Tensor Core 是如何完成大矩阵计算的

2.3.1 Block-level 矩阵乘

在 Tensor Core 中一个矩阵乘的计算，也就是所谓的 GEMM，其中一次只能计算一个小的矩阵块。在实际的运算中，也就需要把矩阵 A 切分出来一个小块，把矩阵 B 也切分出来个小块，算出来一个小的矩阵 C，如下图所示：

那这个时候在整体的软件编程时，就会沿着每一个维度，如上图的 N 维和 K 维进行拆分为小的矩阵进行计算，最后对结果进行累积，代码如下：

for (int mb = O; mb < M; mb += Mtile)
  for (int nb = O; nb < N; nb += NtiLe)
    for(int kb = O; kb < K; kb += Ktile)
    {
      //compute Mtile-by-Ntile-by-Ktile matrix product
      for (int k = O; k < Ktile; ++k)
        for(int i= O; i< Mtile; ++i)
          for (int j= O; j< Ntile; ++j)
          {
            int row = mb + i;
            int col = nb + j;
            C[row][col] += A[row][kb + k] * B[kb + k][col];
          }
    }

在上述代码中，我们沿每个维度将循环嵌套 Loop nest 划分为块 blocks，然后划分成 Mtile-by-Ntile 的独立矩阵乘，最后通过累积 Mtile-by-Ntile-by-Ktile 的矩阵乘积来计算每个乘积

在 GPU 计算时，使用 CUDA Kernel grid 将 CUDA 线程块分配给输出矩阵 C 的每个分区，CUDA 线程块并行计算 Mtile-by-Ntile-by-Ktile 矩阵乘，在 K 维上进行迭代，执行累积 Mtile-by-Ntile-by-Ktile 矩阵乘的结果。可以看出这里计算主要是在线程块也就是 Thread Block 里面去进行的并行计算

2.3.2 Warp-level 矩阵乘

在 CUDA 编程模型中，当我们在线程块（Block）内执行矩阵乘法操作时，这些操作实际上是在 Warp 级别上被分配和执行的。Warp 是 GPU 上的一个执行单元，它由固定数量的线程（通常是 32 个线程）组成，这些线程协同工作以执行相同的指令

在进行矩阵乘法时，为了加速计算并减少内存访问延迟，通常会将矩阵 A 和矩阵 B 的部分数据加载到共享内存（Shared Memory，简称 SMEM）中。共享内存是线程块内所有线程都可以访问的一块快速内存，它允许线程与线程之间进行数据交换和协作，而不必每次都从全局内存（Global Memory）中读取数据

在矩阵乘法中，每个 Warp 会负责计算结果矩阵 C 的一个或多个部分。这通常通过将结果矩阵 C 的不同块分配给不同的 Warp 来实现，每个 Warp 独立地计算其分配到的部分。由于 Warp 内的线程是同步执行的，因此它们可以共同协作，使用共享内存中的数据来完成它们的计算任务。这种分配方式充分利用了 GPU 的并行计算能力，并减少了内存访问的延迟，从而提高了矩阵乘法的性能

当执行矩阵乘法时，Warp 中的线程会协同工作，完成一系列乘法和加法操作。这个过程涉及到从共享内存（SMEM）加载数据到寄存器（RF）进行计算，然后将结果存储回寄存器或全局内存中

下面我们详细解释一下这个过程，首先，多个线程组成一个 Warp，协同工作以处理矩阵乘法中的一部分。这些线程共同合作，通过执行一系列乘法和加法操作，能够高效地计算出结果

其次，为了进行计算，Warp 中的线程需要从共享内存中加载矩阵 A 和 B 的片段到它们的寄存器中。这些片段是矩阵的一小部分，加载到寄存器中可以实现快速的数据访问和计算。这要求数据从共享内存到寄存器的加载速度足够快，以避免计算线程等待数据，从而保持计算的高效性

然后，每个线程在寄存器上执行矩阵乘法操作，计算结果矩阵 C 的一个或多个元素。这些元素暂存于线程的寄存器中，直到所有必要的乘法和加法操作完成。在计算过程中，为了最大化线程的计算效率，共享内存中的数据按照特定维度（K 维）进行排序。这种排序方式有助于减少内存访问延迟，使得线程能够更高效地访问所需的数据

2.3.3 Thread-level 矩阵乘

在 Tensor Core 里面并行执行的就是上述的形式，矩阵 A 乘以矩阵 B 等于 C 矩阵这么一个简单最核心操作

Tensor Core 是英伟达 GPU 的硬件，CUDA 编程模型提供了 WMMA（Warp-level Matrix Multiply-Accumulate）API，这个 API 是专门为 Tensor Core 设计的，它允许开发者在 CUDA 程序中直接利用 Tensor Core 的硬件加速能力。通过 WMMA API，开发者可以执行矩阵乘法累积操作，并管理数据的加载、存储和同步

在 GEMM（General Matrix Multiply，通用矩阵乘法）的软硬件分层中，数据复用是一个非常重要的概念。由于矩阵通常很大，包含大量的数据，因此有效地复用这些数据可以显著提高计算效率。在每一层中，都会通过不同的内存层次结构（如全局内存、共享内存和寄存器）来管理和复用数据

具体来说，大矩阵通常被分割成小块，并存储在全局内存中。然后，在计算过程中，这些小块数据会被加载到共享内存或寄存器中，以便进行高效的计算。通过这些方式，可以最大限度地减少内存访问延迟，并提高计算吞吐量。在计算完每一块矩阵乘法后，得到的结果通常也是一小块数据。为了得到最终的完整矩阵乘法结果，需要将所有小块结果累积起来。这通常涉及到将中间结果从寄存器或共享内存写回到全局内存中，并在必要时进一步的同步和累加操作

最终，通过一系列这样的计算和数据管理操作，可以完成整个 GEMM 计算任务，并将结果写回到输出矩阵 C 中。整个过程充分利用了 Tensor Core 的硬件加速能力和 CUDA 编程模型的灵活性，从而实现了高效的矩阵乘法计算

2.3.4 累积矩阵结果

下面我们再来详细展开 Tensor Core 是如何完成累积矩阵并写出最终结果的

存在 register file 中的临时结果的回传是通过 WMMA 的 API 完成的，在 Tensor Core 提供的 WMMA API 里面有个 store matrix sync 的 API。这个 API 工作就是把所有的数据都搬到共享内存，也就是 SMEM 里面

在 SMEM 里会做大量的累积的操作，把所有的这些数据累积起来后，再存放在全局内存里面。通过全局内存把一个块一个块的拼接起来，把数据回传写出结果

2.3.5 整体计算过程

下面我们来总结一下整个的计算过程：

首先，矩阵在进行 GEMM 计算之前会被分块，这些分块后的矩阵存储在全局内存（Global Memory）中。全局内存是 GPU 上最大的内存区域，用于存储计算过程中需要访问的大量数据

接下来，在计算开始前，程序会将需要参与计算的矩阵分块加载到共享内存（Shared Memory）中。共享内存是每个线程块中所有线程都可以访问的低延迟内存池，它使得同一线程块中的线程能够高效地共享和重用数据

然后，当实际执行矩阵乘法运算时，线程会将共享内存中的数据加载到其私有的寄存器（Register）中。寄存器是 GPU 上访问速度最快的内存空间，每个线程都有自己独立的寄存器文件。在 Tensor Core 中执行矩阵乘法运算时，数据会存储在 Tensor Core 的寄存器文件中，并在这里进行计算

计算完成后，结果会写回到共享内存中，而不是继续存储在 Tensor Core 的寄存器文件中。这是因为共享内存更适合用于中间结果的存储和线程间的数据交换。在写回共享内存的过程中，通过 CUDA 的 WMMA API 提供了诸如 store matrix sync 等函数，用于确保数据的正确同步和累积

在共享内存中进行结果累积后，这些累积的结果最终会被写回到全局内存中。这个过程可能涉及多个线程块的协作，因为整个矩阵乘法运算可能需要多个线程块共同完成。通过全局内存，不同线性块计算得到的结果块可以被拼接起来，形成最终的完整矩阵乘法结果

3. 开发环境及框架搭建

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

cuda_learn/
├── CMakeLists.txt
└── tensor_core_wmma
    ├── CMakeLists.txt
    ├── common
    │   ├── common.h
    │   ├── cuda_timer.h
    │   ├── logging.h
    │   ├── matrix.h
    │   ├── ptx.h
    │   ├── tester.h
    │   └── util.h
    ├── hgemm_v1_wmma_m16n16k16_naive_kernel.cu
    ├── hgemm_v2_wmma_m16n16k16_mma4x2_kernel.cu
    ├── hgemm_v3_wmma_m16n16k16_mma4x2_warp2x4_kerne.cu
    └── hgemm_v4_wmma_m16n16k16_mma4x2_warp2x4_dbuf_async_kernel.cu

3 directories, 13 files

其中 common 文件夹来自于 cuda_hgemm，hgemm 实现代码来自于 LeetCUDA，具体的内容我们后续会依次分析

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

cmake_minimum_required(VERSION 3.20.0)
project(cuda_practice VERSION 0.1.0 LANGUAGES CUDA CXX C)
set(CMAKE_CUDA_ARCHITECTURES 89)
find_package(CUDAToolkit)
add_subdirectory(tensor_core_wmma)

Note：显卡的 CUDA Compute Capability 可查询官方文档：https://developer.nvidia.com/cuda-gpus

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

add_executable(hgemm_v1_wmma_m16n16k16_naive_kernel hgemm_v1_wmma_m16n16k16_naive_kernel.cu)
target_link_libraries(hgemm_v1_wmma_m16n16k16_naive_kernel PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
    target_compile_options(hgemm_v1_wmma_m16n16k16_naive_kernel PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
endif()
# target_compile_options(hgemm_v1_wmma_m16n16k16_naive_kernel PRIVATE -lineinfo)

add_executable(hgemm_v2_wmma_m16n16k16_mma4x2_kernel hgemm_v2_wmma_m16n16k16_mma4x2_kernel.cu)
target_link_libraries(hgemm_v2_wmma_m16n16k16_mma4x2_kernel PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
    target_compile_options(hgemm_v2_wmma_m16n16k16_mma4x2_kernel PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
endif()

add_executable(hgemm_v3_wmma_m16n16k16_mma4x2_warp2x4_kerne hgemm_v3_wmma_m16n16k16_mma4x2_warp2x4_kerne.cu)
target_link_libraries(hgemm_v3_wmma_m16n16k16_mma4x2_warp2x4_kerne PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
    target_compile_options(hgemm_v3_wmma_m16n16k16_mma4x2_warp2x4_kerne PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
endif()

add_executable(hgemm_v4_wmma_m16n16k16_mma4x2_warp2x4_dbuf_async_kernel hgemm_v4_wmma_m16n16k16_mma4x2_warp2x4_dbuf_async_kernel.cu)
target_link_libraries(hgemm_v4_wmma_m16n16k16_mma4x2_warp2x4_dbuf_async_kernel PRIVATE CUDA::cudart ${CUDA_cublas_LIBRARY})
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
    target_compile_options(hgemm_v4_wmma_m16n16k16_mma4x2_warp2x4_dbuf_async_kernel PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-G>)
endif()

编译运行指令如下：

cd cuda_learn
mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Debug && make -j24
./tensor_core_wmma/hgemm_v1_wmma_m16n16k16_naive_kernel

我们来具体看看代码中都做了些什么，我们以 hgemm_v1_wmma_m16n16k16_naive_kernel.cu 为例来讲解，代码如下：

#include <iostream>
#include <cuda_runtime.h>
#include "common/tester.h"
#include "common/common.h"

using namespace nvcuda;

// only 1 warp per block(32 threads), m16n16k16. A, B, C: all row_major
template<const int WMMA_M = 16, const int WMMA_N = 16, const int WMMA_K = 16>
__global__ void hgemm_wmma_m16n16k16_naive_kernel(half* A, half* B, half* C, int M, int N, int K){

    const int NUM_K_TILES = div_ceil(K, WMMA_K);
    const int load_gmem_a_m = blockIdx.y * WMMA_M;
    const int load_gmem_b_n = blockIdx.x * WMMA_N;

    if(load_gmem_a_m >= M && load_gmem_b_n >= N){
        return;
    }

    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, half> C_frag;
    wmma::fill_fragment(C_frag, 0.0);

#pragma unroll
    for(int k = 0; k < NUM_K_TILES; ++k){
        
        wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> A_frag;
        wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> B_frag;

        wmma::load_matrix_sync(A_frag, A + load_gmem_a_m * K + k * WMMA_K, K);
        wmma::load_matrix_sync(B_frag, B + (k * WMMA_K) * N + load_gmem_b_n, N);

        wmma::mma_sync(C_frag, A_frag, B_frag, C_frag);

        __syncthreads();
    }
    
    wmma::store_matrix_sync(C + load_gmem_a_m * N + load_gmem_b_n, C_frag, N, wmma::mem_row_major);
}

void hgemm_wmma_m16n16k16_naive(half* A, half* B, half* C, int M, int N, int K){

    constexpr int WMMA_M = 16;
    constexpr int WMMA_N = 16;
    constexpr int WMMA_K = 16;
    dim3 block(32);
    dim3 grid(div_ceil(N, WMMA_N), div_ceil(M, WMMA_M));

    hgemm_wmma_m16n16k16_naive_kernel<WMMA_M, WMMA_N, WMMA_K><<<grid, block>>>(A, B, C, M, N, K);
}

int main(int argc, char* argv[]){

    Tester tester(512, 2048, 1024, 1, 10, 100, true);
    tester.evaluate(hgemm_wmma_m16n16k16_naive, "hgemm_wmma_m16n16k16_naive");
    return 0;
}

下面我们按调用流程，从最顶层的 main 一路往下，来分段解析整个运行流程：(from ChatGPT)

1. main 函数（hgemm_v1_wmma_m16n16k16_naive_kernel.cu）

int main(int argc, char* argv[]){
    Tester tester(512, 2048, 1024, 1, 10, 100, true);
    tester.evaluate(hgemm_wmma_m16n16k16_naive, "hgemm_wmma_m16n16k16_naive");
    return 0;
}

• 创建一个 Tester 对象，指定矩阵维度 (M=512, N=2048, K=1024)，以及 warmup 次数、profile 次数，sleep 时长、是否开启结果校验
• 调用 tester.evaluate(...)，传入待测函数 hgemm_wmma_m16n16k16_naive 及其名称字符串，用于后续日志输出和性能对比

2. Tester 构造与基准（tester.h）

// Copyright 2023. All Rights Reserved.
// Author: Bruce-Lee-LY
// Date: 20:42:28 on Sun, Feb 12, 2023
//
// Description: tester

#pragma once

#include <memory>

#include "cuda_timer.h"
#include "matrix.h"

class Tester {
public:
    explicit Tester(size_t M = 512, size_t N = 2048, size_t K = 1024, size_t warmup_iterations = 1,
                    size_t profiling_iterations = 10, size_t sleep_duration = 100, bool enable_check = false)
        : m_M(M),
          m_N(N),
          m_K(K),
          m_warmup_iterations(warmup_iterations),
          m_profiling_iterations(profiling_iterations),
          m_sleep_duration(sleep_duration),
          m_enable_check(enable_check) {
        HGEMM_CHECK_GT(m_M, 0);
        HGEMM_CHECK_GT(m_N, 0);
        HGEMM_CHECK_GT(m_K, 0);
        HGEMM_CHECK_GT(m_warmup_iterations, 0);
        HGEMM_CHECK_GT(m_profiling_iterations, 0);
        HGEMM_CHECK_GT(m_sleep_duration, 0);

        m_A = std::make_shared<Matrix>(m_M, m_K, "Matrix A");
        HGEMM_CHECK(m_A);
        m_B = std::make_shared<Matrix>(m_K, m_N, "Matrix B");
        HGEMM_CHECK(m_B);
        m_C = std::make_shared<Matrix>(m_M, m_N, "Matrix C");
        HGEMM_CHECK(m_C);
        m_base = std::make_shared<Matrix>(m_M, m_N, "Matrix Base");
        HGEMM_CHECK(m_base);

        if (m_enable_check) {
            m_cuda_timer.start();
            cublas_tensor_op(m_A->getDevPtr(), m_B->getDevPtr(), m_base->getDevPtr(), m_M, m_N, m_K);
            HLOG("Cublas-Tensor-Op use: %.3f ms", m_cuda_timer.end());
            m_base->moveToHost();
            m_base->memSetDevice();
        }
    }

    ~Tester() {}

    template <typename Func>
    void evaluate(Func &&hgemm, const std::string &name) {
        HLOG("----------------- Evaluating %s -----------------", name.c_str());
        usleep(m_sleep_duration * 1000);
        m_C->tearUp(m_base.get());

        // warm up
        m_cuda_timer.start();
        for (size_t i = 0; i < m_warmup_iterations; ++i) {
            hgemm(m_A->getDevPtr(), m_B->getDevPtr(), m_C->getDevPtr(), m_M, m_N, m_K);
        }
        m_warmup_time = static_cast<double>(m_cuda_timer.end()) / static_cast<double>(m_warmup_iterations);
        HLOG("Warm up time: %.3f ms", m_warmup_time);

        if (m_enable_check) {
            m_C->moveToHost();
            m_C->checkValue(m_base.get());
        }

        profile(std::forward<Func>(hgemm), name);
    }

private:
    void cublas_tensor_op(half *A, half *B, half *C, size_t M, size_t N, size_t K) {
        cublasHandle_t handle = nullptr;
        HGEMM_CHECK_CUBLAS_ERROR(cublasCreate(&handle));
        HGEMM_CHECK_CUBLAS_ERROR(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));

        half alpha = 1.0;
        half beta = 0.0;

        HGEMM_CHECK_CUBLAS_ERROR(cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_16F, K, A,
                                              CUDA_R_16F, K, &beta, C, CUDA_R_16F, N, CUBLAS_COMPUTE_16F,
                                              CUBLAS_GEMM_DEFAULT_TENSOR_OP));
        HGEMM_CHECK_CUBLAS_ERROR(cublasDestroy(handle));
    }

    template <typename Func>
    void profile(Func &&hgemm, const std::string &name) {
        m_cuda_timer.start();
        for (size_t i = 0; i < m_profiling_iterations; ++i) {
            hgemm(m_A->getDevPtr(), m_B->getDevPtr(), m_C->getDevPtr(), m_M, m_N, m_K);
        }
        m_profiling_time = static_cast<double>(m_cuda_timer.end()) / static_cast<double>(m_profiling_iterations);
        m_throughput =
            static_cast<double>(m_M * m_N * m_K * 2) * 1e-12 / (static_cast<double>(m_profiling_time) * 1e-3);

        if ((std::abs(m_base_time) <= 1e-6) && (std::abs(m_base_throughput) <= 1e-6)) {
            m_base_time = m_profiling_time;
            m_base_throughput = m_throughput;
        }

        HLOG("%s exit, profiling time: %.3f ms (%.2f%%), throughput: %.3f TFLOPS (%.2f%%)", name.c_str(),
             m_profiling_time, m_profiling_time / m_base_time * 100, m_throughput,
             m_throughput / m_base_throughput * 100);
    }

    const size_t m_M = 512;
    const size_t m_N = 2048;
    const size_t m_K = 1024;
    const size_t m_warmup_iterations = 1;
    const size_t m_profiling_iterations = 10;
    const size_t m_sleep_duration = 100;
    const bool m_enable_check = false;

    std::shared_ptr<Matrix> m_A = nullptr;     // row major, M * K
    std::shared_ptr<Matrix> m_B = nullptr;     // col major, K * N
    std::shared_ptr<Matrix> m_C = nullptr;     // row major, M * N
    std::shared_ptr<Matrix> m_base = nullptr;  // row major, M * N, base result, init matrix C before each hgemm

    CudaTimer m_cuda_timer;

    double m_warmup_time = 0.0;
    double m_profiling_time = 0.0;
    double m_throughput = 0.0;
    double m_base_time = 0.0;        // cublas tensor op default
    double m_base_throughput = 0.0;  // cublas tensor op default

    HGEMM_DISALLOW_COPY_AND_ASSIGN(Tester);
};

• 成员初始化
  • 存储各维度与迭代次数参数，并检查其合法性
• 矩阵分配
  • 用 Matrix 类分别分配并初始化：
  • m_A：尺寸 M x K
  • m_B：尺寸 K x N
  • m_C：尺寸 M x N（用于最终结果）
  • m_base：尺寸 M x N（用于存储 cuBLAS 基准结果）
• 基准计算（可选）
  • 若 enable_check == true：
  • 1. 调用 cublas_tensor_op 用 cuBLAS-Tensor-Op 计算一次基准矩阵乘，计时并打印耗时
  • 2. 将 m_base 拷贝回 host，用于后续结果校验

3. evaluate 方法流程（tester.h）

template <typename Func>
void Tester::evaluate(Func &&hgemm, const std::string &name) {
    // 1. 打印评测标题
    // 2. usleep(sleep_duration) —— 避免上下文冲突
    // 3. m_C.tearUp(m_base) —— 将 m_C 设为基准值，方便差异检测

    // —— Warmup 阶段 ——  
    // 4. 启动计时器
    // 5. 连续调用 hgemm(...) warmup_iterations 次
    // 6. 结束计时，计算并打印平均 warmup 时间
    // 7.（若开启校验）将 m_C 拷回 host，与 m_base 对比，输出最大/平均差异

    // —— Profiling 阶段 ——  
    // 8. 调用 profile(hgemm, name)
}

4. 基准与性能统计（tester.h）

template <typename Func>
void Tester::profile(Func &&hgemm, const std::string &name) {
    // 1. 启动计时器
    // 2. 重复执行 hgemm(...) profiling_iterations 次
    // 3. 结束计时，计算平均 profiling 时间
    // 4. 计算 TFLOPS = (2*M*N*K) / (time_in_s * 1e12)
    // 5. 与基准（第一次 profile 即为基准）比较，打印相对比率和 TFLOPS
}

• 吞吐量计算：将 M x N x K x 2（乘加算子数）换算到每秒 TFLOPS
• 对比：首轮 profile 结果作为 “base”，后续与之比较百分比

5. Matrix 类：数据管理（matrix.h）

// Copyright 2023. All Rights Reserved.
// Author: Bruce-Lee-LY
// Date: 20:42:28 on Sun, Feb 12, 2023
//
// Description: matrix

#pragma once

#include <random>

#include "common.h"

class Matrix {
public:
    Matrix(size_t row, size_t col, const std::string &name = "Matrix", float min = -1.0, float max = 1.0)
        : m_row(row), m_col(col), m_name(name), m_min(min), m_max(max) {
        HGEMM_CHECK_GT(m_row, 0);
        HGEMM_CHECK_GT(m_col, 0);

        m_elem_num = m_row * m_col;
        HGEMM_CHECK_GT(m_elem_num, 0);

        m_host_ptr = new half[m_elem_num];
        HGEMM_CHECK(m_host_ptr);
        HGEMM_CHECK_CUDART_ERROR(cudaMalloc((void **)&m_dev_ptr, m_elem_num * sizeof(half)));
        HGEMM_CHECK(m_dev_ptr);

        std::random_device rd;
        std::default_random_engine engine{rd()};
        std::uniform_real_distribution<float> uniform(m_min, m_max);
        for (size_t i = 0; i < m_elem_num; ++i) {
            m_host_ptr[i] = __float2half(uniform(engine));
        }

        HGEMM_CHECK_CUDART_ERROR(cudaMemcpy(m_dev_ptr, m_host_ptr, m_elem_num * sizeof(half), cudaMemcpyHostToDevice));

        HLOG("%s: %zu * %zu, cpu: %p, gpu: %p", m_name.c_str(), m_row, m_col, m_host_ptr, m_dev_ptr);
    }

    ~Matrix() {
        if (m_host_ptr) {
            delete[] m_host_ptr;
            m_host_ptr = nullptr;
        }

        if (m_dev_ptr) {
            HGEMM_CHECK_CUDART_ERROR(cudaFree((void *)m_dev_ptr));
            m_dev_ptr = nullptr;
        }
    }

    size_t getRow() const {
        return m_row;
    }

    size_t getCol() const {
        return m_col;
    }

    size_t getElemNum() const {
        return m_elem_num;
    }

    half *getHostPtr() const {
        return m_host_ptr;
    }

    half *getDevPtr() const {
        return m_dev_ptr;
    }

    void tearUp(Matrix *base) {
        HGEMM_CHECK(base);
        HGEMM_CHECK_EQ(m_row, base->getRow());
        HGEMM_CHECK_EQ(m_col, base->getCol());

        HGEMM_CHECK_CUDART_ERROR(
            cudaMemcpy(m_dev_ptr, base->getDevPtr(), m_elem_num * sizeof(half), cudaMemcpyDeviceToDevice));
    }

    void moveToHost() {
        HGEMM_CHECK_CUDART_ERROR(cudaMemcpy(m_host_ptr, m_dev_ptr, m_elem_num * sizeof(half), cudaMemcpyDeviceToHost));
    }

    void moveToDevice() {
        HGEMM_CHECK_CUDART_ERROR(cudaMemcpy(m_dev_ptr, m_host_ptr, m_elem_num * sizeof(half), cudaMemcpyHostToDevice));
    }

    void memSetHost() {
        memset(m_host_ptr, 0, m_elem_num * sizeof(half));
    }

    void memSetDevice() {
        HGEMM_CHECK_CUDART_ERROR(cudaMemset(m_dev_ptr, 0, m_elem_num * sizeof(half)));
    }

    void checkValue(Matrix *base) {
        HGEMM_CHECK(base);
        HGEMM_CHECK_EQ(m_row, base->getRow());
        HGEMM_CHECK_EQ(m_col, base->getCol());

        m_max_diff = 0.0;
        m_avg_diff = 0.0;
        double diff = 0.0;
        for (size_t i = 0; i < m_elem_num; ++i) {
            diff = static_cast<double>(std::abs(__half2float(m_host_ptr[i]) - __half2float(base->getHostPtr()[i])));
            m_max_diff = std::max(m_max_diff, diff);
            m_avg_diff += diff;
        }

        m_avg_diff /= static_cast<double>(m_elem_num);

        HLOG("Max diff: %f, avg diff: %f", m_max_diff, m_avg_diff);
    }

private:
    const size_t m_row = 0;
    const size_t m_col = 0;
    const std::string m_name = "Matrix";
    // the threshold of the random matrix will affect the difference of the hgemm results
    const float m_min = -1.0;
    const float m_max = 1.0;

    size_t m_elem_num = 0;
    half *m_host_ptr = nullptr;
    half *m_dev_ptr = nullptr;

    double m_max_diff = 0.0;
    double m_avg_diff = 0.0;

    HGEMM_DISALLOW_COPY_AND_ASSIGN(Matrix);
};

• 构造
  • 1. 在 CPU 上分配版进度数组 half* m_host_ptr，并随机填充
  • 2. 在 GPU 上 cudaMalloc 分配 half* m_dev_ptr，并一次性拷贝至设备
• 常用方法
  • tearUp(base)：在每次 evaluate 前，将 m_C 的设备内存快速 reset 为基准值（Device→Device 拷贝）
  • moveToHost() / moveToDevice()：Host↔Device 拷贝
  • memSetHost() / memSetDevice()：置零
  • checkValue(base)：在 Host 上逐元素比较，与基准 base 计算最大/平均差异

OK，整个框架搭建完成后我们就来看看 GPU 上如何通过 tensor core 来实现 hgemm

4.【cublas】cublasGemmEx 详解

在正式开始之前，我们先来看看如何用最简单的方式来使用 tensor core 实现 hgemm 矩阵乘法的运算

那其实我们可以直接使用 cuBLAS 库的 API 即 cublasGemmEx 来快速实现，这种方式应该是性能最好且速度最快的一种方式了，并且它也最简单。在前面框架搭建小节中我们也提到过和基准 baseline 比较，其中的基准 baseline 其实就是 cuBLAS 库提供的 API 接口实现

NVIDIA 提供的 cuBLAS（Basic Linear Algebra） 库是一个用于加速人工智能和高性能计算应用的 GPU 加速库。它包含多个 API 扩展，可提供可直接插入的行业标准 BLAS API 和 GEMM API，并支持针对 NVIDIA GPU 高度优化的融合。cuBLAS 库还包含用于批处理操作、跨多个 GPU 执行以及混合和低精度执行的扩展，并提供额外的调整以获得最佳性能

下面我们基于 NVIDIA cuBLAS 官方文档，对 cublasGemmEx 接口做一个全面、结构化的说明，包括函数原型、参数意义、数据类型支持、示例用法以及注意事项

4.1 函数原型

C 接口

cublasStatus_t cublasGemmEx(
    cublasHandle_t      handle,
    cublasOperation_t   transa,
    cublasOperation_t   transb,
    int                 m,
    int                 n,
    int                 k,
    const void         *alpha,
    const void         *A,
    cudaDataType_t      Atype,
    int                 lda,
    const void         *B,
    cudaDataType_t      Btype,
    int                 ldb,
    const void         *beta,
    void               *C,
    cudaDataType_t      Ctype,
    int                 ldc,
    cublasComputeType_t computeType,
    cublasGemmAlgo_t    algo
);

这是 cublasGemmEx 的主要签名，允许用户为 A、B、C 三个矩阵分别指定数据类型，并单独设定计算精度和算法

C++ 兼容变体

cublasStatus_t cublasGemmEx(
    cublasHandle_t    handle,
    cublasOperation_t transa,
    cublasOperation_t transb,
    int               m,
    int               n,
    int               k,
    const void       *alpha,
    const void       *A,
    cudaDataType      Atype,
    int               lda,
    const void       *B,
    cudaDataType      Btype,
    int               ldb,
    const void       *beta,
    void             *C,
    cudaDataType      Ctype,
    int               ldc,
    cudaDataType      computeType,  // 注意这里是 cudaDataType 而非 cublasComputeType_t
    cublasGemmAlgo_t  algo
);

为兼容早期 C++ 代码，提供了第二个 computeType 参数为 cudaDataType 的重载版本

下面我们先对 cublasGemmEx 接口对应的数学表达式进行分析，后续再来看看各个参数的具体含义：(from ChatGPT)

一、GEMM 的数学表达式

cuBLAS 中的通用矩阵乘加（GEMM）操作遵循经典 BLAS 规范，数学公式写作：

$$\mathbf{C}=\alpha \mathrm{op}(\mathbf{A})\mathrm{op}(\mathbf{B})+\beta \mathbf{C}$$

• $\alpha ,\beta$：标量因子，用来分别缩放矩阵乘积结果和原始输出矩阵 $C$
• $\mathrm{op}(\mathbf{A}),\mathrm{op}(\mathbf{B})$：对输入矩阵的 “操作”（不转置、转置或共轭转置）。由 transa/transb 参数决定
• 最终结果累加到原始的 $\mathbf{C}$ 上

二、矩阵存储格式与维度约定

1. 列主序（column-major）存储

• cuBLAS 遵循 Fortran/BLAS 传统：矩阵在内存中按列连续存储
• 比如一个 $m\times n$ 矩阵，先是第一列的 m 个元素，再是第 2 列的 m 个元素，依此类推

2. 维度匹配

• 若不考虑转置，$\mathbf{A}$ 是 $m\times k$，$\mathbf{B}$ 是 $k\times n$，$\mathbf{C}$ 是 $m\times n$
• 更严谨地说，先对 $\mathbf{A}$ 做 transa 操作得到 $\mathrm{op}(\mathbf{A})$，它的维度要是 $m\times k$；同理 $\mathrm{op}(\mathbf{B})$ 必须是 $k\times n$，才可相乘并产生 $m\times n$ 的结果，最后再加到 $\mathbf{C}$ 上

三、$\mathrm{op}(\mathbf{A})$、$\mathrm{op}(\mathbf{B})$ 的三种取值

$$\mathrm{op}(\mathbf{A})=\{\begin{array}{ll}\mathbf{A}& \text{if transa}==\text{CUBLAS\_OP\_N}\\ {\mathbf{A}}^{\mathbf{T}}& \text{if transa}==\text{CUBLAS\_OP\_T}\\ {\mathbf{A}}^{\mathbf{H}}& \text{if transa}==\text{CUBLAS\_OP\_C}\end{array}$$

对于矩阵 $\mathbf{A}$，根据 transa 参数有三种可能：

transa 值	操作 $\mathrm{op}(\mathbf{A})$	物理含义
CUBLAS_OP_N	$\mathrm{op}(\mathbf{A})=\mathbf{A}$	不转置
CUBLAS_OP_T	$\mathrm{op}(\mathbf{A})={\mathbf{A}}^T$	转置
CUBLAS_OP_C	$\mathrm{op}(\mathbf{A})={\mathbf{A}}^H$	共轭转置（复数共轭再转置）

对 $\mathbf{B}$ 则同理，取决于 transb

四、总结

• 公式 $\mathbf{C}=\alpha \mathrm{op}(\mathbf{A})\mathrm{op}(\mathbf{B})+\beta \mathbf{C}$ 是所有 cuBLAS GEMM 调用的核心
• 通过 transa/transb 灵活指定矩阵转置类型，配合列主序存储和 leading dimensions（lad/ldb/ldc），即可对任意形状的矩阵做高效乘加

4.2 参数详解

参数	类型	输入/输出	含义
handle	cublasHandle_t	输入	cuBLAS 上下文句柄，由 cublasCreate() 创建
transa	cublasOperation_t	输入	矩阵 A 的变换：CUBLAS_OP_N（不转置），CUBLAS_OP_T（转置），CUBLAS_OP_C（共轭转置）
transb	cublasOperation_t	输入	矩阵 B 的变换，意义同上
m, n, k	int	输入	对应于计算 C = α·op(A)·op(B) + β·C 中 C(m×n) = A(m×k) × B(k×n)
alpha, beta	const void*	输入	缩放因子 α、β 指针，可位于主机或设备端（受 cublasPointerMode 控制）
A, B, C	const void*/void*	输入/输出	指向设备内存中 A、B、C 矩阵的指针
Atype, Btype, Ctype	cudaDataType_t	输入	矩阵 A、B、C 的数据类型枚举（如 CUDA_R_16F、CUDA_R_32F、CUDA_R_8I 等）
lda, ldb, ldc	int	输入	A、B、C 的“leading dimension”，即存储时的步长
computeType	cublasComputeType_t	输入	指定计算精度类型（如 CUBLAS_COMPUTE_32F、CUBLAS_COMPUTE_64F、CUBLAS_COMPUTE_32I 等）
algo	cublasGemmAlgo_t	输入	算法枚举（如 CUBLAS_GEMM_DEFAULT、CUBLAS_GEMM_DEFAULT_TENSOR_OP、其他内部 ID）

4.3 数据类型与计算模式支持

cublasGemmEx 支持灵活的输入/输出、缩放、计算类型组合。以下为主要组合示例（完整表见文档）：

computeType	缩放因子类型	A/B 数据类型	C 输出类型
CUBLAS_COMPUTE_32F / ..._PEDANTIC	CUDA_R_32F	CUDA_R_16F/16BF/8I/32F	CUDA_R_16F/32F
CUBLAS_COMPUTE_16F / ..._PEDANTIC	CUDA_R_16F	CUDA_R_16F	CUDA_R_16F
CUBLAS_COMPUTE_32I / ..._PEDANTIC	CUDA_R_32I	CUDA_R_8I	CUDA_R_32I
CUBLAS_COMPUTE_64F / ..._PEDANTIC	CUDA_R_64F	CUDA_R_64F	CUDA_R_64F
CUBLAS_COMPUTE_32F_FAST_16F / ..._TF32 等	CUDA_R_32F	CUDA_R_16F / TF32	CUDA_R_32F

4.4 使用示例

// 假设 A, B, C 已在设备上分配并填充数据且 A, B 为 row-major 存储
cublasHandle_t handle;
cublasCreate(&handle);
cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST);

int m = 512, n = 512, k = 512;
int lda = k, ldb = n, ldc = n;
float alpha = 1.0f, beta = 0.0f;

// 使用半精度输入、单精度累加、半精度输出
cublasGemmEx(handle,
    CUBLAS_OP_N, CUBLAS_OP_N,
    n, m, k,
    &alpha,
    B, CUDA_R_16F, ldb,
    A, CUDA_R_16F, lda,
    &beta,
    C, CUDA_R_16F, ldc,
    CUBLAS_COMPUTE_32F,              // 计算精度
    CUBLAS_GEMM_DEFAULT_TENSOR_OP    // 算法选择
);

cublasDestroy(handle);

4.5 常见注意事项

1. 硬件要求：仅支持 Compute Capability ≥ 5.0 的 GPU，否则返回 CUBLAS_STATUS_ARCH_MISMATCH

2. 对齐要求：使用 FAST 或 int32 模式时，A/B 指针需满足 4 或 16 字节对齐

3. 缩放因子指针模式：可由 cublasSetPointerMode(handle, ...) 切换主机/设备指针模式

4. 错误检查：返回值类型为 cublasStatus_t，应检查是否为 CUBLAS_STATUS_SUCCESS，并处理如 INVALID_VALUE、NOT_SUPPORTED 等错误

5. 性能调优：可通过 cublasSetMathMode 启用 TF32 或禁止低精度累积，并可尝试不同 algo 值获取最佳性能

更多细节大家可以查看 cuBLAS 官方文档：https://docs.nvidia.com/cuda/cublas/#cublasgemmex

4.6 补充：shape/stride

由于博主经常对主序、转置弄混，因此这里多说两句

下面我们结合 CuTe 中 shape/stride（张量维度/步幅） 的概念，以小例子说明行主序（row-major）和列主序（column-major）的区别，以及 BLAS 里 “leading dimension”（简称 ld*，也就是步幅）是如何计算的

Note：CuTe（CUDA Tensor Expressions） 是 CUTLASS 3.0 中引入的核心子库，用于 描述和操作线程与数据的层次化多维布局。它提供了 Layout 与 Tensor 两类 C++ CUDA 模板抽象，封装数据类型、形状、内存空间和布局，并自动生成复杂的索引计算逻辑，用户只需关心高层的张量运算描述，无需手动管理线程映射与内存布局

1. shape/stride 的概念

在 cute（或 Pytorch、NumPy）中，一个张量可以描述为：

shape = (D0, D1, …, Dn)
stride = (S0, S1, …, Sn)

对索引 (i0,i1,…,in)，它在一维连续内存上的偏移就是：

offset = i0*S0 + i1*S1 + … + in*Sn

• shape 决定了各维的大小
• stride 决定了跨维跳跃多少元素

2. 行主序 vs 列主序：shape/stride 表示

行主序（C 风格）

• 对于一个 $m\times k$ 矩阵，行主序最常见的 stride 是

$$\mathrm{shape}=(m,k),\mathrm{stride}=(k,1)$$

• 意味着：同一行内相邻元素偏移 1；不同行之间相邻行起点相差 $k$ 个元素

列主序（Fortran/BLAS 风格）

• 同样的 $m\times k$ 矩阵，用列主序就是

$$\mathrm{shape}=(m,k),\mathrm{stride}=(1,m)$$

• 意味着：同一列内相邻元素偏移 1；不同列之间起点相差 $m$ 个元素

3. 数值例子：$2\times 3$ 矩阵

假设矩阵内容按 “逻辑” 顺序编号为：

$$\left[\begin{array}{ccc}{a}_{00}& a_{01}& a_{02}\\ a_{10}& a_{11}& a_{12}\end{array}\right]$$

内存存储：

索引 $(i,j)$	行主序 offset=i·3+j	列主序 offset=i·1+j·2
(0,0)	0·3+0 = 0	0·1+0·2 = 0
(0,1)	0·3+1 = 1	0·1+1·2 = 2
(0,2)	0·3+2 = 2	0·1+2·2 = 4
(1,0)	1·3+0 = 3	1·1+0·2 = 1
(1,1)	1·3+1 = 4	1·1+1·2 = 3
(1,2)	1·3+2 = 5	1·1+2·2 = 5

• 行主序内存序列：

$$[a_{00},a_{01},a_{02},a_{10},a_{11},a_{12}]$$

• 列主序内存序列：

$$[a_{00},a_{10},a_{01},a_{11},a_{02},a_{12}]$$

4. shape/stride 等价示例

通过上面的分析，我们知道其实 shape $(m,k)$-stride $(k,1)$ 跟 shape $(k,m)$-stride $(1,k)$ 在内存上是一模一样的

• $\mathrm{shape}=(m,k),\mathrm{stride}=(k,1)$：按 row-major 存储
• $\mathrm{shape}=(k,m),\mathrm{stride}=(1,k)$：按 column-major 对 $(k\times m)$ 矩阵存储

实际上它们底层都是一段连续长度 $m\cdot k$ 的内存，且偏移公式都退化为：

$$\mathrm{offset}=i\cdot k+j⟺\mathrm{offset}=i+j\cdot k$$

只是前者视为 “第 0 维长 $k$”，后者视为 “第 1 维跨越 $k$”，看起来形状不同、stride 不同，内存排列却一致

5. BLAS 里的 “leading dimension”（lad/ldb/ldc）

在 cuBLAS、Fortan BLAS 中：

• 默认都是列主序
• 对于 $\mathrm{shape}=(m,k)$ 的矩阵 A，内存 $\mathrm{stride}=(1,\mathrm{lda})$
• 第一维（行）跨 1
• 第二维（列）跨 lda

6. 总结

1. shape/stride 刚好抓住了 “内存映射” 的本质：

• 行主序 $⟹$ stride = (列数, 1)
• 列主序 $⟹$ stride = (1, 行数)

2. shape $(m,k)$-stride $(k,1)$ 与 shape $(k,m)$-stride $(1,k)$ 虽然 “看” 成不同的维度，但内存布局完全相同

3. 在 cuBLAS 中，lda、ldb、ldc 就是列主序下的第 2 维 stride，通常取各自矩阵的行数

5.【cublas】cublasGemmEx 使用

有了上面的基础后，我们回过头再来看看 cublasGemmEx() 接口的使用

这边主要是通过阅读 UP 主推荐的 有关CUBLAS中的矩阵乘法函数 这篇文章来学习，强烈建议大家阅读原文

Note：以下内容均 Copy 自 有关CUBLAS中的矩阵乘法函数

cublasGemmEx() 接口定义如下：

cublasStatus_t cublasGemmEx(cublasHandle_t handle,
                           cublasOperation_t transa,
                           cublasOperation_t transb,
                           int m,
                           int n,
                           int k,
                           const void     *alpha,
                           const void     *A,
                           cudaDataType   Atype,
                           int lda,
                           const void     *B,
                           cudaDataType   Btype,
                           int ldb,
                           const void     *beta,
                           void           *C,
                           cudaDataType   Ctype,
                           int ldc,
                           cudaDataType   computeType,
                           cublasGemmAlgo_t algo)

测试用例如下：

$$A=\left(\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\end{array}\right),B=\left(\begin{array}{cccc}1& 2& 3& 4\\ 5& 6& 7& 8\\ 9& 10& 11& 12\end{array}\right),C=AB=\left(\begin{array}{cccc}38& 44& 50& 56\\ 83& 98& 113& 128\end{array}\right)$$

其中维度 $m=2,n=4,k=3$，即 $C_{m\times n}=A_{m\times k}{B}_{k\times n}⟹C_{2\times 4}=A_{2\times 3}=B_{3\times 4}$

5.1 错误示范

直接将 A、B 放入函数中，转置参数均取 CUBLAS_OP_N，lda、ldb、ldc 正常地取各矩阵的行数

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, *A, CUDA_R_16F, m, *B, CUDA_R_16F, k, &beta, *C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

cuBLAS 计算过程如下：

其中绿色的虚线框代表的是没有进入到 cuBLAS 的计算，红色的虚线框代表的是 cuBLAS 内部的计算

当把 C/C++ 中行优先保存的二维数组 $A,B$ 输入 cuBLAS 时，cuBLAS 自动将输入的 $A,B$ 理解为列优先存储，而且保持 $A,B$ 的行、列数保持不变，即将其当作图中的 $A_1,B_1$，注意这里的 $A_1,B_1$ 并不等于 $A,B$ 的转置

接着内部调用 tensor core 等硬件资源计算乘积 $C_1$ 作为输出，最后将计算的结果通过 cudaMemcpy 从设备端（Device）复制回主机端（Host），但主机端（C/C++）会将输出 $C_1$ 当成是行优先的矩阵，所以我们最终在 Host 上拿到的矩阵乘积结果就变成了 $C_2$ 的模样，显然这样的计算结果显然不是我们期望的 $C$

Note：图中的 “列优先” 和 “行优先” 分别代表进入和退出 cuBLAS 时发生的强制转换，不要理解为其他意思

关于主序（行主序和列主序）和主维（ld*）这边博主再啰嗦两句

对于矩阵 $A=\left(\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\end{array}\right)$ 如果是不同的主序，存储时的内存排布是不同的

• 行主序（Row-Major）存储时
  • 内存排布：$1,2,3,4,5,6$
  • 每一行内元素是连续的，要从一行跳到下一行，就要跨过 “列数” 个元素
  • 因此它的 leading dimension 主维 $\mathrm{lda}=\mathrm{列数}=3$
• 列主序（Col-Major）存储时
  • 内存排布：$1,4,2,5,3,6$
  • 每一列内元素是连续的，要从一列跳到下一列，就要跨过 “行数” 个元素
  • 因此它的 leading dimension 主维 $\mathrm{lda}=\mathrm{行数}=2$

下面我们可以简单验证一下下标映射

• Row-Major 下标 $(i,j)$ 的线性偏移

$$\mathrm{offset}=i\times \mathrm{lda}+j=i\times 3+j$$

例如 $A(1,1)=5$ 在内存中的线性索引（index）就是 $1\cdot 3+1=4$，与连续数组 $1,2,3,4,5,6$ 中第 4 个元素（从 0 开始计数）吻合

• Col-Major 下标 $(i,j)$ 的线性偏移

$$\mathrm{offset}=i+j\times \mathrm{lda}=i+j\times 2$$

例如 $A(1,1)=5$ 在列主序的偏移是 $1+1\cdot 2=3$，在 $1,4,2,5,3,6$ 中却是第 3 个元素

Note：这里需要特别注意，行列主序的转换并不等于转置

这是关于主序的差异，而主维（leading dimension, ld）的作用就是 告诉库函数在一段连续的一维内存里 “跨行” 或 “跨列” 要跳过多少元素，从而正确地把它当作二维矩阵来索引

假设二维矩阵 $A$ 在内存中数据排布是 $1,2,3,4,5,6$，如果矩阵 $A$ 是列主序且 lda 主维等于 2，则二维矩阵 $A$ 的形状是：

$$A_{2\times 3}=\left(\begin{array}{ccc}1& 3& 5\\ 2& 4& 6\end{array}\right)$$

如果矩阵 $A$ 是列主序但 lda 主维等于 3，则二维矩阵 $A$ 的形状是：

$$A_{3\times 2}=\left(\begin{array}{cc}1& 4\\ 2& 5\\ 3& 6\end{array}\right)$$

我们也可以看出来对于同一段内存，不同的 ld 值会映射成不同形状的矩阵，这是关于主维的一些理解

5.2 正确过程

一般教程上的调用过程利用了转置的性质即 $AB=(B^T{A}^T{)}^T$。我们把一个行优先的矩阵看作列优先的同时，交换其行、列数，其结果等价于得到该矩阵的转置，列优先转行优先的原理相同（为什么呢？🤔）

博主对这句话想了很久，还是不能理解，下面我们来简要解释下：(from ChatGPT)

这句话本质上是对 “行主序 ↔ 列主序” 两种内存布局下同一段物理数据的 不同解读，这种解读上的 “互换行列” 恰好等价于数学意义上的矩阵转置。下面分几步剖析它为什么成立，并结合 shape-stride 来看一个具体例子

1. 行主序 vs 列主序 的下标映射

• 行主序（Row-Major）：元素 $(i,j)$ 的线性内存下标

$${\mathrm{offset}}_{\mathrm{row}}(i,j)=i\times \mathrm{ida}+j$$

其中 $\mathrm{lda}$（leading dimension）通常等于列数 $k$

• 列主序（Col-Major）：元素 $(i,j)$ 的线性内存下标

$${\mathrm{offset}}_{\mathrm{col}}(i,j)=i+j\times \mathrm{ida}$$

这里的 $\mathrm{lda}$ 通常等于行数 $m$

2. 互换行列并看作另一种布局，等价于转置

假设我们有一个 $m\times k$ 的矩阵 $A$，用 行主序 存储，$\mathrm{lda}=k$。此时

$$A[i,j]\mathrm{的地址}=i\cdot k+j$$

如果我们不移动数据，把这段内存当作一个 $k\times m$ 的 列主序 矩阵来读，令新的下标为 $(r,c)$，${\mathrm{lda}}^{\prime }=m$。此时

$${\mathrm{offset}}_{\mathrm{col}}(r,c)=r+c\cdot m$$

为了地址一致，我们对应地让

$$r=j,c=i$$

也就是说，新矩阵在位置 $(r,c)=(j,i)$ 处的值，正好是原矩阵 $A$ 在 $(i,j)$ 的元素。这正是转置矩阵 $A^T$ 的定义：

$$(A^T)[j,i]=A[i,j]$$

由此可以得出一个结论：把内存不变地从 “$m\times k$ 行主序” 视作 “$k\times m$ 列主序”，程序读出来的值就等于原矩阵的转置

理解了之后我们就可以在调用该函数的时候，先放入 B 交换参数 k 和 n 的位置，再放入 A 并交换参数 m 和 k 的位置，这样就顺理成章得到了结果 C（所有转换由 cuBLAS 完成，不需要手工调整数组）

// old
// cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, *A, CUDA_R_16F, m, *B, CUDA_R_16F, k, &beta, *C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

// new
cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, *B, CUDA_R_16F, n, *A, CUDA_R_16F, k, &beta, *C, CUDA_R_16F, n, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

cuBLAS 计算过程如下：

cuBLAS 读取 B 和 A 的时候虽然进行了重排，但是我们交换了行列数，因此得到的是原矩阵的转置，即 $C_1=B^T{A}^T$，接着交换 $C_1$ 的行列数，最终得到的就是 $C_1$ 的转置即正常的结果 $C=C_1^T=(B^T{A}^T{)}^T=AB$

5.3 手动转置

如果我们不想使用 5.2 中的奇淫技巧，而是希望按顺序将 A、B 输入 cuBLAS 中运算，我们可以先尝试手动将 A、B 转化为列优先存储，然后再调用该函数，最后将输出的 C 转换回行优先即可。

转换示例代码如下：

// 行优先转列优先，只改变存储方式，不改变矩阵尺寸
void rowToCol(float* a, int row, int col){

    float* temp = (float*)malloc(row * col * sizeof(float));    // 存储列优先存储的临时数组

    for(int i = 0; i < row * col; ++i){
        temp[i] = a[i / row + i % row * col];    // 找出列优先的第 i 个位置对应行优先坐标进行赋值
    }

    for(int i = 0; i < row * col; ++i){
        a[i] = temp[i];    // 覆盖原数组
    }

    free(temp);
    return;
}

// 列优先转行优先
void colToRow(float* a, int row, int col){

    float* temp = (float*)malloc(row * col * sizeof(float));

    for(int i = 0; i < row * col; ++i){
        temp[i] = a[i / col + i % col * row];    // 找出行优先的第 i 个位置对应列优先坐标进行赋值
    }

    for(int i = 0; i < row * col; ++i){
        a[i] = temp[i];
    }

    free(temp);
    return;
}

这样一来，我们的调用参数如下（跟 1 中参数位置一模一样）：

rowToCol(h_A, m ,k);
rowToCol(h_B, k, n);

...// 准备工作

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, m, n, k, &alpha, d_A, CUDA_R_16F, m, d_B, CUDA_R_16F, k, &beta, d_C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

...// 回收工作

colToRow(h_CUBLAS, m, n);

注意上面的转换函数中我们并没有交换矩阵 A 的行数、列数，所以在上面的调用中仍然有 $m=2,n=4,k=3$

cuBLAS 计算过程如下：

注意经过 rowToCol 处理得到的矩阵，以 $A_1$ 为例，它相当于用列优先的方式（竖着走）把 A 遍历，得到的轨迹按行优先的方式存放（横着放）。可见矩阵 A、B、C 一波三折，但是最终获得了我们期望的结果

5.4 使用 cuBLAS 自动转置

有了 3 的基础，我们就能很好的理解参数 cublasOperation_t transa, cublasOperation transb 的作用了。这两个参数就是在问，是否要在 cuBLAS 内部完成上面的 rowToCol 过程，如果我们选择参数 CUBLAS_OP_T，那么不再需要手动地将 A、B 进行转换（但是仍需要对输出 C 进行转换），直接上代码：

cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_T, m, n, k, &alpha, d_A, CUDA_R_16F, k, d_B, CUDA_R_16F, n, &beta, d_C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

...// 回收工作

colToRow(h_CUBLAS, m, n);

cuBLAS 计算过程如下：

我们发现主维数 lda、ldb 已经用上了，它们分别等于 $A_2,B_2$ 的行数（在列优先存储中矩阵的行数是维度的第一个分量，故称主维）

其实图中可以省略 $\to A_2$ 这一中间步骤，但是为了突出 “转置” 的意味，我们把它画出来，再经过 CUBLAS_OP_T 处理，方便理解

最终我们获得了期望的结果 $C_3$，我们发现这种方式省去了对 A、B 的 rowToCol 操作，但是 $C_2$ 的 colToRow 操作还是逃不掉，没有 2 来的巧妙

5.5 其它组合

有了以上的说明，我们就可以对 A、B 的不同情况加以利用，得到下面两种组合，它们都能得到正确的结果：

// A 转 B 不转
// rowToCol(h_A, m ,k);
rowToCol(h_B, k, n);

...// 准备工作

cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, m, n, k, &alpha, d_A, CUDA_R_16F, k, d_B, CUDA_R_16F, k, &beta, d_C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

...// 回收工作

colToRow(h_CUBLAS, m, n);

// B 转 A 不转
rowToCol(h_A, m ,k);
// rowToCol(h_B, k, n);

...// 准备工作

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, m, n, k, &alpha, d_A, CUDA_R_16F, m, d_B, CUDA_R_16F, n, &beta, d_C, CUDA_R_16F, m, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

...// 回收工作

colToRow(h_CUBLAS, m, n);

5.6 主维（leading dimension）参数

大家可以思考下，如果我们想要计算矩阵 $A_{m\times k}\times B_{k\times n}=C_{m\times n}$，那 $m,n,k$ 三个参数已经固定了矩阵所有尺寸，为什么还要一组主维参数呢？🤔

看完了上面部分的分析后我们会发现，输入的矩阵 $A,B$ 在整个计算过程中会发生变形，包括行列优先变换和转置变换，所以需要一组参数告诉该函数，矩阵变形后应该具有什么样的尺寸，这就是主维参数的作用

当参数 transa 为 CUBLAS_OP_N 时，矩阵 $A$ 在 cuBLAS 中的尺寸实际为 $\mathrm{Ida}\times k$，此时要求 $\mathrm{Ida}\ge \max (1,m)$，否则该函数直接报错，输出全零的结果

当参数 transa 为 CUBLAS_OP_T 或 CUBLAS_OP_C 时，矩阵 $A$ 在 cuBLAS 中的尺寸实际为 $\mathrm{Ida}\times m$，此时要求 $\mathrm{Ida}\ge \max (1,k)$

transb 为 CUBLAS_OP_N 时，$B$ 尺寸为 $\mathrm{Idb}\times n$，要求 $\mathrm{Idb}\ge \max (1,k)$；transb 为 CUBLAS_OP_T 或 CUBLAS_OP_C 时，$B$ 尺寸为 $\mathrm{Idb}\times k$，此时要求 $\mathrm{Idb}\ge \max (1,n)$

$C$ 尺寸为 $\mathrm{Idc}\times n$，要求 $\mathrm{Idc}\ge \max (1,m)$

可见，是否需要该函数帮我们转置矩阵 $A$ 会影响 $A$ 在函数中的存储，而主维正是在这一过程中起作用的。特别的，如果按照一般教程中的方法来调用该函数，三个主维参数是不同特别处理的

上面的讲解中只要求了三个主维参数的下界，下面我们来尝试调高三个参数看看对于最终结果有什么样的变化

1. Idb + 1

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, d_B, CUDA_R_16F, n + 1, d_A, CUDA_R_16F, k, &beta, *C, CUDA_R_16F, n, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

cuBLAS 计算过程如下：

注意 $B_1$ 中多出来的部分被用 0 来填充了，计算 $C_1$ 的时候是依其尺寸来进行的，若 $C_1$ 的主维尺寸没有加 1，则最后一行不进行计算

2. Ida + 1

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, d_B, CUDA_R_16F, n, d_A, CUDA_R_16F, k + 1, &beta, *C, CUDA_R_16F, n, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

cuBLAS 计算过程如下：

注意 $A_1$ 在参与乘法时去掉了多出的行，其实应该是计算矩阵乘法时 $\mathrm{行乘长}=\min (B_1\mathrm{列数},A_1\mathrm{行数})=3$

3. Idc + 1

cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, d_B, CUDA_R_16F, n, d_A, CUDA_R_16F, k, &beta, *C, CUDA_R_16F, n + 1, CUBLAS_COMPUTE_16F, CUBLAS_GEMM_DEFAULT_TENSOR_OP)

cuBLAS 计算过程如下：

注意计算出来的 $C_1$ 自动补充了 0 行，但是在输出的时候连着 0 一起输出，并且把最后的两个元素给挤掉了

6.【cublas】baseline 算法修改

这里有个点需要大家注意，那就是如果我们按照第 3 小节中原始框架代码去执行程序并对比结果时，会发现我们自己实现的 hgemm 的结果和 baseline 的结果差异比较大，如下图所示：

那这其实并不是我们自己实现的 hgemm 代码存在问题，而是本身对比的 baseline 算法存在问题，下面我们就一起来看看

工具类以及对比的 baseline 算法的实现来自于 cuda_hgemm 仓库，具体代码在 common/tester.h 文件中，也就是调用的 cuBLAS 的 cublasGemmEx 接口，如下所示：

HGEMM_CHECK_CUBLAS_ERROR(cublasGemmEx(handle, CUBLAS_OP_T, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_16F, K, A,
                                        CUDA_R_16F, K, &beta, C, CUDA_R_16F, N, CUBLAS_COMPUTE_16F,
                                        CUBLAS_GEMM_DEFAULT_TENSOR_OP));  

首先我们的输入矩阵 A、B 默认都是行主序，且矩阵 A 维度是 $m\times k$，矩阵 B 维度是 $k\times n$，那上面这种调用方式最终结果其实并不符合我们的预期，我们可以举个简单的例子来看下

假设 $A=\left(\begin{array}{ccc}1& 2& 3\\ 4& 5& 6\end{array}\right)$，$B=\left(\begin{array}{cccc}1& 2& 3& 4\\ 5& 6& 7& 8\\ 9& 10& 11& 12\end{array}\right)$，$m=2,k=3,n=4$

则在上面这种调用方式下，cuBLAS 计算过程如下：

可以看到最终的结果与我们实际的期望有所偏差

那正确调用方式如下所示：

HGEMM_CHECK_CUBLAS_ERROR(cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, B, CUDA_R_16F, N, A,
                                        CUDA_R_16F, K, &beta, C, CUDA_R_16F, N, CUBLAS_COMPUTE_16F,
                                        CUBLAS_GEMM_DEFAULT_TENSOR_OP));

这种调用方式也正是我们在 5.2 小节中提到的利用转置的性质来完成 $C=AB=(B^T{A}^T{)}^T$ 矩阵乘法的运算

在这种调用方式下 cuBLAS 计算过程如下：

修改 cuBLAS 调用方式后再重新编译执行程序进行结果对比，会发现此时差异非常小，如下图所示：

可以看到二者结果接近，由于我们是半精度矩阵乘法，因此有些差异也正常

结语

这篇文章我们首先了解了下 tensor core，它是 NVIDIA GPU 上的一个硬件，可以用来实现混合精度计算并加速矩阵运算，CUDA 通过 CUDA C++ WMMA API 向外提供了 Tensor Core 在 Warp 级别上的计算操作支持

接着我们对开发环境和框架进行了搭建，其中具体的实现代码来自于 LeetCUDA 这个仓库，工具和测试类的实现来自于 cuda_hgemm 这个仓库

最后我们对 cuBLAS 中的接口 cublasGemmEx 进行了详细的分析，我们用这个接口作为 baseline 来衡量我们自己实现的 hgemm 是否存在问题。我们参考官方文档对接口参数和注意事项进行了讲解，并参考相关文章详细介绍了其使用方法，包括行主序、列主序区别、转置、主维等参数的理解和使用

篇幅原因，后续的 hgemm 具体实现的代码分析我们放在下篇文章中讲解，敬请期待🤗

下载链接

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

参考

• 【CUDA进阶】Tensor Core实战教程（已完结）
• 有关CUBLAS中的矩阵乘法函数
• https://github.com/xlite-dev/LeetCUDA
• https://github.com/Bruce-Lee-LY/cuda_hgemm
• https://docs.nvidia.com/cuda/cublas
• https://developer.nvidia.com/blog/programming-tensor-cores-cuda-9/
• https://chatgpt.com