在cuda中使用tensor core计算GEMM(上)

 

从CUDA9.0开始就已经支持代码中调用tensor core进行计算,tensor core是NVIDIA的volta架构中新处理单元,分布于各个流处理器(SM)中,其在物理层支持如下形式的运算:

在cuda中使用tensor core计算GEMM(上)_第1张图片

其中矩阵乘法中的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的矩阵运算。

// The only dimensions currently supported by WMMA
const int WMMA_M = 16;
const int WMMA_N = 16;
const int WMMA_K = 16;

__global__ void wmma_example(half *a, half *b, float *c, 
                             int M, int N, int K, 
                             float alpha, float beta) 
{

    // Leading dimensions. Packed with no transpositions.
    int lda = M;
    int ldb = K;
    int ldc = M;
    
    // Tile using a 2D grid
    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
    int warpN = (blockIdx.y * blockDim.y + threadIdx.y);

 

  • 声明fragment变量
    // Declare the fragments
    wmma::fragment a_frag;
    wmma::fragment b_frag;
    wmma::fragment acc_frag;
    wmma::fragment c_frag;

 fragement变量用于对应一次wmma操作中的小矩阵块,声明时所填的参数表明了该fragment在计算中的用途和角色(用于乘法的matrix_a和matrix_b,附加项和输出矩阵则定义为accumulator),以及三个参数用于表示数据的维度形状,如果是matrix_a或matrix_b还必须说明数据在内存中的存储形式是按行还是按列以方便GPU在读取数据时。

  • 循环计算一系列tile的乘积

计算之前需要加载数据,用到的也是wmma命名空间下得api,前两个参数为fragment变量和其实地址,最后一个参数为leading dimension,表示数据在按行存储时矩阵的总行数或是按列存储时的总列数,这样知道了输入矩阵的存储方式和leading dimension在加载数据时GPU才知道如何在不连续的数据中取出我们需要的tile。

    // Loop over the K-dimension
    for (int i = 0; i < K; i += WMMA_K) {
        int aRow = warpM * WMMA_M;
        int aCol = i;
        int bRow = i;
        int bCol = warpN * WMMA_N;
        
        // Bounds checking
        if (aRow < M && aCol < K && bRow < K && bCol < N) {
            // Load the inputs
            wmma::load_matrix_sync(a_frag, a + aRow + aCol * lda, lda);
            wmma::load_matrix_sync(b_frag, b + bRow + bCol * ldb, ldb);

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

首先利用warp的ID来决定这个warp所计算的tile,然后将加载计算所需的数据,之前定义了leading dimensions就是为了这里让GPU知道该如何从不连续的数据中得到计算所需要的。为了方便理解单次循环中的计算,简单画了个示意图:

在cuda中使用tensor core计算GEMM(上)_第2张图片

可以看出一次循环中一个warp完成了a和b中一个tile的乘法计算,而每次循环的结果都累加在acc_fag上,最后acc_frag中的结果也就是结果矩阵c中该warp对应的tile。至于为什么大的矩阵运算可以划成tile并将tile看成元素一样运算是线性代数的基本定理,这里不再赘述。

  • scale和附加矩阵

之前提到了tensor core不仅可以一次进行一个tile 的乘法,还可以在乘法结果上加上一个附加项,官方blog中同时也提到,在进行加法计算之前,乘法输出和附加矩阵都可以先进行scale也就是乘上一个系数。

    // Load in current value of c, scale by beta, and add to result scaled by alpha
    int cRow = warpM * WMMA_M;
    int cCol = warpN * WMMA_N;
    
    if (cRow < M && cCol < N) {
        wmma::load_matrix_sync(c_frag, c + cRow + cCol * ldc, ldc, wmma::mem_col_major);
        
        for(int i=0; i < c_frag.num_elements; i++) {
            c_frag.x[i] = alpha * acc_frag.x[i] + beta * c_frag.x[i];
        }

这里num_elements和x[]都是fragment类的成员变量,方便用户对fragment的数据进行逐一单独访问。

//num_elements和x[]在fragment类中的定义
enum fragment::num_elements;
T fragment::x[num_elements];
  • 最后保存计算结果
        // Store the output
        wmma::store_matrix_sync(c + cRow + cCol * ldc, c_frag, ldc, wmma::mem_col_major);
    }
}

官方blog指出保存的地址可以为GPU上任意可见地址(及shared memory和global memory)。

你可能感兴趣的:(CUDA,tensor,core)