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

cuda编程笔记（20）-- 混合精度计算

news 2025/9/17 10:08:52

背景：为什么要用混合精度

问题：现代深度学习模型越来越大，计算量巨大，显存占用高。
目标：
1. 降低显存占用 → FP16 (半精度) 占用 FP32 的一半。
2. 提升计算吞吐量 → NVIDIA GPU 上 Tensor Core 对 FP16 计算有专门加速。
3. 保持训练稳定性 → 梯度累加或缩放保持 FP32 精度的数值稳定。

所以出现了 混合精度训练 (Mixed Precision Training)：

前向、反向大部分用 FP16 计算。
梯度累加和权重更新用 FP32 以避免精度丢失。

精度介绍

在 C++ / CUDA 里：

float → 单精度浮点 (FP32)
- 1 bit 符号位
- 8 bit 阶码 (exponent)
- 23 bit 尾数 (mantissa)
- 大约能表示 7 位有效十进制数字
- 范围约：1e-38 ~ 1e+38

double → 双精度浮点 (FP64)
- 1 bit 符号位
- 11 bit 阶码
- 52 bit 尾数
- 大约能表示 15~16 位有效十进制数字
- 范围约：1e-308 ~ 1e+308

__half (CUDA 专有) → 半精度浮点 (FP16)
- 1 bit 符号位
- 5 bit 阶码
- 10 bit 尾数
- 大约能表示 3 位有效十进制数字
- 范围约：6e-8 ~ 6e+4（远比 FP32 窄）

所以：

平时的 float 就是 FP32，对应 cublasSgemm 里的那个 S。
double 就是 FP64，对应 cublasDgemm 里的 D。
half（或 __half）才是 FP16，要用 cublasGemmEx / Tensor Core 才能用起来。

Tensor Core

Tensor Core 是 GPU 上专用矩阵乘法单元，对 FP16-运算特别优化。
特点：
- 输入矩阵：FP16
- 累加寄存器：FP32
- 输出矩阵：FP16 或 FP32
速度：在 Volta 及以后 GPU，Tensor Core 可以实现 FP16 矩阵运算吞吐量比 FP32 高 8~16 倍。

混合精度训练的典型做法

前向传播：
- 输入、权重矩阵都用 FP16。
- GEMM 运算通过 Tensor Core 或 cublasGemmEx 完成。
反向传播：
- 梯度计算用 FP16 或 FP32。
- 需要用 FP32 累加以避免数值下溢。
梯度缩放 (loss scaling)：
- 因为 FP16 范围小，梯度可能非常小 → 下溢。
- 训练框架会把 loss 放大一个系数（scale），梯度也放大，再更新权重。
权重更新：
- 梯度累加到 FP32 master weight（主权重）。
- 更新完，再转换成 FP16 存储在显存里。

和普通做法的区别

在之前的文章，我们也实现过LeNet的训练流程，这两有什么区别呢？

普通 FP32 训练（之前做的）

forward/backward 都在 FP32。
反向传播结果：grad_weight (FP32)、grad_bias (FP32)。
update：直接 weight -= lr * grad_weight。
所有数据类型一致，不需要额外处理。

混合精度（FP16 参与）

为了省显存，模型的参数（权重、激活、梯度）存在显存里一般是 FP16。
但是！FP16 表示范围小，容易溢出/下溢。
所以框架（比如 Apex AMP、PyTorch AMP）会偷偷维护一份 FP32 master weights：
- master_weight (FP32)：更新时真正使用的参数副本。
- model_weight (FP16)：forward/backward 时用的低精度参数。
backward 算法会得到 FP16 梯度（因为输入、权重都是 FP16）。
框架会把这个 FP16 梯度 cast 成 FP32，然后：
```
master_grad += grad.to(FP32)
master_weight -= lr * master_grad
```
更新完 master_weight 后，再把它 cast 回 FP16，写到 model_weight 里，供下一次 forward 用。

为什么会有 `+=` 的「梯度累加」？

在最标准的 mini-batch SGD 里：

