从CUDA9.0开始就已经支持代码中调用tensor core进行计算,tensor core是NVIDIA的volta架构中新处理单元,分布于各个流处理器(SM)中,其在物理层支持如下形式的运算:
其中矩阵乘法中的A,B数据类型必须为FP16,而累加矩阵C和最终输出数据类型可以为FP16也可以为FP32。
要在自己的kernel中使用tensor core必须包含头文件mma.h,相关变量和结构位于叫做wmma(Warp Matrix Multiply Accumulate) 的namespace中。需要注意的是虽然tensor core物理基本结构为4*4*4的单元,但一般用的时候直接进行16*16大小的矩阵乘法(事实上官方的例子中也提到了WMMA的操作维度目前只支持16*16)。
下面以官方的例子进行说明:
- 定义和初始化
单次16*16矩阵的乘法和累加称为wmma操作,首先定义wmma在各个方向上的维度(实际上就是单次计算的小矩阵(tile)的size),这里输入的M,N,K,分别为矩阵A的行数,矩阵B的列数,以及A的列数也就是B的行数,注意这里A,B的size不局限于16*16,因为16*16只是单次wmma可运算的矩阵大小,而通过多个warp(SM中的基本执行代码单元,包含32个thread)和循环实现更大size的矩阵运算。