cutlass cuda c++实现快速线性代数
矩阵乘法是许多科学应用中的核心计算，特别是在深度学习中。现代深度神经网络中的许多操作要么被定义为矩阵乘法，要么可以被转换为矩阵乘法。
例如，NVIDIA cuDNN库使用各种形式的矩阵乘法实现了神经网络的卷积，例如直接卷积的经典公式，即图像到列和滤波器数据集之间的矩阵乘积。当基于快速傅里叶变换（FFT）或Winograd方法[3]计算卷积时，矩阵乘法也是核心例程。
在构建cuDNN时，我们从cuBLAS库中通用矩阵乘法（GEMM）的高性能实现开始，对其进行补充和定制，以高效计算卷积。如今，我们调整这些GEMM策略和算法的能力对于为深度学习中的许多不同问题和应用提供最佳性能至关重要。
CUTALSS为在CUDA C++中使用高性能GEMM构建块开发使用新算法所需的技术和结构。密集线性代数的灵活高效应用在深度学习和更广泛的GPU计算生态系统中至关重要
CUTLASS介绍
这篇文章将介绍CUTALSS（线性代数子程序的CUDA模板）的预览版，这是CUDA C++模板和抽象的集合，用于在CUDA内核中实现各级和各种规模的高性能GEMM计算。与其他用于密集线性代数的模板GPU库（例如MAGMA库）不同，CUTALSS的目的是将GEMM的“运动部分”分解为C++模板类抽象的基本组件，使程序员能够在自己的CUDA内核中轻松定制和实例化它们。
我们的CUTLASS原语包括对混合精度计算的广泛支持，提供专门的数据移动和乘法累加抽象，用于处理8位整数、半精度浮点（FP16）、单精度浮点（FP32）和双精度浮点（FP64）类型。CUTLASS最令人兴奋的功能之一是使用WMMA API在Volta架构中的新Tensor内核上运行矩阵乘法。Tesla V100的Tensor Core是可编程的矩阵乘法和累加单元，可以高效地提供高达125个Tensor TFLOP/s。
GPU上的高效矩阵乘法
GEMM计算C=αA*B+βC，其中A、B和C是矩阵。A是M乘K矩阵，B是K乘N矩阵，C是M乘N矩阵。为简单起见，让我们在以下示例中假设标量α=β=1。稍后，我们将展示如何使用支持任意缩放函数的CUTLASS实现自定义元素操作
最简单的一种实现可以使用3层嵌套循环
for (int i = 0; i < M; ++i)
for (int j = 0; j < N; ++j)
for (int k = 0; k < K; ++k)
C[i][j] += A[i][k] * B[k][j];位置（i，j）处的C元素是A的第i行和B的第j列的K元素点积。理想情况下，性能应受到处理器算术吞吐量的限制。事实上，对于M=N=K的大型方阵，矩阵乘积中的数学运算次数为O（N3），而所需的数据量为O（N2），计算强度约为N。然而，达到理论计算强度需要重用每个元素O（N）次。不幸的是，上述“内积”算法依赖于在快速片上缓存中保存大量工作集，这会在M、N和K增长时产生抖动。
更好的公式是通过在K维上将循环结构化为最外层的循环来排列循环嵌套。这种计算形式加载一列a和一行B一次，计算其外积，并将该外积的结果累积在矩阵C中。之后，这列a和这行B将不再使用。
for (int k = 0; k < K; ++k) // K dimension now outer-most loop
for (int i = 0; i < M; ++i)
for (int j = 0; j < N; ++j)
C[i][j] += A[i][k] * B[k][j];这种方法的一个问题是，它要求C的所有M-by-N元素都是活的，以存储每条乘法累加指令的结果，理想情况下，存储在可以像乘法累加指令一样快速写入的内存中。我们可以通过按Ntile将C划分为大小为Mtile的块来减小C的工作集大小，这些块保证适合片上存储器。然后，我们将“外产品”配方应用于每块tile。这将导致以下循环嵌套。
for (int m = 0; m < M; m += Mtile) // iterate over M dimension
for (int n = 0; n < N; n += Ntile) // iterate over N dimension
for (int k = 0; k < K; ++k)
for (int i = 0; i < Mtile; ++i) // compute one tile
for (int j = 0; j < Ntile; ++j) {
int row = m + i;
int col = n + j;
C[row][col] += A[row][k] * B[k][col];
}对于C的每个图块，A和B的图块被精确地提取一次，从而达到O（N）的计算强度。C的每个图块的大小可以被选择为与目标处理器的L1缓存或寄存器的容量相匹配，并且嵌套的外部循环可以被简单地并行化。这是一个很大的进步！
进一步的重组为利用局部性和并行性提供了额外的机会。我们可以通过在块中逐步遍历K维来累积矩阵的乘积，而不是只累积向量外积。我们通常将这个概念称为累积矩阵积。
分层GEMM结构
CUTLASS通过将计算分解为线程块tile、warp tile和线程 tile的层次结构，并应用累积矩阵乘积的策略，将图块结构应用于GPU，以高效地实现GEMM。此层次结构与NVIDIA CUDA编程模型非常相似，如图1所示。在这里，您可以看到数据从全局内存移动到共享内存（矩阵到线程块块图块），从共享内存移动到寄存器文件（线程块图块到warp图块）以及从寄存器文件移动到CUDA核心进行计算（warp图块到线程图块）。
线程块Tile
每个线程块通过迭代地从输入矩阵加载矩阵数据块并计算累积矩阵乘积（C+=A*B）来计算其输出GEMM的一部分。图2显示了单线程块执行的计算，并突出显示了其主循环的一次迭代中使用的数据块。
CUDA线程块图块结构进一步划分为warp（以SIMT方式一起执行的线程组）。Warps为GEMM计算提供了一个有用的组织，并且是WMMA API的明确组成部分，我们将在稍后讨论。
下图显示了一个块级矩阵乘积结构的详细视图。A和B的图块从全局内存加载，并存储到所有扭曲可访问的共享内存中。如图3所示，线程块的输出图块在warp上进行了空间分区。我们将此输出图块的存储称为累加器，因为它存储了累加矩阵乘积的结果。每个累加器在每次数学运算中都会更新一次，因此它需要驻留在SM中最快的内存中：寄存器文件(register file)。
参数BlockItems{X，Y，K}是程序员指定的编译时常数，用于调整目标处理器的GEMM计算和特定GEMM配置的纵横比（例如M，N，K，数据类型等）。在图中，我们展示了一个8 warp、256线程的线程块，这是CUTALSS中实现的大型SGEMM（FP32 GEMM）图块的典型特征。
warp tile
一旦数据存储在共享内存中，每个warp都会通过迭代线程块图块的K维、从共享内存加载子矩阵（或片段）并计算累积的外积来计算一系列累积的矩阵积。图4显示了详细视图。片段的大小在K维上通常非常小，以最大化相对于从共享内存加载的数据量的计算强度，从而避免共享内存带宽成为瓶颈。
上图还描述了几个warp之间共享内存的数据共享。同一行线块中的warp加载相同的A片段，同一列中的warp则加载相同的B片段。
我们注意到，GEMM结构的以warp为中心的组织在实现高效的GEMM内核方面是有效的，但并不依赖于隐式warp同步执行来实现同步。CUTLASS GEMM内核与__syncthreads（）的调用同步良好。
Thread tile
CUDA编程模型是根据线程块和单个线程定义的。因此，warp结构被映射到由单个线程执行的操作上。线程不能访问彼此的寄存器，因此我们必须选择一个组织，使寄存器中保存的值能够被同一线程执行的多个数学指令重用。这导致线程内出现2D平铺结构，如图5中的详细视图所示。每个线程向CUDA内核发出一系列独立的数学指令，并计算累积的外积。
在上图中，warp的左上象限以灰色着色。32个单元格对应于warp内的32个线程。这种安排导致同一行或同一列中的多个线程分别获取A和B片段的相同元素。为了最大限度地提高计算强度，可以复制这种基本结构以形成完整的warp水平累加器图块，从而产生一个8乘8的整体线程图块，该图块由8乘1和1乘8片段的外积计算得出。这由绿色显示的四个累加器图块说明。
WMMA GEMM
warp tile结构可以用CUDA 9中引入的CUDA warp矩阵多累加API（WMMA）来实现，以针对Volta V100 GPU的Tensor内核。有关WMMA API的更多详细信息，请参阅CUDA9中的Tensor内核编程
每个Tensor Core都提供一个4x4x4矩阵处理数组，执行操作D=a*B+C，其中a、B、C和D是4×4矩阵，如图6所示。矩阵乘法输入A和B是FP16矩阵，而累加矩阵C和D可以是FP16或FP32矩阵。
下图显示了以CUDA WMMA API为目标的warp tile结构。对wmma:：load_matrix_sync的调用将A和B的片段加载到nvcuda:：wmma:；fragment<>模板的实例中，warp图块的累加器元素被构造为nvcuda：：wmma：fragment<accumulator>对象的数组。这些片段存储了一个分布在warp线之间的二维矩阵。最后，对每个累加器片段（以及来自A和B的相应片段）调用nvcuda:：wmma:：mma_sync（），使用Tensor Core计算warp范围的矩阵乘法累加运算。
CUTLASS在文件block_task_WMMA.h中基于WMMA API实现GEMM。warp tile的维度必须是目标CUDA计算能力的nvcuda:：WMMA模板定义的矩阵多累加形状的倍数。在CUDA 9.0中，WMMA的基本尺寸为16乘16乘16。
完成GEMM
完整的GEMM结构可以表示为由线程块的线程执行的嵌套循环，如下表所示。除最外层的“主”循环外，所有循环都有恒定的迭代计数，编译器可以完全展开。为简洁起见，这里省略了地址和索引计算，但在CUTLASS源代码中进行了解释。
// Device function to compute a thread block’s accumulated matrix product
__device__ void block_matrix_product(int K_dim) {
// Fragments used to store data fetched from SMEM
value_t frag_a[ThreadItemsY];
value_t frag_b[ThreadItemsX];
// Accumulator storage
accum_t accumulator[ThreadItemsX][ThreadItemsY];
// GEMM Mainloop - iterates over the entire K dimension - not unrolled
for (int kblock = 0; kblock < K_dim; kblock += BlockItemsK) {
// Load A and B tiles from global memory and store to SMEM
//
// (not shown for brevity - see the CUTLASS source for more detail)
...
__syncthreads();
// Warp tile structure - iterates over the Thread Block tile
#pragma unroll
for (int warp_k = 0; warp_k < BlockItemsK; warp_k += WarpItemsK) {
// Fetch frag_a and frag_b from SMEM corresponding to k-index
//
// (not shown for brevity - see CUTLASS source for more detail)
...
// Thread tile structure - accumulate an outer product
#pragma unroll
for (int thread_x = 0; thread_x < ThreadItemsX; ++thread_x) {
#pragma unroll
for (int thread_y=0; thread_y < ThreadItemsY; ++thread_y) {
accumulator[thread_x][thread_y] += frag_a[y]*frag_b[x];
}
}
}
__syncthreads();
}
}WarpItemsK是指目标数学运算的点积大小。对于SGEMM（FP32 GEMM）、DGEM（FP64）和HGEMM（FP 16），标量乘累加指令的点积长度为1。对于IGEMM（8位整数GEMM），CUTLASS的目标是WarpItemsK=4的四元素整数点积指令（IDP4A）。对于基于WMMA的GEMM，我们选择WMMA:：fragment<>模板的K维。目前，这被定义为WarpItemsK=16。
软件流水线
分片矩阵乘广泛使用寄存器文件来保存片段和累加器分片以及大型共享内存分配。相对较高的片上存储需求限制了占用率，即一个SM上可以并发运行的线程块的最大数量。因此，GEMM实现可以在每个SM中容纳比GPU计算工作负载典型的更少的warp和线程块。我们使用软件流水线来隐藏数据移动延迟，方法是在一个循环中同时执行GEMM层次结构的所有阶段，并在下一次迭代中将每个阶段的输出输送到其依赖阶段，如下图所示。
GEMM CUDA内核在管道中发出三个并发的操作流，这些操作流对应于GEMM层次结构中的数据流阶段（图1）。图中每个阶段的相对大小表示操作的延迟是长还是短，橙色箭头突出显示了每个流阶段之间的数据依赖关系。在数据存储到共享内存后调用__syncthreads（）会同步所有warp，这样它们就可以在没有竞争条件的情况下读取共享内存。流水线的最后一个数学阶段与来自共享内存的负载重叠，共享内存将数据将输送到下一个主循环迭代的第一个数学阶段。
实际上，CUDA程序员通过在程序文本中为每个阶段交织CUDA语句，并依靠CUDA编译器在编译后的代码中发出正确的指令调度，来实现管道阶段之间的指令级并发。#pragma unroll和编译时常量的广泛使用使CUDA编译器能够展开循环并将数组元素映射到寄存器，这两者对于可调、高效的实现至关重要。有关示例，请参见block_task:：consume_tile（）。
我们在GEMM层次结构的每个级别使用双缓冲，使上游流水线级能够在依赖流水线级从其存储元素加载时将数据写入共享内存或寄存器。值得注意的是，这消除了第二个__syncthreads（），因为一个共享内存缓冲区被写入，而另一个被读取。双缓冲的成本是共享内存容量的两倍，是用于保存共享内存读取的寄存器数量的两倍。
实际可用的延迟隐藏量取决于线程块、warp和线程分片的大小，以及SM内活动数学功能单元的吞吐量。虽然较大的分片为数据重用提供了更多机会，并可能提供更多的延迟隐藏，但SM寄存器文件和共享内存的物理容量限制了最大分片的大小。幸运的是，NVIDIA GPU有足够的存储资源来执行足够大的GEMM分片，以限制数学运算！
CUTLASS
CUTALSS是分层GEMM结构作为CUDA C++模板类的实现。我们打算将这些模板包含在现有的设备端CUDA内核和函数中，但我们也提供了一个示例内核和启动接口，以便快速启动和运行。与CUB一样，广泛使用模板参数和编译时常量使CUTALSS具有可调性和灵活性。
CUTLASS实现了高效GEMM实现所需操作的抽象。专门的“图块加载器”将数据从全局内存高效地移动到共享内存中，既适应源数据的布局，又能够高效、无冲突地加载到寄存器中。对于某些布局，IGEM需要对数据进行一些重构，以针对CUDA的4元素整数点积指令，这是在数据存储到SMEM时完成的。
CUTLASS GEMM Device Functions
下面来自dispatch.h的示例定义了一个block_task类型，并为浮点数据实例化了一个GEMM，假设列为主要输入矩阵。block_task_policy_t定义了GEMM图块大小，下一节将详细讨论。
/// CUTLASS SGEMM example
__global__ void gemm_kernel(
float *C,
float const *A,
float const *B,
int M,
int N,
int K) {
// Define the GEMM tile sizes - discussed in next section
typedef block_task_policy <
128, // BlockItemsY: Height in rows of a tile
32, // BlockItemsX - Width in columns of a tile
8, // ThreadItemsY - Height in rows of a thread-tile
4, // ThreadItemsX - Width in columns of a thread-tile
8, // BlockItemsK - Depth of a tile
true, // UseDoubleScratchTiles - whether to double-buffer SMEM
block_raster_enum::Default // Block rasterization strategy
> block_task_policy_t;
// Define the epilogue functor
typedef gemm::blas_scaled_epilogue<float, float, float> epilogue_op_t ;
// Define the block_task type.
typedef block_task <
block_task_policy_t,
float,
float,
matrix_transform_t::NonTranspose,
4,
matrix_transform_t::NonTranspose,
4,
epilogue_op_t,
4,
true
> block_task_t;
// Declare statically-allocated shared storage
__shared__ block_task_t::scratch_storage_t smem;
// Construct and run the task
block_task_t(
reinterpret_cast(&smem),
&smem,
A,
B,
C,
epilogue_op_t(1, 0),
M,
N,
K).run();
}block_task_t实例使用共享内存分配smem在矩阵乘积计算中存储线程块级图块。
epilogue_op_t是一个模板参数，它指定了一个函数子，用于在矩阵乘法运算完成后更新输出矩阵。这使您可以使用自定义元素操作轻松组合矩阵乘法，我们稍后将对此进行更详细的描述。CUTLASS提供了gemm:：blas_scaled.epilogue函数器实现，用于计算熟悉的gemm操作C=alpha*AB+beta*C（在epilogue_function.h中定义）。
CUTLASS GEMM策略
CUTLASS组织编译时常量，在GEMM层次结构的每个级别指定图块大小，作为GEMM:：block_task_policy模板的特例化，该模板具有以下声明。
template <
int BlockItemsY, /// Height in rows of a tile in matrix C
int BlockItemsX, /// Width in columns of a tile in matrix C
int ThreadItemsY, /// Height in rows of a thread-tile in C
int ThreadItemsX, /// Width in columns of a thread-tile in C
int BlockItemsK, /// Number of K-split subgroups in a block
bool UseDoubleScratchTiles, /// Whether to double buffer shared memory
grid_raster_strategy::kind_t RasterStrategy /// Grid <a href="https://developer.nvidia.com/discover/ray-tracing">rasterization</a> strategy
> struct block_task_policy;dispatch_Policies.h中定义了几个有效GEMM分块结构的策略，我们在下面展示了一个这样的策略。此策略将矩阵乘法运算分解为CUDA块，每个CUDA块跨越输出矩阵的128乘32块。存储a和B的线程块块块的大小分别为128乘8和8乘32。该策略针对GEMM计算进行了优化，其中C矩阵的N维相对较窄。
/// GEMM task policy specialization for tall SGEMM
template <>
struct gemm_policy<float, float, problem_size_t::Tall> :
block_task_policy<
128, // BlockItemsY - Height in rows of a tile
32, // BlockItemsX - Width in columns of a tile
8, // ThreadItemsY - Height in rows of a thread-tile
4, // ThreadItemsX - Width in columns of a thread-tile
8, // BlockItemsK - Depth of a tile
true, // UseDoubleScratchTiles - whether to double-buffer SMEM
grid_raster_strategy::Default> // Grid rasterization strategy
{};线程图块片段的大小分别为ThreadItems Y-by-1和ThreadItems X-by-1。在上述示例的情况下，这些被给出为来自A的8乘1矢量和来自B的4乘1矢量。
定义了策略类型后，我们可以定义gemm:：block_task的类型，这是一个CUTLASS gemm。此模板具有以下参数列表。
template <
/// Parameterization of block_task_policy
typename block_task_policy_t,
/// Multiplicand value type (matrices A and B)
typename value_t,
/// Accumulator value type (matrix C and scalars)
typename accum_t,
/// Layout enumerant for matrix A
matrix_transform_t::kind_t TransformA,
/// Alignment (in bytes) for A operand
int LdgAlignA,
/// Layout enumerant for matrix B
matrix_transform_t::kind_t TransformB,
/// Alignment (in bytes) for B operand
int LdgAlignB,
/// Epilogue functor applied to matrix product
typename epilogue_op_t,
/// Alignment (in bytes) for C operand
int LdgAlignC,
/// Whether GEMM supports matrix sizes other than mult of BlockItems{XY}
bool Ragged
> struct block_task;value_t和accum_t分别指定源操作数和累加器矩阵的类型。TransformA和TransformB分别指定操作数A和B的布局。虽然我们没有详细讨论矩阵布局，但CUTLASS支持行主和列主输入矩阵的所有组合。
LdgAlignA和LdgAlign B指定了有保证的对齐，这使得CUTLASS设备代码能够使用向量内存操作。例如，8字节的对齐允许CUTLASS在两个元素向量中加载float类型的元素。这通过减少GPU内正在进行的内存操作的数量来减少代码大小并提高性能。更关键的是，参差不齐的处理表明矩阵维度是否可以是任意大小（满足对齐保证）。如果此模板参数为false，则矩阵A、B和C的维度都应该是block_task_policy中图块参数的倍数。
将元素操作与SGEMM融合
深度学习计算通常在GEMM计算之后执行简单的元素操作，例如计算激活函数。这些带宽受限的层可以融合到GEMM操作的末尾，以消除额外的内核启动，避免全局内存的往返。
以下示例演示了GEMM模板的简单应用，该模板将偏差项添加到缩放矩阵乘积中，然后应用ReLU函数将结果箝位为非负值。通过将结语分离为一个functor，传递参数（如指向其他矩阵和张量参数的指针或其他比例因子）是简单的，不会妨碍GEMM实现。
首先，我们定义一个实现gemm:：epilogue_op概念的类。这里没有显示构造函数和其他方法，但在函数调用运算符的实现中显示了元素偏置和ReLU操作。
template <typename accum_t, typename scalar_t, typename output_t>
struct fused_bias_relu_epilogue {
// Data members pass additional arguments to epilogue
scalar_t const *Bias;
accum_t threshold;
/// Constructor callable on host and device initializes data members
inline __device__ __host__
fused_bias_relu_epilogue(
scalar_t const *Bias,
accum_t threshold
): Bias(Bias), threshold(threshold) { }
/// Applies bias + ReLu operation
inline __device__ __host__
output_t operator()(
accum_t accumulator, /// element of matrix product result
output_t c, /// element of source accumulator matrix C
size_t idx /// index of c element; may be used to load
/// elements from other identically-
/// structured matrices
) const {
// Compute the result by scaling the matrix product, adding bias,
// and adding the scaled accumulator element.
accum_t result = output_t(
alpha * scalar_t(accumulator) +
Bias[i] + // load and add the bias
beta * scalar_t(c)
);
// apply clamping function
return max(threshold, result);
}
};然后我们应用该运算符作为结束操作。
// New: define type for custom epilogue functor
typedef fused_bias_relu_epilogue_t<float, float, float>
bias_relu_epilogue_t;
/// Computes GEMM fused with Bias and ReLu operation
__global__ void gemm_bias_relu(
..., /// GEMM parameters not shown
bias_relu_epilogue_t bias_relu_op) { /// bias_relu_op constructed
/// by caller
// Define the block_task type.
typedef block_task<
block_task_policy_t, // same policy as previous example
float,
float,
matrix_transform_t::NonTranspose,
4,
matrix_transform_t::NonTranspose,
4,
bias_relu_epilogue_t, // New: custom epilogue functor type
4,
true
> block_task_t ;
// Declare statically-allocated shared storage
__shared__ block_task_t::scratch_storage_t smem;
// Construct and run the task
block_task_t(
reinterpret_cast(&smem),
&smem,
A,
B,
C,
bias_relu_op, // New: custom epilogue object
M,
N,
K).run();
}这个简单的例子展示了将通用编程技术与高效的GEMM实现相结合的价值。
Tesla V100 (Volta) Performance
CUTLASS非常高效，在标量GEMM计算方面的性能可与cuBLAS相媲美。图9显示了CULASS相对于CUDA 9.0编译的cuBLAS的性能，该cuBLAS在NVIDIA Tesla V100 GPU上运行，适用于大矩阵尺寸（M=10240，N=K=4096）。图9显示了CUTALSS支持的每种计算数据类型的相对性能，以及输入操作数的行主布局和列主布局的所有排列。
在大多数情况下，CUTLASS C++的性能与cuBLAS中手动调优的汇编内核的性能相差不到百分之几。对于WMMA GEMM（图9中的WGEMM），CUTALSS尚未达到与cuBLAS相同的性能，但我们正在与CUDA编译器和GPU架构团队密切合作，开发技术，以在CUDA代码中达到类似的性能水平。
Try CUTLASS Today!
我们在这篇博客文章中没有提到许多有趣的细节，所以我们建议您查看CUTLASS存储库并亲自尝试CUTLASS。cutlass_test示例程序演示了如何调用CUTALSS GEMM内核，验证其结果并测量其性能。