如何使用張量核在CUDA C++設(shè)備代碼中編程

新 Volta GPU 架構(gòu)的一個(gè)定義性特征是它的 張量核 ，它使 Tesla V100 加速器的峰值吞吐量是上一代 Tesla P100 的 32 位浮點(diǎn)吞吐量的 12 倍。張量核心使人工智能程序員能夠使用 混合精度 來(lái)實(shí)現(xiàn)更高的吞吐量而不犧牲精度。

張量核心已經(jīng)在主版本或許多深度學(xué)習(xí)框架（包括 PyTorch 、 TensorFlow 、 MXNet 和 Caffe2 ）中通過(guò) pull 請(qǐng)求支持 深度學(xué)習(xí) 培訓(xùn)。有關(guān)在使用這些框架時(shí)啟用張量核心的更多信息，請(qǐng)查看 混合精度訓(xùn)練指南 。

在這篇博客文章中，我們展示了如何使用 CUDA 庫(kù)在自己的應(yīng)用程序中使用張量核，以及如何直接在 CUDA C ++設(shè)備代碼中編程。

什么是張量核？

Tesla V100 的張量核心是可編程的矩陣乘法和累加單元，可為訓(xùn)練和推理應(yīng)用提供多達(dá) 125 個(gè)張量 TFLOP 。 Tesla V100GPU 包含 640 個(gè)張量核心：每平方米 8 個(gè)。張量核心及其相關(guān)數(shù)據(jù)路徑都是定制的，可以顯著提高浮點(diǎn)計(jì)算吞吐量，只需適度的面積和功耗成本。時(shí)鐘門(mén)控廣泛用于最大限度地節(jié)省電力。

每個(gè)張量核提供一個(gè) 4x4x4 矩陣處理數(shù)組，該數(shù)組執(zhí)行運(yùn)算 D = A * B + C ，其中 答：， B 、 C 和 D 是 4 × 4 矩陣，如圖 1 所示。矩陣乘法輸入 A 和 B 是 FP16 矩陣，而累加矩陣 C 和 D 可以是 FP16 或 FP32 矩陣。

圖 1 ：張量核 4x4x4 矩陣乘法和累加。

每個(gè)張量核心對(duì)每個(gè)時(shí)鐘執(zhí)行 64 個(gè)浮點(diǎn) FMA 混合精度運(yùn)算（ FP16 輸入乘法全精度乘積， FP32 累加，如圖 2 所示），一個(gè) SM 中的 8 個(gè)張量核心每個(gè)時(shí)鐘執(zhí)行 1024 個(gè)浮點(diǎn)運(yùn)算。與使用標(biāo)準(zhǔn) FP32 操作的 Pascal GP100 相比，每 SM 深度學(xué)習(xí)應(yīng)用程序的吞吐量顯著提高了 8 倍，導(dǎo)致 Volta V100 GPU 的吞吐量比 Pascal P100 GPU 提高了 12 倍。張量核對(duì) FP16 輸入數(shù)據(jù)進(jìn)行 FP32 累加運(yùn)算。對(duì)于 4x4x4 矩陣乘法， FP16 乘法會(huì)產(chǎn)生一個(gè)全精度的結(jié)果，該結(jié)果在 FP32 運(yùn)算中與給定點(diǎn)積中的其他乘積累加，如圖 8 所示。

圖 2 ： Volta GV100 張量核心操作。

在程序執(zhí)行過(guò)程中，多個(gè)張量核被一個(gè)完整的執(zhí)行過(guò)程并發(fā)使用。扭曲中的線程提供了一個(gè)更大的 16x16x16 矩陣運(yùn)算，由張量核心處理。 CUDA 將這些操作暴露為 CUDA C ++ WMMA API 中的扭曲級(jí)別矩陣操作。這些 C ++接口提供專門(mén)的矩陣加載、矩陣乘法和累加運(yùn)算以及矩陣存儲(chǔ)操作，以有效地利用 CUDA C ++程序中的張量核。

