CUDA 9中张量核（Tensor Cores）编程

CUDA 9中张量核（Tensor Cores）编程

Programming Tensor Cores in CUDA 9

一．概述

新的Volta GPU架构的一个重要特点是它的Tensor核，使Tesla V100加速器的峰值吞吐量是上一代Tesla P100的32位浮点吞吐量的12倍。Tensor内核使人工智能程序员能够使用混合精度来获得更高的吞吐量，而不牺牲精度。            

Tensor核心已经在许多深度学习框架（包括Tensorflow、PyTorch、MXNet和Caffe2）中支持深度学习训练，无论是在主版本中还是通过pull请求。有关在使用这些框架时启用Tensor核心的更多信息，请参阅《混合精度训练指南》。对于深度学习推理，最近的TensorRT 3版本也支持Tensor核心。             

本文将展示如何使用CUDA库在自己的应用程序中使用张量核，以及如何在CUDA C++设备代码中直接编程。

 二．什么是张量核（Tensor Cores）？             

特斯拉V100的张量核心是可编程的矩阵乘法和累加单元，可以提供多达125 Tensor tflop的训练和推理应用。特斯拉V100 GPU包含640个Tensor Cores：8/SM。Tensor内核及其相关的数据路径是定制的，以显著提高浮点计算吞吐量，只需适当的区域和功耗。时钟选通广泛应用于最大限度地节省功耗。             

每个张量核提供一个4x4x4矩阵处理数组，它执行操作D=a*B+C，其中a、B、C和D是4×4矩阵，如图1所示。矩阵乘法输入A和B是FP16矩阵，而累积矩阵C和D可以是FP16或FP32矩阵。

 

  每个张量核执行64个浮点FMA混合精度操作每个时钟（FP16输入乘法与全精度积和FP32累加，如图2所示）和8张量核在一个SM执行总共1024个浮点操作每个时钟。与使用标准FP32操作的Pascal GP100相比，每SM深度学习应用程序的吞吐量显著增加了8倍，因此Volta V100 GPU的吞吐量与Pascal P100 GPU相比总共增加了12倍。张量核对FP16输入数据进行运算，FP32累加。如图2所示，对于4x4x4矩阵乘法，FP16乘法产生的全精度结果是在FP32运算中与给定点积中的其他乘积累积的结果。

  三.  CUDA库中的张量核              

使用Tensor核的两个CUDA库是cuBLAS和cuDNN。cuBLAS使用张量核加速GEMM计算（GEMM是矩阵-矩阵乘法的BLAS术语）；cuDNN使用张量核加速卷积和递归神经网络（RNNs）。             

许多计算应用程序使用GEMM：信号处理、流体动力学等等。随着这些应用程序的数据大小呈指数级增长，这些应用程序需要在处理速度上进行匹配。图3中的混合精度GEMM性能图显示，张量核显然满足了这一需求。             

提高卷积速度的需求同样巨大；例如，深神经网络（DNNs）使用了许多层卷积。人工智能研究人员每年都在设计越来越深的神经网络；最深的神经网络中的卷积层现在有几十个。训练DNNs需要卷积层在正向和反向传播期间重复运行。

图4中的卷积性能图显示，张量核满足卷积性能的需要。

两个性能图表都显示，特斯拉V100的张量核心提供了数倍于上一代特斯拉P100的性能。性能改进这一巨大的变化如何在计算领域工作：使交互性成为可能，启用“假设”方案研究，或减少服务器场使用。如果在应用程序中使用GEMM或卷积，请使用下面的简单步骤来提高工作效率。

四．如何在立方体中使用张量核             

通过对现有cuBLAS代码进行一些更改，可以利用张量核。这些变化是在使用cuBLAS API时的小变化。             

下面的示例代码应用一些简单的规则来指示cuBLAS应该使用张量核；这些规则在代码后面显式枚举。             

示例代码             

下面的代码与以前的架构中用于调用cuBLAS中GEMM的通用代码基本相同。
// First, create a cuBLAS handle:

cublasStatus_t cublasStat = cublasCreate(&handle);

// Set the math mode to allow cuBLAS to use Tensor Cores:

cublasStat = cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH);

// Allocate and initialize your matrices (only the A matrix is shown):

size_t matrixSizeA = (size_t)rowsA * colsA;

T_ELEM_IN **devPtrA = 0;

cudaMalloc((void**)&devPtrA[0], matrixSizeA * sizeof(devPtrA[0][0]));

T_ELEM_IN A  = (T_ELEM_IN *)malloc(matrixSizeA * sizeof(A[0]));

memset( A, 0xFF, matrixSizeA* sizeof(A[0]));

status1 = cublasSetMatrix(rowsA, colsA, sizeof(A[0]), A, rowsA, devPtrA[i], rowsA);

// ... allocate and initialize B and C matrices (not shown) ...

// Invoke the GEMM, ensuring k, lda, ldb, and ldc are all multiples of 8,

