Tensor Cores 使用介绍

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

• 例程必须是一个GEMM；目前只有GEMM支持Tensor Core执行。

• 数学模式必须设置为CUBLAS_TENSOR_OP_MATH。

• k、lda、ldb和ldc必须是8的倍数；m必须是4的倍数。Tensor Core数学例程以8个值为一步跨越输入数据，因此矩阵的维度必须是8的倍数。

• 矩阵的输入和输出数据类型必须是半精度或单精度。

• 不满足上述规则的GEMM将回退到非Tensor Core实现。

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

• 卷积算法必须是ALGO_1（前向传播时为IMPLICIT_PRECOMP_GEMM）。除了ALGO_1之外的其他卷积算法可能会在未来的cuDNN版本中使用Tensor Cores。

• 数学类型必须设置为CUDNN_TENSOR_OP_MATH,与cuBLAS类似.

• 输入和输出通道的维度都必须是8的倍数。与cuBLAS类似，Tensor Core数学例程以8个值为一步跨越输入数据，因此输入数据的维度必须是8的倍数。

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

• 不满足上述规则的卷积将回退到非Tensor Core实现。

template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment;// 等待直到所有warp lanes都到达load_matrix_sync，然后从内存中加载矩阵片段a。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 lanes都到达store_matrix_sync，然后将矩阵片段a存储到内存中。void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);// 使用常量值v填充一个矩阵片段。void fill_fragment(fragment<...> &a, const T& v);// 等待直到所有warp lanes都到达mma_sync，然后执行warp同步的矩阵乘累加操作D = A * B + C。void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);

#include <mma.h>using namespace nvcuda;

// The only dimensions currently supported by WMMAconst 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);

// 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;// set o in accumulator fragment   wmma::fill_fragment(acc_frag, 0.0f);

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

    // 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];
        }
        // Store the output        wmma::store_matrix_sync(c + cRow + cCol * ldc, c_frag, ldc, wmma::mem_col_major);
    }}