但是在我們深入了解張量核心的低級(jí)編程細(xì)節(jié)之前，讓我們看看如何通過(guò) CUDA 庫(kù)訪問(wèn)它們的性能。

CUDA 庫(kù)中的張量核

使用張量核的兩個(gè) CUDA 庫(kù)是 cuBLAS 和 cuDNN 。 cuBLAS 使用張量核來(lái)加速 GEMM 計(jì)算（ GEMM 是矩陣矩陣乘法的 BLAS 項(xiàng)）； cuDNN 使用張量核來(lái)加速卷積和 遞歸神經(jīng)網(wǎng)絡(luò) 。

許多計(jì)算應(yīng)用都使用 GEMMs ：信號(hào)處理、流體力學(xué)和許多其他的。隨著這些應(yīng)用程序的數(shù)據(jù)大小呈指數(shù)級(jí)增長(zhǎng)，這些應(yīng)用程序需要匹配地提高處理速度。圖 3 中的混合精度 GEMM 性能圖表明張量核明確地滿足了這一需求。

提高卷積速度的需求同樣大；例如，今天的深度 神經(jīng)網(wǎng)絡(luò) （ DNNs ）使用了許多層卷積。人工智能研究人員每年都在設(shè)計(jì)越來(lái)越深的神經(jīng)網(wǎng)絡(luò)；現(xiàn)在最深的網(wǎng)絡(luò)中的卷積層數(shù)量已經(jīng)有幾十個(gè)。訓(xùn)練 dnn 需要在前向和反向傳播期間重復(fù)運(yùn)行卷積層。圖 4 中的卷積性能圖顯示張量核滿足了卷積性能的需要。（您或許也對(duì) 混合精度神經(jīng)網(wǎng)絡(luò)訓(xùn)練的有效技術(shù) 上的這篇文章感興趣）

兩個(gè)性能圖表都顯示， Tesla V100 的張量核心的性能是上一代 Tesla P100 的數(shù)倍。性能改進(jìn)這一巨大的改變了計(jì)算領(lǐng)域的工作方式：使交互成為可能，啟用“假設(shè)”場(chǎng)景研究，或者減少服務(wù)器場(chǎng)的使用。如果您在應(yīng)用程序中使用 GEMMs 或卷積，請(qǐng)使用下面的簡(jiǎn)單步驟來(lái)加速您的工作。

如何在 cuBLAS 中使用張量核

您可以利用張量核心，對(duì)現(xiàn)有的 cuBLAS 代碼進(jìn)行一些更改。這些更改是您使用 cuBLAS API 時(shí)所做的微小更改。

下面的示例代碼應(yīng)用了一些簡(jiǎn)單的規(guī)則來(lái)指示 cuBLAS 應(yīng)該使用張量核；這些規(guī)則在代碼后面顯式地枚舉。

示例代碼

下面的代碼在很大程度上與以前的架構(gòu)上用于調(diào)用 cuBLAS 中 GEMM 的通用代碼相同。

下面的代碼在很大程度上與以前的架構(gòu)上用于調(diào)用 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 ldcare 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);

一些簡(jiǎn)單的規(guī)則

cuBLAS 用戶會(huì)注意到他們現(xiàn)有的 cuBLAS GEMM 代碼有一些變化：

例程必須是 GEMM ；目前，只有 GEMM 支持 Tensor 核心執(zhí)行。

數(shù)學(xué)模式必須設(shè)置為 CUBLAS_TENSOR_OP_MATH 。浮點(diǎn)數(shù)學(xué)是非關(guān)聯(lián)的，因此張量核心數(shù)學(xué)例程的結(jié)果與類似的非張量核心數(shù)學(xué)例程的結(jié)果不完全對(duì)等。 cuBLAS 要求用戶選擇使用張量核。

