如何使用张量核在CUDA C++设备代码中编程

新 Volta GPU 架构的一个定义性特征是它的 张量核 ，它使 Tesla V100 加速器的峰值吞吐量是上一代 Tesla P100 的 32 位浮点吞吐量的 12 倍。张量核心使人工智能程序员能够使用 混合精度 来实现更高的吞吐量而不牺牲精度。

张量核心已经在主版本或许多深度学习框架（包括 PyTorch 、 TensorFlow 、 MXNet 和 Caffe2 ）中通过 pull 请求支持 深度学习 培训。有关在使用这些框架时启用张量核心的更多信息，请查看 混合精度训练指南 。

在这篇博客文章中，我们展示了如何使用 CUDA 库在自己的应用程序中使用张量核，以及如何直接在 CUDA C ++设备代码中编程。

什么是张量核？

Tesla V100 的张量核心是可编程的矩阵乘法和累加单元，可为训练和推理应用提供多达 125 个张量 TFLOP 。 Tesla V100GPU 包含 640 个张量核心：每平方米 8 个。张量核心及其相关数据路径都是定制的，可以显著提高浮点计算吞吐量，只需适度的面积和功耗成本。时钟门控广泛用于最大限度地节省电力。

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

图 1 ：张量核 4x4x4 矩阵乘法和累加。

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

图 2 ： Volta GV100 张量核心操作。

在程序执行过程中，多个张量核被一个完整的执行过程并发使用。扭曲中的线程提供了一个更大的 16x16x16 矩阵运算，由张量核心处理。 CUDA 将这些操作暴露为 CUDA C ++ WMMA API 中的扭曲级别矩阵操作。这些 C ++接口提供专门的矩阵加载、矩阵乘法和累加运算以及矩阵存储操作，以有效地利用 CUDA C ++程序中的张量核。

但是在我们深入了解张量核心的低级编程细节之前，让我们看看如何通过 CUDA 库访问它们的性能。

CUDA 库中的张量核

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

许多计算应用都使用 GEMMs ：信号处理、流体力学和许多其他的。随着这些应用程序的数据大小呈指数级增长，这些应用程序需要匹配地提高处理速度。图 3 中的混合精度 GEMM 性能图表明张量核明确地满足了这一需求。

提高卷积速度的需求同样大；例如，今天的深度 神经网络 （ DNNs ）使用了许多层卷积。人工智能研究人员每年都在设计越来越深的神经网络；现在最深的网络中的卷积层数量已经有几十个。训练 dnn 需要在前向和反向传播期间重复运行卷积层。图 4 中的卷积性能图显示张量核满足了卷积性能的需要。（您或许也对 混合精度神经网络训练的有效技术 上的这篇文章感兴趣）

两个性能图表都显示， Tesla V100 的张量核心的性能是上一代 Tesla P100 的数倍。性能改进这一巨大的改变了计算领域的工作方式：使交互成为可能，启用“假设”场景研究，或者减少服务器场的使用。如果您在应用程序中使用 GEMMs 或卷积，请使用下面的简单步骤来加速您的工作。

如何在 cuBLAS 中使用张量核

您可以利用张量核心，对现有的 cuBLAS 代码进行一些更改。这些更改是您使用 cuBLAS API 时所做的微小更改。

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

示例代码

下面的代码在很大程度上与以前的架构上用于调用 cuBLAS 中 GEMM 的通用代码相同。

下面的代码在很大程度上与以前的架构上用于调用 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);

一些简单的规则

cuBLAS 用户会注意到他们现有的 cuBLAS GEMM 代码有一些变化：

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

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

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

矩阵的输入和输出数据类型必须是半精度或单精度。（上面只显示了 CUDA_R_16F ，但也支持 CUDA_R_32F 。）

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

GEMM 性能

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

图 3 。使用张量核的 Tesla V100 （ Volta ）与 Tesla P100 （ Pascal ）的矩阵矩阵乘法（ GEMM ）的性能比较。输入矩阵是半精度的，计算是单精度的。

如何在 cuDNN 中使用张量核

在 cuDNN 中使用张量核也很简单，而且只涉及对现有代码的细微更改。

示例代码

在 cuDNN 中使用张量核心的示例代码可以在 cuDNN samples 目录的 conv_sample.cpp 中找到；我们复制了下面的一些摘录。（ 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) );

