C++ Warp矩阵运算利用Tensor Core来加速D = A * B + C形式的矩阵问题，这些运算支持计算能力7.0或者更高设备的混合精度浮点数据。

1. 概述

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

（1）fragment

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

fragment：包含分布在Warp中所有线程上的矩阵片段。

Use：可用值：matrix_a、matrix_b、accumulator。matrix_a表示当前fragment用作第一个被乘数，A；matrix_b表示当前fragment用作第二个被乘数，B；accumulator表示当前fragment用作源或目标累加器（C或者D）。

m、n、k：描述参与乘法累加运算的Warp宽度矩阵块的形状，matrix_a块的尺寸为m*k；matrix_b的尺寸为k*n；accumulator块的尺寸为m*n。

数据T表示数据类型，对于A和B数据类型T可以是double、float、__half、__nv_bfloat16、char或者unsigned char；对于C和D数据类型T可以是double、float、int或者__half。

必须为matrix_a和matrix_b指定Layout参数，row_major 或 col_major 分别指示矩阵行或列中的元素在内存中是连续的。 accumulator 矩阵的 Layout 参数应保留默认值 void 。仅当如下所述加载或存储累加器时才指定行或列布局。

（2）load_matrix_sync

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

等到所有 Warp 通道都到达 load_matrix_sync 时，才从内存加载矩阵fragment a。

mptr 必须是指向内存中矩阵第一个元素的 mptr 位对齐指针 ldm 描述连续行（对于行主布局）或列（对于列主布局）之间元素的步幅，并且对于 __half 元素类型必须是 8 的倍数，对于 float 元素类型必须是 4 的倍数。（即，两种情况下均为 16 字节的倍数）。

如果片段是 accumulator ，则必须将 layout 参数指定为 mem_row_major 或 mem_col_major 。对于 matrix_a 和 matrix_b 片段，布局是从片段的 layout 参数推断出来的。

mptr、ldm 、 layout 和 a 的所有模板参数的值对于 Warp 中的所有线程必须相同。

Warp 中的所有线程都必须调用此函数，否则结果未定义。

（3）store_matrix_sync

void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);

等待所有 Warp 通道都到达 store_matrix_sync 后，将矩阵片段 a 存储到内存中。

mptr 必须是指向内存中矩阵首元素的 mptr 位对齐指针 ldm 描述连续行（对于行主布局）或列（对于列主布局）之间元素的步幅，对于 __half 元素类型，必须是 8 的倍数；对于 float 元素类型，必须是 4 的倍数。（即，两种情况下均为 16 字节的倍数）。

输出矩阵的布局必须指定为 mem_row_major 或 mem_col_major 。

所有 Warp 线程的 mptr 、 ldm 、 layout 和 a 的所有模板参数的值必须相同。

（4）fill_fragment

void fill_fragment(fragment<...> &a, const T& v);

用常数值 v 填充矩阵片段。由于矩阵元素到每个片段的映射未指定，因此该函数通常由 warp 中的所有线程使用公共的 v 值来调用。

（5）mma_sync

void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);

等待所有 Warp 通道到达 mma_sync 后，执行 Warp 同步矩阵乘法累加运算 D=A*B+C 。同时支持就地运算 C=A*B+C 。Warp 中所有线程的每个矩阵片段的 satf 值和模板参数值必须相同。此外，片段 A 、 B 、 C 和 D 的模板参数 m 、 n 和 k 必须匹配。Warp 中的所有线程都必须调用此函数，否则结果未定义。

如果 satf （饱和为有限值）模式为 true ，则以下附加数值属性适用于目标累加器：

如果元素结果为 +Infinity，则相应的累加器将包含 +MAX_NORM

如果元素结果为-Infinity，则相应的累加器将包含 -MAX_NORM

如果元素结果为 NaN，则相应的累加器将包含 +0

以下代码在单个 warp 中实现了 16x16x16 矩阵乘法。

#include <mma.h>
using namespace nvcuda;

__global__ void wmma_ker(half *a, half *b, float *c) {
   // Declare the fragments
   wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag;
   wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag;
   wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

   // Initialize the output to zero
   wmma::fill_fragment(c_frag, 0.0f);

   // Load the inputs
   wmma::load_matrix_sync(a_frag, a, 16);
   wmma::load_matrix_sync(b_frag, b, 16);

   // Perform the matrix multiplication
   wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

   // Store the output
   wmma::store_matrix_sync(c, c_frag, 16, wmma::mem_row_major);
}

Tensor Core 支持多种元素类型和矩阵大小。下表列出了支持的 matrix_a 、 matrix_b 和 accumulator 矩阵的各种组合：

Matrix A  矩阵 A	Matrix B  矩阵 B	Accumulator  累加器	Matrix Size (m-n-k)  矩阵大小（mnk）
__half  __一半	__half  __一半	float  漂浮	16x16x16
__half  __一半	__half  __一半	float  漂浮	32x8x16
__half  __一半	__half  __一半	float  漂浮	8x32x16
__half  __一半	__half  __一半	__half  __一半	16x16x16
__half  __一半	__half  __一半	__half  __一半	32x8x16
__half  __一半	__half  __一半	__half  __一半	8x32x16
unsigned char  无符号字符	unsigned char  无符号字符	int  整数	16x16x16
unsigned char  无符号字符	unsigned char  无符号字符	int  整数	32x8x16
unsigned char  无符号字符	unsigned char  无符号字符	int  整数	8x32x16
signed char  有符号字符	signed char  有符号字符	int  整数	16x16x16
signed char  有符号字符	signed char  有符号字符	int  整数	32x8x16
signed char  有符号字符	signed char  有符号字符	int  整数	8x32x16

Matrix A  矩阵 A	Matrix B  矩阵 B	Accumulator  累加器	Matrix Size (m-n-k)  矩阵大小（mnk）
__nv_bfloat16	__nv_bfloat16	float  漂浮	16x16x16
__nv_bfloat16	__nv_bfloat16	float  漂浮	32x8x16
__nv_bfloat16	__nv_bfloat16	float  漂浮	8x32x16
precision::tf32  精度::tf32	precision::tf32  精度::tf32	float  漂浮	16x16x8

Matrix A  矩阵 A	Matrix B  矩阵 B	Accumulator  累加器	Matrix Size (m-n-k)  矩阵大小（mnk）
double  双倍的	double  双倍的	double  双倍的	8x8x4

 

Matrix A  矩阵 A	Matrix B  矩阵 B	Accumulator  累加器	Matrix Size (m-n-k)  矩阵大小（mnk）
precision::u4  精度::u4	precision::u4  精度::u4	int  整数	8x8x32
precision::s4  精度::s4	precision::s4  精度::s4	int  整数	8x8x32
precision::b1  精度::b1	precision::b1  精度::b1	int  整数	8x8x128