你取一个 batch（比如 128 张图片）。
前向、反向传播一次，得到这个 batch 的 平均梯度。
直接 param -= lr * grad 更新权重。
这里确实不需要额外的梯度累加，因为一次反传就把整个 batch 的梯度算完了。

有可能这样：

梯度累积 (Gradient Accumulation)

有时候显存不够，你不能一次性放下 128 张图片。
你就把 batch 拆成 4 次，每次 32 张：
这样叫 累积梯度模拟大 batch。
这样每次计算的一次梯度是FP16的，需要转化成FP32存入主梯度
写成+=,同一套代码既能支持「单 batch 更新」也能支持「梯度累积」。

cublas的混合精度

在 CUDA 里有两种常见的 “half”：

__half → CUDA 内置的半精度类型，定义在 <cuda_fp16.h> 里。
half → 这是 __half 的 typedef，但默认 C++ 并不认识，需要 #include <cuda_fp16.h>。

更关键的是：

__half / half 没有定义 operator<<，所以 std::cout 输出会报错。

`float2half` 和 `half2float`

这两个函数定义在 <cuda_fp16.h>，作用是 float 和 half 的互转。

__float2half(float x)
→ 把 32 位浮点数 转换为 16 位浮点数（__half）。
例如：
```
__half h = __float2half(3.14f);
```
__half2float(__half h)
→ 把 16 位浮点数 转换为 32 位浮点数。
例如：
```
float f = __half2float(h);
```

原因：

在 CUDA 里，__half 只是 16 bit 容器，本身没有太多算术运算支持。
大多数运算要么用专门的 half intrinsic，要么转回 float 计算。
打印时更是必须转 float，因为 std::cout 不认识 __half。