k 、 lda 、 ldb 和 ldc 都必須是 8 的倍數(shù)； m 必須是 4 的倍數(shù)。張量核心數(shù)學(xué)例程以八個(gè)值的步長(zhǎng)跨越輸入數(shù)據(jù)，因此矩陣的維數(shù)必須是 8 的倍數(shù)。

矩陣的輸入和輸出數(shù)據(jù)類型必須是半精度或單精度。（上面只顯示了 CUDA_R_16F ，但也支持 CUDA_R_32F 。）

不滿足上述規(guī)則的 gemm 將返回到非張量核心實(shí)現(xiàn)。

GEMM 性能

如前所述， Tensor 內(nèi)核提供的 GEMM 性能是以前硬件的數(shù)倍。圖 3 顯示了 GP100 （ Pascal ）與 GV100 （ Volta ）硬件的比較性能。

圖 3 。使用張量核的 Tesla V100 （ Volta ）與 Tesla P100 （ Pascal ）的矩陣矩陣乘法（ GEMM ）的性能比較。輸入矩陣是半精度的，計(jì)算是單精度的。

如何在 cuDNN 中使用張量核

在 cuDNN 中使用張量核也很簡(jiǎn)單，而且只涉及對(duì)現(xiàn)有代碼的細(xì)微更改。

示例代碼

在 cuDNN 中使用張量核心的示例代碼可以在 cuDNN samples 目錄的 conv_sample.cpp 中找到；我們復(fù)制了下面的一些摘錄。（ cuDNN 樣本目錄 與文檔一起打包。）

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

一些簡(jiǎn)單的規(guī)則

注意一些與普通 cuDNN 用法不同的地方：

卷積算法必須是 ALGO_1 （ IMPLICIT_PRECOMP_GEMM 表示正向）。除了 ALGO_1 之外的其他卷積算法可能在未來(lái)的 cuDNN 版本中使用張量核。

數(shù)學(xué)類型必須設(shè)置為 CUDNN_TENSOR_OP_MATH 。與 cuBLAS 一樣，張量核心數(shù)學(xué)例程的結(jié)果與類似的非張量核心數(shù)學(xué)例程的結(jié)果并不完全等價(jià)，因此 cuDNN 要求用戶“選擇”使用張量核心。

輸入和輸出通道尺寸都必須是 8 的倍數(shù)。同樣，在 cuBLAS 中，張量核心數(shù)學(xué)例程以八個(gè)值的步長(zhǎng)跨越輸入數(shù)據(jù)，因此輸入數(shù)據(jù)的維數(shù)必須是 8 的倍數(shù)。

卷積的輸入、過(guò)濾和輸出數(shù)據(jù)類型必須為半精度。

不滿足上述規(guī)則的卷積將返回到非張量核心實(shí)現(xiàn)。

上面的示例代碼顯示了 NCHW 數(shù)據(jù)格式，請(qǐng)參見(jiàn) conv_sample.cpp NHWC 支持示例。

卷積性能

如前所述，張量核心的卷積性能是以前硬件的數(shù)倍。圖 4 顯示了 GP100 （ Pascal ）與 GV100 （ Volta ）硬件的比較性能。

圖 4 。張量核的 Tesla V100 （ Volta ）卷積與 Tesla P100 （ Pascal ）卷積的性能比較。比較來(lái)自每個(gè)神經(jīng)網(wǎng)絡(luò)的 卷積 層運(yùn)行時(shí)間的幾何平均值。 V100 和 P100 都使用 FP16 輸入/輸出數(shù)據(jù)和 FP32 計(jì)算； V100 使用張量核心，而 P100 使用 FP32 融合乘法加法（ FMA ）。

CUDA 9.0 中張量核的編程訪問(wèn)

通過(guò) CUDA 9.0 訪問(wèn)內(nèi)核中的張量核是一個(gè)預(yù)覽功能。這意味著本節(jié)中描述的數(shù)據(jù)結(jié)構(gòu)、 api 和代碼在未來(lái)的 CUDA 版本中可能會(huì)發(fā)生變化。