一些简单的规则

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

卷积算法必须是 ALGO_1 （ IMPLICIT_PRECOMP_GEMM 表示正向）。除了 ALGO_1 之外的其他卷积算法可能在未来的 cuDNN 版本中使用张量核。

数学类型必须设置为 CUDNN_TENSOR_OP_MATH 。与 cuBLAS 一样，张量核心数学例程的结果与类似的非张量核心数学例程的结果并不完全等价，因此 cuDNN 要求用户“选择”使用张量核心。

输入和输出通道尺寸都必须是 8 的倍数。同样，在 cuBLAS 中，张量核心数学例程以八个值的步长跨越输入数据，因此输入数据的维数必须是 8 的倍数。

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

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

上面的示例代码显示了 NCHW 数据格式，请参见 conv_sample.cpp NHWC 支持示例。

卷积性能

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

图 4 。张量核的 Tesla V100 （ Volta ）卷积与 Tesla P100 （ Pascal ）卷积的性能比较。比较来自每个神经网络的 卷积 层运行时间的几何平均值。 V100 和 P100 都使用 FP16 输入/输出数据和 FP32 计算； V100 使用张量核心，而 P100 使用 FP32 融合乘法加法（ FMA ）。

CUDA 9.0 中张量核的编程访问

通过 CUDA 9.0 访问内核中的张量核是一个预览功能。这意味着本节中描述的数据结构、 api 和代码在未来的 CUDA 版本中可能会发生变化。

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

图 5 ： warp 执行 D = A * B + C ，其中 A 、 B 、 C 和 D 是 16 × 16 矩阵。（注意图 1 中编号的变化：多个张量核心操作由 WMMA API 组合，以执行 16 × 16 矩阵乘法和累加运算。）

让我们看一个简单的例子，它展示了如何使用 WMMA （ Warp Matrix Multiply Accumulate ） API 来执行矩阵乘法。注意，这个例子并没有针对高性能进行调整，主要是作为 API 的演示。为了获得更好的性能， MIG ht 应用于此代码的优化示例，请查看 CUDA 工具箱中的 cudaTensorCoreGemm 示例。为了获得最高的生产性能，应该使用 cuBLAS 代码，如上所述。

标题和命名空间

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

#include 
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 是一个 warp 范围的操作，这些寄存器分布在 warp 的线程中，每个线程持有整个矩阵的 片段 。单个矩阵参数与片段之间的映射是不透明的，因此您的程序不应对此进行假设。

在 CUDA 中，片段是一种模板化类型，其模板参数描述了片段持有的矩阵（ a 、 B 或累加器）、整体 WMMA 操作的形状、数据类型，以及对于 a 和 B 矩阵，数据是行还是列主。最后一个参数可用于执行 A 或 B 矩阵的换位。这个例子没有换位，所以两个矩阵都是列 major ，这是 GEMM 的标准。

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

内环

我们用一个矩阵来计算每一个输出的扭曲策略。为此，我们需要循环 A 矩阵的行和 B 矩阵的列。这是沿着两个矩阵的 K 维生成一个 MxN 输出块。 loadmatrix 函数从内存（在本例中是全局内存，尽管可以是任何内存空间）中获取数据并将其放入片段中。加载的第三个参数是矩阵内存中的“前导维度”；我们加载的 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 可见的任何内存空间，并且必须指定内存中的前导维度。还有一个选项可以指定输出是写在行还是列 major 。

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

这样，矩阵乘法就完成了。我在这篇博文中省略了主机代码，不过是一个 完整的工作示例可以在 Github 上找到 。

今天就从 CUDA 9 中的张量核心开始吧

希望这个例子能让您了解如何在应用程序中使用张量核。

关于作者

Jeremy Appleyard 是 NVIDIA 欧洲开发人员技术团队的一名开发人员。他位于英国牛津附近，与开发人员一起加速 GPUs 上的应用程序。他拥有克兰菲尔德大学计算流体力学博士学位。

Scott Yokim 是 NVIDIA 的 CUDA 库团队的高级软件工程师。他于 2008 年加入 NVIDIA ，在此之前，他是多家公司的计算机图形程序员。斯科特拥有弗吉尼亚理工大学数学硕士学位。

审核编辑：郭婷