`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);

参数解释（和普通 cublasSgemm 对比）：

Atype / Btype / Ctype
→ 输入、输出矩阵的数据类型，比如 CUDA_R_16F（FP16）、CUDA_R_32F（FP32）。
computeType
→ 计算时内部使用的数据类型
- CUDA_R_16F → 直接 FP16 运算（精度低，几乎不用）
- CUDA_R_32F → FP16 输入，FP32 累加（主流，Tensor Core）
algo
→ 算法选择，常见 CUBLAS_GEMM_DEFAULT_TENSOR_OP（启用 Tensor Core）。

也就是说，cublasGemmEx 的最大特点：
允许输入/输出/计算精度不同，这就是“混合精度”的入口。

如果是FP16*FP16,最后存在FP32里：

从 global memory 里把 FP16 加载到寄存器。
Tensor Core 单元做 FP16 × FP16 乘法，但累加器（accumulator）是 FP32。
- 每次乘法得到的结果先提升为 float，再加到 FP32 累加器里。
- 这样避免了 FP16 精度损失在多次加法里不断放大。
最终结果保存在 FP32 内存里。

如果结果也设置为FP16呢？

内部依旧 FP32 累加（数值稳定）。
但是 最后写回内存时强制 cast 到 FP16。
你得到的结果会有 更大舍入误差，尤其是大 batch、累加次数多时。

示例代码

#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <iostream>
#include <vector>int main() {const int batch=2,in_features=3,out_features=4;cublasHandle_t handle;cublasCreate(&handle);std::vector<__half> h_A{__float2half(1.f),__float2half(2.f),__float2half(3.f),__float2half(4.f),__float2half(5.f),__float2half(6.f)};std::vector<__half> h_B{__float2half(1.f),__float2half(2.f),__float2half(3.f),__float2half(4.f),__float2half(5.f),__float2half(6.f),__float2half(7.f),__float2half(8.f)};std::vector<float> h_C(in_features*out_features,0); // 用 float 保存结果__half *d_A, *d_B;float *d_C;cudaMalloc(&d_A, batch*in_features*sizeof(__half));cudaMalloc(&d_B, batch*out_features*sizeof(__half));cudaMalloc(&d_C, in_features*out_features*sizeof(float));cudaMemcpy(d_A, h_A.data(), batch*in_features*sizeof(__half), cudaMemcpyHostToDevice);cudaMemcpy(d_B, h_B.data(), batch*out_features*sizeof(__half), cudaMemcpyHostToDevice);const float alpha = 1.0f;const float beta = 0.0f;cublasGemmEx(handle,CUBLAS_OP_N, CUBLAS_OP_T,out_features, in_features, batch,&alpha,d_B, CUDA_R_16F, out_features,d_A, CUDA_R_16F, in_features,&beta,d_C, CUDA_R_32F, out_features,   // 输出用 floatCUDA_R_32F, CUBLAS_GEMM_DEFAULT_TENSOR_OP);cudaMemcpy(h_C.data(), d_C, in_features*out_features*sizeof(float), cudaMemcpyDeviceToHost);for (int i = 0; i < in_features*out_features; i++) {std::cout << h_C[i] << " ";}std::cout << std::endl;cudaFree(d_A);cudaFree(d_B);cudaFree(d_C);cublasDestroy(handle);
}

cuBLASLt的混合精度

cublasLtMatmul ——这是 cublasLt (Lightweight cuBLAS) API 的核心函数。

它比 cublasGemmEx 灵活得多，是 NVIDIA 专门为 Tensor Core + 混合精度 做的高性能矩阵乘法接口。

所有新出现的类型和函数都属于#include<cublasLt.h>头文件；链接要加上cublasLt

cublasStatus_t cublasLtMatmul(cublasLtHandle_t                        lightHandle,cublasLtMatmulDesc_t                    operationDesc,const void*                             alpha,const void*                             A,const cublasLtMatrixLayout_t            Adesc,const void*                             B,const cublasLtMatrixLayout_t            Bdesc,const void*                             beta,const void*                             C,const cublasLtMatrixLayout_t            Cdesc,void*                                   D,const cublasLtMatrixLayout_t            Ddesc,const cublasLtMatmulPreference_t        preference,void*                                   workspace,size_t                                  workspaceSize,cudaStream_t                            stream);

句柄
- cublasLtHandle_t lightHandle
  cuBLASLt 的上下文句柄，类似 cublasHandle_t。
操作描述符
- cublasLtMatmulDesc_t operationDesc
  指定运算类型，比如：
  - 矩阵转置/共轭转置
  - 数据类型（半精度/单精度/BF16/INT8 等）
  - 计算模式（是否用 Tensor Core）
  - 算法选择（默认、启发式、指定 kernel）
矩阵数据和描述符
- A, Adesc：矩阵 A + 布局描述符
- B, Bdesc：矩阵 B + 布局描述符
- C, Cdesc：输入矩阵 C（类似 GEMM 的 beta * C）
- D, Ddesc：输出矩阵 D（结果会写到 D）
⚠️ 注意：D 可以和 C 相同，也可以不同，这让它支持更复杂的 fused operations（例如 GEMM + bias + activation）。
缩放系数
- alpha、beta 和 GEMM 一样，分别缩放 A*B 和 C。
性能调优
- cublasLtMatmulPreference_t preference
  可以设置：
  - 是否允许 Tensor Core
  - 是否倾向于占用少内存
  - 是否倾向于低延迟（短 kernel）
  - 工作空间大小上限
- workspace / workspaceSize
  给 cuBLASLt 额外临时显存空间，用于算法优化（比如更快的 tile 分块）。
流
- cudaStream_t stream
  可以放进任意 CUDA 流里执行，支持异步。

cublasLtHandle_t

作用：cuBLASLt 的全局上下文，类似 cublasHandle_t。
创建/销毁：

cublasLtHandle_t ltHandle;
cublasLtCreate(&ltHandle);
cublasLtDestroy(ltHandle);

cublasLtMatmulDesc_t

型：描述符 (descriptor)
作用：描述矩阵乘法操作的属性，比如：
- 矩阵 A、B 是否转置
- 数据计算类型（FP16、BF16、TF32、FP32、INT8）
- 是否使用 Tensor Core
创建/销毁：

cublasLtMatmulDesc_t matmulDesc;
cublasLtMatmulDescCreate(&matmulDesc, CUDA_R_32F); // 计算精度 FP32
cublasLtMatmulDescDestroy(matmulDesc);

设置属性（类似键值对接口），所有属性都是这么设置的：

cublasOperation_t transa = CUBLAS_OP_N;
cublasLtMatmulDescSetAttribute(matmulDesc,CUBLASLT_MATMUL_DESC_TRANSB,&transa, sizeof(transa));

属性	数据类型	说明
`CUBLASLT_MATMUL_DESC_TRANSA`	`cublasOperation_t`	矩阵 A 是否转置（`CUBLAS_OP_N`/`CUBLAS_OP_T`）
`CUBLASLT_MATMUL_DESC_TRANSB`	`cublasOperation_t`	矩阵 B 是否转置
`CUBLASLT_MATMUL_DESC_EPILOGUE`	`cublasLtEpilogue_t`	后处理操作，比如加 bias/激活（`CUBLASLT_EPILOGUE_DEFAULT`、`CUBLASLT_EPILOGUE_BIAS` 等）
`CUBLASLT_MATMUL_DESC_SCALE_TYPE`	`cudaDataType`	alpha/beta 的数据类型，通常 `CUDA_R_32F`
`CUBLASLT_MATMUL_DESC_COMPUTE_TYPE`	`cudaDataType`	内部计算精度（FP16 输入可以 FP32 累加）
`CUBLASLT_MATMUL_DESC_TILE_ID`	`int`	选择特定 tile / Tensor Core 算法的内部实现

`cublasLtMatrixLayout_t`

类型：描述符
作用：描述矩阵的内存布局，包括：
- 行数、列数
- leading dimension (lda)
- 批大小（batched GEMM）
- stride（batch 间偏移）
- 数据类型（FP16/FP32/INT8）

cublasStatus_t cublasLtMatrixLayoutCreate(cublasLtMatrixLayout_t *matLayout,cudaDataType dataType,uint64_t rows,        // 矩阵行数uint64_t cols,        // 矩阵列数int64_t ld            // leading dimension
);

ld: leading dimension，指的是 相邻两行（column-major）或相邻两列（row-major）之间的跨度。

如果是 列优先（默认情况），ld = rows。
如果设置了 row-major (CUBLASLT_ORDER_ROW)，ld = cols。

属性	数据类型	说明
`CUBLASLT_MATRIX_LAYOUT_ORDER`	`cublasLtOrder_t`	行主序 / 列主序（`CUBLASLT_ORDER_ROW` / `CUBLASLT_ORDER_COL`）
`CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT`	`int`	batch 大小
`CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET`	`long long`	每个 batch 在内存里的偏移量（元素数量，不是字节数）
`CUBLASLT_MATRIX_LAYOUT_LD`	`int`	leading dimension (`lda`)
`CUBLASLT_MATRIX_LAYOUT_DATA_TYPE`	`cudaDataType`	数据类型，例如 `CUDA_R_16F`

batch GEMM 必须设置 batch count 和 stride。
row/col order 属性可以控制是按行优先还是列优先访问。
leading dimension 是矩阵“内存间隔”，对 row/col order 都非常关键。

cublasLtMatmulPreference

告诉 cuBLASLt 你想在执行 matmul 时偏向哪种策略，例如内存限制、速度优先等。

属性	数据类型	说明
`CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES`	`size_t`	分配给内部优化算法的工作空间大小（显存）
`CUBLASLT_MATMUL_PREF_MAX_ALGO_COUNT`	`int`	查找候选算法数量的上限
`CUBLASLT_MATMUL_PREF_SEARCH_MODE`	`cublasLtAlgoSearch_t`	算法搜索模式： `CUBLASLT_ALGO_SEARCH_DEFAULT`（快速） `CUBLASLT_ALGO_SEARCH_EXHAUSTIVE`（穷尽搜索，找最快）
`CUBLASLT_MATMUL_PREF_EPSILON`	`float`	精度容忍度（浮点比较）

preference 并不改变数学计算结果，只影响选择哪种算法 / tile / Tensor Core 组合。
如果显存紧张，可以调小 MAX_WORKSPACE_BYTES。
搜索模式决定了 cuBLASLt 是否遍历所有算法找到最快的。

示例代码

#include <cublas_v2.h>
#include<cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <iostream>
#include <vector>int main() {const int batch=2,in_features=3,out_features=4;cublasLtHandle_t ltHandle;cublasLtCreate(&ltHandle);//2*3std::vector<__half> h_A{__float2half(1.f),__float2half(2.f),__float2half(3.f),__float2half(4.f),__float2half(5.f),__float2half(6.f)};//2*4std::vector<__half> h_B{__float2half(1.f),__float2half(2.f),__float2half(3.f),__float2half(4.f),__float2half(5.f),__float2half(6.f),__float2half(7.f),__float2half(8.f)};//C=AT*B 3*4std::vector<float> h_C(in_features*out_features,0); // 用 float 保存结果__half *d_A, *d_B;float *d_C;cudaMalloc(&d_A, batch*in_features*sizeof(__half));cudaMalloc(&d_B, batch*out_features*sizeof(__half));cudaMalloc(&d_C, in_features*out_features*sizeof(float));cudaMemcpy(d_A, h_A.data(), batch*in_features*sizeof(__half), cudaMemcpyHostToDevice);cudaMemcpy(d_B, h_B.data(), batch*out_features*sizeof(__half), cudaMemcpyHostToDevice);const float alpha = 1.0f;const float beta = 0.0f;// 创建 matmul 描述符cublasLtMatmulDesc_t matmulDesc = nullptr;cublasLtMatrixLayout_t layoutA = nullptr, layoutB = nullptr, layoutC = nullptr;cublasLtMatmulDescCreate(&matmulDesc,CUBLAS_COMPUTE_32F,CUDA_R_32F);//计算精度32F，系数精度32F// 设置 row-major,我们要用行优先，所以ld要传列数cublasLtMatrixLayoutCreate(&layoutA, CUDA_R_16F, batch, in_features, in_features);  // A: batch x in_featurescublasLtMatrixLayoutCreate(&layoutB, CUDA_R_16F, batch, out_features, out_features); // B: batch x out_featurescublasLtMatrixLayoutCreate(&layoutC, CUDA_R_32F, in_features, out_features, out_features); // C: in_features x out_features//设置行优先cublasLtOrder_t order=CUBLASLT_ORDER_ROW;cublasLtMatrixLayoutSetAttribute(layoutA,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(cublasLtOrder_t));cublasLtMatrixLayoutSetAttribute(layoutB,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(cublasLtOrder_t));cublasLtMatrixLayoutSetAttribute(layoutC,CUBLASLT_MATRIX_LAYOUT_ORDER,&order,sizeof(cublasLtOrder_t));//设置A要转置cublasOperation_t transA = CUBLAS_OP_T;cublasOperation_t transB = CUBLAS_OP_N;cublasLtMatmulDescSetAttribute(matmulDesc,CUBLASLT_MATMUL_DESC_TRANSA, &transA, sizeof(transA));cublasLtMatmulDescSetAttribute(matmulDesc,CUBLASLT_MATMUL_DESC_TRANSB, &transB, sizeof(transB));cublasLtMatmul(ltHandle,matmulDesc,&alpha,d_A, layoutA,d_B, layoutB,&beta,d_C, layoutC,d_C, layoutC,  // 输出 buffer，可复用nullptr,nullptr, 0,    // 没有preference、workspace0);            // streamcudaMemcpy(h_C.data(), d_C, in_features*out_features*sizeof(float), cudaMemcpyDeviceToHost);for (int i = 0; i < in_features*out_features; i++) {std::cout << h_C[i] << " ";}std::cout << std::endl;cublasLtMatmulDescDestroy(matmulDesc);cublasLtMatrixLayoutDestroy(layoutA);cublasLtMatrixLayoutDestroy(layoutB);cublasLtMatrixLayoutDestroy(layoutC);cudaFree(d_A);cudaFree(d_B);cudaFree(d_C);cublasLtDestroy(ltHandle);
}

cuDNN的混合精度

关键点在descriptor，cudnnSetTensor4dDescriptor函数里


cudnnStatus_t cudnnSetTensor4dDescriptor(cudnnTensorDescriptor_t tensorDesc,cudnnTensorFormat_t format,   // NCHW 或 NHWCcudnnDataType_t dataType,     // float, half, double...int n, int c, int h, int w    // 维度
);

设定dataType枚举量，比如

CUDNN_DATA_FLOAT	float32	最常用数据类型，精度高
CUDNN_DATA_DOUBLE	float64	双精度，GPU 上使用少，性能低
CUDNN_DATA_HALF	float16	半精度，适合 Tensor Core 加速

mma.sync指令与WMMA API

Tensor Core 是 NVIDIA GPU 上的专用矩阵运算单元（从 Volta 架构开始有）。
- 它的核心指令就是 MMA (Matrix Multiply-Accumulate)，即做：
  D=A×B+C
  其中 A、B 是矩阵片段（tile），C 是累加结果，D 是输出。
Tensor Core 一次操作的 tile 是固定大小的（如 16x8x8、16x16x16 等），运算精度和输入/输出类型是可选的。

`mma.sync` PTX 指令格式

mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 d, a, b, c;

mma.sync：warp-level 同步的矩阵乘加运算（所有 32 线程参与一个 tile）。
aligned：说明内存对齐要求（通常是 128bit 对齐）。
m16n8k8：tile 形状，含义是：
- m = 16：结果矩阵有 16 行
- n = 8：结果矩阵有 8 列
- k = 8：乘法维度（A 的列数 / B 的行数）
row.col：矩阵存储布局
- row → A 按行主序
- col → B 按列主序
f16.f16.f16.f16：数据类型说明
- 第一个 f16：A 的元素类型
- 第二个 f16：B 的元素类型
- 第三个 f16：累加器 C 的类型
- 第四个 f16：输出 D 的类型
  常用组合：
- f16.f16.f16.f16 → 全半精度
- f16.f16.f16.f32 → FP16 输入，FP32 累加（最常用）
- s8.s8.s32.s32 → INT8 输入，INT32 累加

执行过程

以 mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f32 为例：

A：16×8，FP16
B：8×8，FP16
C/D：16×8，FP32
运算：

$D_{16*8}=A_{16*8}*B_{8*8}+C_{16*8}$

整个 tile 是由 一个 warp（32个线程）合作完成 的，每个线程负责 tile 中一小部分的元素。

WMMA API（CUDA C++ 封装）

需要包含头文件#include<mma.h>；链接cudart这个库

里面的所有类的作用域都在nvcuda下

wmma::fragment

  template<typename Use, // 用途（matrix_a, matrix_b, accumulator）
int m, int n, int k, // 矩阵乘法的 tile 大小 (M×K * K×N = M×N)
typename T, // 数据类型（half、float、tf32等）
typename Layout=void>// 矩阵存储布局（row_major 或 col_major，accumulator 不需要）
class fragment;

注意这里m,n,k不是随便啥组合都行的；头文件里有模板类的特化，没有特化的组合是用不了的。

wmma::load_matrix_sync

template <typename Use, int M, int N, int K, typename T, typename Layout>
__device__ void load_matrix_sync(fragment<Use, M, N, K, T, Layout> &a,const T *ptr,int ldm
);

参数

a：目标 fragment（传引用）。
ptr：指向 device 内存的矩阵起始地址。
ldm：leading dimension（行主时=列数，列主时=行数）。

功能

从 global memory 把矩阵数据加载进 fragment。
加载时会根据 row_major 或 col_major 的布局解释数据。

返回值

无返回值，结果写入 fragment。

wmma::fill_fragment

template <typename Use, int M, int N, int K, typename T>
__device__ void fill_fragment(fragment<Use, M, N, K, T> &a,const T &value
);

参数

a：目标 fragment（accumulator）。
value：要填充的常数。

功能

把 accumulator fragment 的所有元素初始化为 value。

返回值

无返回值，结果写入 fragment。

`wmma::mma_sync`

template <int M, int N, int K,typename T1, typename T2, typename T3,typename LayoutA, typename LayoutB>
__device__ void mma_sync(fragment<wmma::accumulator, M, N, K, T3> &d,const fragment<wmma::matrix_a, M, N, K, T1, LayoutA> &a,const fragment<wmma::matrix_b, M, N, K, T2, LayoutB> &b,const fragment<wmma::accumulator, M, N, K, T3> &c
);

参数

d：输出 accumulator fragment。
a：输入矩阵 A fragment（一般是 half）。
b：输入矩阵 B fragment（一般是 half）。
c：输入矩阵 C fragment（float，累加器）。

功能

执行矩阵乘加运算：
D=A×B+C
内部运算模式：FP16 × FP16 → FP32 累加。

返回值

无返回值，结果写入 d。

wmma::store_matrix_sync

template <int M, int N, int K, typename T, typename Layout>
__device__ void store_matrix_sync(T *ptr,const fragment<wmma::accumulator, M, N, K, T> &a,int ldm,Layout mem_layout
);

参数

ptr：目标 device 内存地址。
a：accumulator fragment（源数据）。
ldm：leading dimension。
mem_layout：wmma::mem_row_major 或 wmma::mem_col_major。

功能

把 fragment 中的结果写回全局内存。

返回值

无返回值，结果写入内存。

示例代码

#include <cublas_v2.h>
#include <cublasLt.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <iostream>
#include <vector>
#include <mma.h>
using namespace nvcuda;
__global__ void test_wmma_8x32x16(const half *A, const half *B, float *C) {// 声明 fragmentwmma::fragment<wmma::matrix_a, 8, 32, 16, half, wmma::row_major> a_frag;wmma::fragment<wmma::matrix_b, 8, 32, 16, half, wmma::col_major> b_frag;wmma::fragment<wmma::accumulator, 8, 32, 16, float> c_frag, d_frag;// 初始化wmma::load_matrix_sync(a_frag, A, 16); // lda = K = 16wmma::load_matrix_sync(b_frag, B, 16); // ldb = K = 16wmma::fill_fragment(c_frag, 0.0f);// 矩阵乘wmma::mma_sync(d_frag, a_frag, b_frag, c_frag);// 存回结果wmma::store_matrix_sync(C, d_frag, 32, wmma::mem_row_major); // ldc = N = 32
}int main() {const int M = 8, N = 32, K = 16;// host 数据std::vector<half> h_A(M*K);std::vector<half> h_B(K*N);std::vector<float> h_C(M*N);// 初始化 A = 1, B = 1for (int i = 0; i < M*K; i++) h_A[i] = __float2half(1.0f);for (int i = 0; i < K*N; i++) h_B[i] = __float2half(1.0f);// device 内存half *d_A, *d_B;float *d_C;cudaMalloc(&d_A, M*K*sizeof(half));cudaMalloc(&d_B, K*N*sizeof(half));cudaMalloc(&d_C, M*N*sizeof(float));cudaMemcpy(d_A, h_A.data(), M*K*sizeof(half), cudaMemcpyHostToDevice);cudaMemcpy(d_B, h_B.data(), K*N*sizeof(half), cudaMemcpyHostToDevice);// 启动 1 个 warp (32 threads)test_wmma_8x32x16<<<1, 32>>>(d_A, d_B, d_C);cudaMemcpy(h_C.data(), d_C, M*N*sizeof(float), cudaMemcpyDeviceToHost);// 打印结果std::cout << "C (8x32):\n";for (int i = 0; i < M; i++) {for (int j = 0; j < N; j++) {std::cout << h_C[i*N + j] << " ";}std::cout << "\n";}cudaFree(d_A);cudaFree(d_B);cudaFree(d_C);return 0;
}

为什么B是列主序？

对常用的 16×16×16 tile：