雖然 cuBLAS 和 cuDNN 覆蓋了張量核的許多潛在用途，但是您也可以直接在 nvcuda：：wmma C ++中編程它們。張量核心通過(guò) CUDA 命名空間中的一組函數(shù)和類型在 CUDA 9 。 0 中公開(kāi)。它們?cè)试S您將值加載或初始化為張量核心所需的特殊格式，執(zhí)行矩陣乘法累加（ MMA ）步驟，并將值存儲(chǔ)回內(nèi)存。在程序執(zhí)行過(guò)程中，一個(gè)完整的扭曲同時(shí)使用多個(gè)張量核。這允許 warp 在非常高的吞吐量下執(zhí)行 16x16x16mma （圖 5 ）。

圖 5 ： warp 執(zhí)行 D = A * B + C ，其中 A 、 B 、 C 和 D 是 16 × 16 矩陣。（注意圖 1 中編號(hào)的變化：多個(gè)張量核心操作由 WMMA API 組合，以執(zhí)行 16 × 16 矩陣乘法和累加運(yùn)算。）

讓我們看一個(gè)簡(jiǎn)單的例子，它展示了如何使用 WMMA （ Warp Matrix Multiply Accumulate ） API 來(lái)執(zhí)行矩陣乘法。注意，這個(gè)例子并沒(méi)有針對(duì)高性能進(jìn)行調(diào)整，主要是作為 API 的演示。為了獲得更好的性能， MIG ht 應(yīng)用于此代碼的優(yōu)化示例，請(qǐng)查看 CUDA 工具箱中的 cudaTensorCoreGemm 示例。為了獲得最高的生產(chǎn)性能，應(yīng)該使用 cuBLAS 代碼，如上所述。

標(biāo)題和命名空間

WMMA API 包含在 mma.h 頭文件中。完整的名稱空間是 nvcuda：：wmma：：* ，但是在代碼中保持 wmma 的顯式是很有用的，所以我們只使用 nvcuda 名稱空間。

#include 
using namespace nvcuda;

設(shè)計(jì)和初始化

完整的 GEMM 規(guī)范允許算法處理 a 或 b 的換位，并使數(shù)據(jù)跨距大于矩陣中的跨距。為了簡(jiǎn)單起見(jiàn)，讓我們假設(shè) a 和 b 都不是換位的，并且內(nèi)存和矩陣的前導(dǎo)維度是相同的。

我們將采用的策略是讓一個(gè) warp 負(fù)責(zé)輸出矩陣的單個(gè) 16 × 16 部分。通過(guò)使用二維網(wǎng)格和線程塊，我們可以有效地在二維輸出矩陣上平鋪扭曲。

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

在執(zhí)行 MMA 操作之前，操作數(shù)矩陣必須在 GPU 的寄存器中表示。由于 MMA 是一個(gè) warp 范圍的操作，這些寄存器分布在 warp 的線程中，每個(gè)線程持有整個(gè)矩陣的 片段 。單個(gè)矩陣參數(shù)與片段之間的映射是不透明的，因此您的程序不應(yīng)對(duì)此進(jìn)行假設(shè)。

在 CUDA 中，片段是一種模板化類型，其模板參數(shù)描述了片段持有的矩陣（ a 、 B 或累加器）、整體 WMMA 操作的形狀、數(shù)據(jù)類型，以及對(duì)于 a 和 B 矩陣，數(shù)據(jù)是行還是列主。最后一個(gè)參數(shù)可用于執(zhí)行 A 或 B 矩陣的換位。這個(gè)例子沒(méi)有換位，所以兩個(gè)矩陣都是列 major ，這是 GEMM 的標(biāo)準(zhǔn)。

 // Declare the fragments wmma::fragment a_frag; wmma::fragment b_frag; wmma::fragment acc_frag; wmma::fragment c_frag;

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

 wmma::fill_fragment(acc_frag, 0.0f);