// and m is a multiple of 4:

cublasStat = cublasGemmEx(handle, transa, transb, m, n, k, alpha,

                          A, CUDA_R_16F, lda,

                          B, CUDA_R_16F, ldb,

                          beta, C, CUDA_R_16F, ldc, CUDA_R_32F, algo);

五．一些简单的规则             

cuBLAS用户将注意到其现有cuBLAS GEMM代码的一些变化：              

例程必须是GEMM；目前，只有GEMM支持Tensor核心执行。             

数学模式必须设置为CUBLAS_TENSOR_OP_math。浮点数学不具有关联性，因此张量核心数学例程的结果与类似的非张量核心数学例程的结果不完全等价。cuBLAS要求用户“选择”使用张量核。             

k、lda、ldb和ldc都必须是8的倍数；m必须是4的倍数。张量核心数学例程以八个值的步骤遍历输入数据，因此矩阵的维数必须是八的倍数。             

矩阵的输入和输出数据类型必须是半精度或单精度。（上面只显示了CUDA_R_16F，但也支持CUDA_R_32F。）不满足上述规则的GEMM将返回到非张量核心实现。              GEMM性能             

如前所述，Tensor内核提供的GEMM性能是以前硬件的几倍。图3显示了GP100（Pascal）和GV100（Volta）硬件的比较性能。

 

  六．如何在cuDNN中使用张量核             

在cuDNN中使用Tensor核也很容易，而且只涉及对现有代码的微小更改。             

示例代码             

在cuDNN中使用张量核的示例代码可以在conv中找到_示例.cpp在cuDNN samples目录中；复制了下面的一些摘录。（cuDNN samples目录与文档打包在一起。）

// Create a cuDNN handle:

checkCudnnErr(cudnnCreate(&handle_));

// Create your tensor descriptors:

checkCudnnErr( cudnnCreateTensorDescriptor( &cudnnIdesc ));

checkCudnnErr( cudnnCreateFilterDescriptor( &cudnnFdesc ));

checkCudnnErr( cudnnCreateTensorDescriptor( &cudnnOdesc ));

checkCudnnErr( cudnnCreateConvolutionDescriptor( &cudnnConvDesc ));

// Set tensor dimensions as multiples of eight (only the input tensor is shown here):

int dimA[] = {1, 8, 32, 32};

int strideA[] = {8192, 1024, 32, 1};

checkCudnnErr( cudnnSetTensorNdDescriptor(cudnnIdesc, getDataType(),

                                          convDim+2, dimA, strideA) );

// Allocate and initialize tensors (again, only the input tensor is shown):

checkCudaErr( cudaMalloc((void**)&(devPtrI), (insize) * sizeof(devPtrI[0]) ));

hostI = (T_ELEM*)calloc (insize, sizeof(hostI[0]) );

initImage(hostI, insize);

checkCudaErr( cudaMemcpy(devPtrI, hostI, sizeof(hostI[0]) * insize, cudaMemcpyHostToDevice));

// Set the compute data type (below as CUDNN_DATA_FLOAT):

checkCudnnErr( cudnnSetConvolutionNdDescriptor(cudnnConvDesc,

                                               convDim,

                                               padA,

                                               convstrideA,

                                               dilationA,

                                               CUDNN_CONVOLUTION,

                                               CUDNN_DATA_FLOAT) );

 

// Set the math type to allow cuDNN to use Tensor Cores:

checkCudnnErr( cudnnSetConvolutionMathType(cudnnConvDesc, CUDNN_TENSOR_OP_MATH) );

// Choose a supported algorithm:

cudnnConvolutionFwdAlgo_t algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;

// Allocate your workspace:

checkCudnnErr( cudnnGetConvolutionForwardWorkspaceSize(handle_, cudnnIdesc,

                                                       cudnnFdesc, cudnnConvDesc,

                                                       cudnnOdesc, algo, &workSpaceSize) );

if (workSpaceSize > 0) {

   cudaMalloc(&workSpace, workSpaceSize);

}

// Invoke the convolution:

checkCudnnErr( cudnnConvolutionForward(handle_, (void*)(&alpha), cudnnIdesc, devPtrI,

                                       cudnnFdesc, devPtrF, cudnnConvDesc, algo,

                                       workSpace, workSpaceSize, (void*)(&beta),

                                       cudnnOdesc, devPtrO) );

七．一些简单的规则             

注意一些与常用cuDNN用法不同的变化：             

卷积算法必须是ALGO 1（前向的隐式预处理）。在将来的cuDNN版本中，除ALGO 1之外的其他卷积算法可能使用张量核。             

数学类型必须设置为CUDNN_TENSOR_OP_math。与cuBLAS一样，张量核数学例程的结果与类似的非张量核数学例程的结果并不完全等价，因此cuDNN要求用户“选择”使用张量核。              输入和输出通道尺寸必须是8的倍数。同样，在cuBLAS中，张量核心数学例程以8个值的步长遍历输入数据，因此输入数据的维数必须是8的倍数。             

