Tensor Cores 使用介紹
共 12652字，需瀏覽 26分鐘
·
2024-04-20 12:32
概要介紹
TensorCore 是從Nvidia Volta 架構(gòu)GPU開始支持的重要特性，使CUDA開發(fā)者能夠使用混合精度來(lái)獲得更高的吞吐量，而不犧牲精度。TensorCore已經(jīng)在許多深度學(xué)習(xí)框架（包括Tensorflow、PyTorch、MXNet和Caffe2）中支持深度學(xué)習(xí)訓(xùn)練。本文將展示如何使用CUDA庫(kù)在自己的應(yīng)用程序中使用張量核，以及如何在CUDA C++設(shè)備代碼中直接編程。
TensorCore
TensorCore是可編程的矩陣乘法和累加單元，可以提供多達(dá)125 Tensor tflop的訓(xùn)練和推理應(yīng)用。TensorCore及其相關(guān)的數(shù)據(jù)路徑是定制的，以顯著提高浮點(diǎn)計(jì)算吞吐量。每個(gè)TensorCore提供一個(gè)4x4x4矩陣處理數(shù)組，它執(zhí)行操作D=A*B+C，其中A、B、C和D是4×4矩陣，如下圖所示。矩陣乘法輸入A和B是FP16矩陣，而累積矩陣C和D可以是FP16或FP32矩陣。
每個(gè)TensorCore每個(gè)時(shí)鐘周期可以執(zhí)行64個(gè)浮點(diǎn)FMA混合精度操作，而在一個(gè)SM中有8個(gè)TensorCore，所以一個(gè)SM中每個(gè)時(shí)鐘可以執(zhí)行1024(8x64x2)個(gè)浮點(diǎn)操作。TensorCore對(duì)FP16輸入數(shù)據(jù)進(jìn)行運(yùn)算，使用FP32累加。如圖下圖所示，對(duì)于4x4x4矩陣乘法，F(xiàn)P16乘法的結(jié)果是一個(gè)完整精度的值，該值在進(jìn)行4x4x4矩陣乘法的點(diǎn)積運(yùn)算中與其他乘積一起累積在FP32操作中。
對(duì)一般用戶來(lái)說(shuō)，可以通過(guò)使用cuBLAS和cuDNN這兩個(gè)CUDA庫(kù)來(lái)間接使用Tensor Cores。cuBLAS利用Tensor Cores加速GEMM計(jì)算（GEMM是BLAS中矩陣乘法的術(shù)語(yǔ)）；cuDNN則利用Tensor Cores加速卷積和循環(huán)神經(jīng)網(wǎng)絡(luò)（RNNs）的計(jì)算。
cuBLAS中使用TensorCore
可以通過(guò)對(duì)現(xiàn)有的cuBLAS代碼進(jìn)行一些更改來(lái)充分利用Tensor Cores。這些更改是對(duì)cuBLAS API的使用進(jìn)行的小修改。以下示例代碼應(yīng)用了一些簡(jiǎn)單的規(guī)則，以指示cuBLAS應(yīng)該使用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);cuBLAS用戶將注意到與現(xiàn)有的cuBLAS GEMM代碼相比有一些變化：
例程必須是一個(gè)GEMM；目前只有GEMM支持Tensor Core執(zhí)行。
數(shù)學(xué)模式必須設(shè)置為
CUBLAS_TENSOR_OP_MATH。k、lda、ldb和ldc必須是8的倍數(shù)；m必須是4的倍數(shù)。Tensor Core數(shù)學(xué)例程以8個(gè)值為一步跨越輸入數(shù)據(jù)，因此矩陣的維度必須是8的倍數(shù)。矩陣的輸入和輸出數(shù)據(jù)類型必須是半精度或單精度。
不滿足上述規(guī)則的GEMM將回退到非Tensor Core實(shí)現(xiàn)。
cuDNN中使用TensorCore
在cuDNN中使用Tensor Cores也很簡(jiǎn)單，而且同樣只需要對(duì)現(xiàn)有代碼進(jìn)行輕微修改。
// 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（前向傳播時(shí)為IMPLICIT_PRECOMP_GEMM）。除了ALGO_1之外的其他卷積算法可能會(huì)在未來(lái)的cuDNN版本中使用Tensor Cores。數(shù)學(xué)類型必須設(shè)置為
CUDNN_TENSOR_OP_MATH,與cuBLAS類似.輸入和輸出通道的維度都必須是8的倍數(shù)。與cuBLAS類似，Tensor Core數(shù)學(xué)例程以8個(gè)值為一步跨越輸入數(shù)據(jù)，因此輸入數(shù)據(jù)的維度必須是8的倍數(shù)。
卷積的輸入、濾波器和輸出數(shù)據(jù)類型必須是半精度。
不滿足上述規(guī)則的卷積將回退到非Tensor Core實(shí)現(xiàn)。
CUDA C++中使用TensorCore
雖然cuBLAS和cuDNN涵蓋了許多Tensor Cores的潛在用途，但用戶還可以直接在CUDA C++中編程。Tensor Cores通過(guò)nvcuda::wmma命名空間中的一組函數(shù)和類型在CUDA 9.0中公開。這些函數(shù)和類型允許您將值加載或初始化到張量核心所需的特殊格式中，執(zhí)行矩陣乘累加（MMA）步驟，并將值存回內(nèi)存。在程序執(zhí)行期間，一個(gè)完整的warp可以同時(shí)使用多個(gè)Tensor Cores，這使得warp能夠以非常高的吞吐量執(zhí)行16x16x16的MMA。核心的API如下所示，詳細(xì)介紹見文檔。
template<typename Use, int m, int n, int k, typename T, typename Layout=void> class fragment;// 等待直到所有warp lanes都到達(dá)load_matrix_sync，然后從內(nèi)存中加載矩陣片段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都到達(dá)store_matrix_sync，然后將矩陣片段a存儲(chǔ)到內(nèi)存中。void store_matrix_sync(T* mptr, const fragment<...> &a, unsigned ldm, layout_t layout);// 使用常量值v填充一個(gè)矩陣片段。void fill_fragment(fragment<...> &a, const T& v);// 等待直到所有warp lanes都到達(dá)mma_sync，然后執(zhí)行warp同步的矩陣乘累加操作D = A * B + C。void mma_sync(fragment<...> &d, const fragment<...> &a, const fragment<...> &b, const fragment<...> &c, bool satf=false);下面來(lái)看個(gè)實(shí)際的簡(jiǎn)單例子。
頭文件引用
#include <mma.h>using namespace nvcuda;聲明和定義
完整的GEMM規(guī)范允許算法在a或b的轉(zhuǎn)置上工作，并且數(shù)據(jù)步幅可以大于矩陣中的步幅。為簡(jiǎn)單起見，我們假設(shè)a和b都沒(méi)有被轉(zhuǎn)置，并且內(nèi)存和矩陣的主導(dǎo)維度相同。我們采用的策略是讓一個(gè)warp負(fù)責(zé)輸出矩陣的一個(gè)16×16的部分。通過(guò)使用二維網(wǎng)格和線程塊，我們可以有效地在二維輸出矩陣上劃分warp。
// 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);在執(zhí)行MMA操作之前，操作數(shù)矩陣必須表示在GPU的寄存器中。由于MMA是一個(gè)warp范圍的操作，這些寄存器分布在warp的各個(gè)線程之間，每個(gè)線程持有整個(gè)矩陣的一個(gè)fragment。在CUDA中，fragment是一個(gè)模板類型，具有描述片段持有的矩陣、整個(gè)WMMA操作的形狀、數(shù)據(jù)類型以及A和B矩陣中數(shù)據(jù)是按行還是按列主序的模板參數(shù)。最后一個(gè)參數(shù)可以用于對(duì)A或B矩陣進(jìn)行轉(zhuǎn)置。這個(gè)示例中沒(méi)有進(jìn)行轉(zhuǎn)置，所以兩個(gè)矩陣都是按列主序的，這是GEMM的標(biāo)準(zhǔn)方式。
// 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);內(nèi)部循環(huán)
我們用于GEMM的策略是每個(gè)warp計(jì)算輸出矩陣的一個(gè)tile。為此，我們需要在A矩陣的行和B矩陣的列上進(jìn)行循環(huán)。這沿著這兩個(gè)矩陣的K維度進(jìn)行，并生成一個(gè)MxN的輸出tile。load矩陣函數(shù)從內(nèi)存中獲取數(shù)據(jù)（在這個(gè)示例中是全局內(nèi)存，盡管它可以是任何內(nèi)存空間），并將其放入一個(gè)fragment中。load的第三個(gè)參數(shù)是矩陣在內(nèi)存中的“主導(dǎo)維度”；我們加載的16×16 tile在內(nèi)存中是不連續(xù)的，因此函數(shù)需要知道連續(xù)列（或行，如果這些是按行主序的片段）之間的跨度。MMA調(diào)用在原地累積，因此第一個(gè)和最后一個(gè)參數(shù)都是我們之前初始化為零的累加器fragment。
// 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);
}
}結(jié)果寫回
現(xiàn)在，acc_frag根據(jù)A和B的乘積持有了這個(gè)warp輸出tile的結(jié)果。完整的GEMM規(guī)范允許對(duì)這個(gè)結(jié)果進(jìn)行縮放，并且可以就地累積到一個(gè)矩陣上。進(jìn)行這種縮放的一種方法是對(duì)fragment執(zhí)行逐元素操作。盡管從矩陣坐標(biāo)到線程的映射未定義，但逐元素操作不需要知道這種映射，因此仍然可以使用片段執(zhí)行這些操作。因此，對(duì)fragment執(zhí)行縮放操作或?qū)⒁粋€(gè)fragment的內(nèi)容添加到另一個(gè)fragment中是合法的，只要這兩個(gè)fragment具有相同的模板參數(shù)。利用這個(gè)特性，我們加載了C中的現(xiàn)有數(shù)據(jù)，并在正確的縮放下將計(jì)算結(jié)果迄今為止與之累積。
最后，我們將數(shù)據(jù)存儲(chǔ)到內(nèi)存中。目標(biāo)指針再次可以是對(duì)GPU可見的任何內(nèi)存空間，并且必須指定內(nèi)存中的主導(dǎo)維度。還有一個(gè)選項(xiàng)可以指定輸出是按行還是按列主序?qū)懭搿?/p>
// 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);
}}
總結(jié)
本文介紹了TensorCore和其api wmma api的使用，和常規(guī)CUDA C開發(fā)不太一樣的地方，就是它是warp-level的，這里需要切換下思路，其實(shí)可以把它想象成synchronized函數(shù)，warp中所有線程都在等待wmma中的流程執(zhí)行完畢，可以從所有api都有sync結(jié)尾來(lái)引導(dǎo)出。wmma api運(yùn)用了線程中的寄存器GPR和TensorCore硬件來(lái)加速矩陣運(yùn)算，絕大多數(shù)場(chǎng)景下通過(guò)cuBLAS或cuDNN使用即可。
參考
Programming Tensor Cores in CUDA 9 | NVIDIA Technical Blog
https://docs.nvidia.com/cuda/cu
- The End -
長(zhǎng)按二維碼關(guān)注我們
本公眾號(hào)專注：
1. 技術(shù)分享；
2. 學(xué)術(shù)交流；
3. 資料共享。
歡迎關(guān)注我們，一起成長(zhǎng)！