內(nèi)環(huán)

我們用一個(gè)矩陣來(lái)計(jì)算每一個(gè)輸出的扭曲策略。為此，我們需要循環(huán) A 矩陣的行和 B 矩陣的列。這是沿著兩個(gè)矩陣的 K 維生成一個(gè) MxN 輸出塊。 loadmatrix 函數(shù)從內(nèi)存（在本例中是全局內(nèi)存，盡管可以是任何內(nèi)存空間）中獲取數(shù)據(jù)并將其放入片段中。加載的第三個(gè)參數(shù)是矩陣內(nèi)存中的“前導(dǎo)維度”；我們加載的 16 × 16 塊在內(nèi)存中是不連續(xù)的，因此函數(shù)需要知道連續(xù)列（或行，如果這些是行的主要片段）之間的跨距。

MMA 調(diào)用就地累積，因此第一個(gè)參數(shù)和最后一個(gè)參數(shù)都是我們先前初始化為零的累加器片段。

 // 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 現(xiàn)在基于 A 和 B 的乘法保存此扭曲的輸出塊的結(jié)果。完整的 GEMM 規(guī)范允許縮放此結(jié)果，并將其累積到適當(dāng)?shù)木仃図敳?。?shí)現(xiàn)這種縮放的一種方法是對(duì)片段執(zhí)行元素級(jí)操作。雖然沒(méi)有定義從矩陣坐標(biāo)到線程的映射，但是元素級(jí)操作不需要知道這個(gè)映射，所以仍然可以使用片段來(lái)執(zhí)行。因此，對(duì)片段執(zhí)行縮放操作或?qū)⒁粋€(gè)片段的內(nèi)容添加到另一個(gè)片段是合法的，只要這兩個(gè)片段具有相同的模板參數(shù)。如果片段具有不同的模板參數(shù)，則結(jié)果未定義。使用這個(gè)特性，我們將現(xiàn)有的數(shù)據(jù)加載到 C 語(yǔ)言中，并使用正確的縮放比例來(lái)累積到目前為止的計(jì)算結(jié)果。

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

最后，我們將數(shù)據(jù)存儲(chǔ)到內(nèi)存中。同樣，目標(biāo)指針可以是 GPU 可見(jiàn)的任何內(nèi)存空間，并且必須指定內(nèi)存中的前導(dǎo)維度。還有一個(gè)選項(xiàng)可以指定輸出是寫(xiě)在行還是列 major 。

 // Store the output wmma::store_matrix_sync(c + cRow + cCol * ldc, c_frag, ldc, wmma::mem_col_major); }
}

這樣，矩陣乘法就完成了。我在這篇博文中省略了主機(jī)代碼，不過(guò)是一個(gè) 完整的工作示例可以在 Github 上找到 。

今天就從 CUDA 9 中的張量核心開(kāi)始吧

希望這個(gè)例子能讓您了解如何在應(yīng)用程序中使用張量核。

關(guān)于作者

Jeremy Appleyard 是 NVIDIA 歐洲開(kāi)發(fā)人員技術(shù)團(tuán)隊(duì)的一名開(kāi)發(fā)人員。他位于英國(guó)牛津附近，與開(kāi)發(fā)人員一起加速 GPUs 上的應(yīng)用程序。他擁有克蘭菲爾德大學(xué)計(jì)算流體力學(xué)博士學(xué)位。

Scott Yokim 是 NVIDIA 的 CUDA 庫(kù)團(tuán)隊(duì)的高級(jí)軟件工程師。他于 2008 年加入 NVIDIA ，在此之前，他是多家公司的計(jì)算機(jī)圖形程序員。斯科特?fù)碛懈ゼ醽喞砉ご髮W(xué)數(shù)學(xué)碩士學(xué)位。

審核編輯：郭婷