卷积的输入、滤波和输出数据类型必须为半精度。             

不满足上述规则的卷积将返回到非张量核心实现。             

上面的示例代码显示了NCHW数据格式，请参见conv_示例.cppNHWC支持的样本。              卷积性能             

如前所述，张量核的卷积性能是以前硬件的几倍。图4显示了GP100（Pascal）和GV100（Volta）硬件的比较性能。

 

 八．在CUDA 9.0中对张量核的编程访问             

通过CUDA9.0访问内核中的Tensor核是一个预览功能。这意味着本节中描述的数据结构、api和代码在将来的CUDA版本中可能会发生更改。             

虽然CuBLAS和CUDNN覆盖了张量核的许多潜在用途，但是也可以直接在CUDA C++中编程它们。张量核在CUDA 9.0中通过nvcuda：：wmma命名空间中的一组函数和类型公开。它们允许将值加载或初始化为张量核所需的特殊格式，执行矩阵乘法累加（MMA）步骤，并将值存储回内存。在程序执行期间，多个张量核被一个完全扭曲同时使用。这允许warp在非常高的吞吐量下执行16x16x16mma（图5）。

 让看一个简单的例子，它展示了如何使用WMMA（Warp Matrix Multiply Accumulate）API执行矩阵乘法。请注意，这个示例并不是为高性能而调整的，它主要用作API的演示。为了获得更好的性能，可以应用于此代码的优化示例，请查看CUDA工具包中的cudatensorcoregem示例。为了获得最高的生产性能，应使用立方块，如上所述。 标题和命名空间             

WMMA API包含在mma.h头文件中。完整的名称空间是nvcuda：：wmma：：*，但是在整个代码中保持wmma显式很有用，因此将只使用nvcuda名称空间。

#include <mma.h>

using namespace nvcuda;

九．声明和初始化             

完整的GEMM规范允许算法处理a或b的转置，并允许数据跨距大于矩阵中的跨距。为了简单起见，假设a和b都没有被转置，并且内存和矩阵的前导维数是相同的。              将采用的策略是让一个warp负责输出矩阵的一个16×16部分。通过使用二维网格和线程块，可以有效地将曲面平铺到二维输出矩阵上。

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

在执行MMA操作之前，操作数矩阵必须在GPU的寄存器中表示。由于MMA是一个全曲速操作，这些寄存器分布在曲速的线程中，每个线程持有整个矩阵的一个片段。各个矩阵参数与其片段之间的映射是不透明的，因此程序不应对此进行假设。              在CUDA中，片段是一种模板类型，模板参数描述片段保存的矩阵（a、B或累加器）、整个WMMA操作的形状、数据类型，以及对于a和B矩阵，数据是主要行还是主要列。最后一个参数可用于执行A或B矩阵的换位。这个例子没有换位，所以两个矩阵都是列主矩阵，这是GEMM的标准。

// Declare the fragments

    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> a_frag;

    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::col_major> b_frag;

    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_frag;

    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;

wmma::fill_fragment(acc_frag, 0.0f);

初始化步骤的最后一部分是用零填充累加器片段。

内循环             

用于GEMM的策略是计算每个曲面的输出矩阵的一个平铺。为此，需要循环遍历矩阵的行和列。这是沿着两个矩阵的K维，并生成一个MxN输出平铺。load matrix函数从内存（在本例中是全局内存，尽管它可以是任何内存空间）获取数据并将其放入片段中。加载的第三个参数是矩阵内存中的“前导维度”；加载的16×16平铺在内存中是不连续的，因此函数需要知道连续列（或行，如果这些列是行的主要片段）之间的跨距。             

MMA调用累积到位，因此第一个和最后一个参数都是先前初始化为零的累加器片段。

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

结束             

acc_frag现在保存基于A和B的乘法的该曲面输出平铺的结果。完整的GEMM规范允许缩放该结果，并将其累积在适当的矩阵上。实现这种缩放的一种方法是对片段执行按元素的操作。虽然没有定义从矩阵坐标到线程的映射，但是元素操作不需要知道这个映射，所以仍然可以使用片段执行。因此，对片段执行缩放操作或将一个片段的内容添加到另一个片段是合法的，只要这两个片段具有相同的模板参数。如果片段具有不同的模板参数，则结果未定义。利用这个特性，我们在C语言中加载现有的数据，并以正确的比例，用它累积到目前为止的计算结果。

// 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];

        }

最后，将数据存储到内存中。目标指针可以是GPU可见的任何内存空间，并且必须指定内存中的前导维度。还有一个选项可以指定输出是写入row还是column major。

        // Store the output

        wmma::store_matrix_sync(c + cRow + cCol * ldc, c_frag, ldc, wmma::mem_col_major);

    }

}