实现高性能#

下面我们介绍一些通用建议和示例，这些建议和示例可能有助于实现高性能。

通用建议#

  • 从库提供的默认设置开始，以获得最佳计算性能。

  • 使用寄存器片段 API 而不是共享内存 API。

  • 计算密集型和内存密集型内核的最佳参数可能不相同。

  • 如果可能，请确保 BLAS 操作是批处理的，以便在网格中运行足够的 CUDA 块来充分利用 GPU 以获得峰值性能。

  • 将相邻的内存密集型内核（预处理和后处理）与 BLAS 内核合并，以节省全局内存访问。

  • 如果可能，请使用具有 MMA 支持 的 精度 组合。

  • 使用张量 execute(…) API 以获得更好的 IO 和矩阵乘法性能。

  • 使用 cublasdx::copy 复制共享内存和全局内存张量。它应自动向量化加载和存储。

  • 使用 cublasdx::copy_fragment 将寄存器片段复制到共享内存和全局内存张量。它应自动向量化加载和存储。

  • 使用 16 字节（128 位）对齐的指针，并使用 MaxAlignment（Alignment<16, 16, 16> 的别名）。

矩阵布局#

我们建议将张量 API 与 get_layout_smem_*()、suggest_layout_smem_*() 和 cublasdx::make_tensor 函数一起使用。这允许为矩阵使用自定义布局，从而可以提供更好的性能，并更好地匹配您的内核。

  • 尝试使用 suggest_layout_smem_*()，特别是对于 A 和 B 矩阵，以获得更好的 GEMM 和 IO 性能。

  • 如果您需要矩阵采用纯列优先或行优先顺序，请使用 get_layout_smem_*()。

  • GEMM 的最佳布局可能与整个内核的最佳布局不同；请进行实验并尝试各种方法。

  • 对于没有专用 MMA 的用例（如 fp32-fp32-fp32），尝试使用 get_layout_smem_*() 和 suggested_leading_dimension_of 来改进共享内存访问模式。

  • 对于寄存器片段 API，始终选择与您的布局匹配的数据分区器

    • suggest_partitioner() 用于 suggest_layout_smem_*()

    • get_partitioner() 用于 get_layout_smem_*()

示例

使用 最大对齐、cublasdx::copy 和 建议布局 用于共享内存以提高性能的示例。

using BLAS = decltype(Size<128, 128, 128>() + Type<type::real>() + Precision<__half, __half, double>() + MaxAlignment() + Block() + ...);

// Tensors with global memory data
auto a_global_tensor = cublasdx::make_tensor(a, BLAS::get_layout_gmem_a());
auto b_global_tensor = cublasdx::make_tensor(b, BLAS::get_layout_gmem_b());
auto c_global_tensor = cublasdx::make_tensor(c, BLAS::get_layout_gmem_c());

// Tensors with shared memory data
auto a_shared_tensor = cublasdx::make_tensor(smem_a, BLAS::suggest_layout_smem_a());
auto b_shared_tensor = cublasdx::make_tensor(smem_b, BLAS::suggest_layout_smem_b());

// 16-byte (128-bit) alignment helps vectorize (if possible) copying between shared and global memory
using blas_alignment = cublasdx::alignment_of<BLAS>; // 16, 16, 16
cublasdx::copy<BLAS, blas_alignment::a>(a_global_tensor, a_shared_tensor);
cublasdx::copy<BLAS, blas_alignment::b>(b_global_tensor, b_shared_tensor);
cublasdx::copy_wait();

// Get suggested partitioner, matching shared memory layouts for A and B
auto partitioner = BLAS::suggest_partitioner();

// Make a register result accumulator for this GEMM execution
auto c_fragment_accumulator = partitioner.make_accumulator_fragment();

// Partition c_global and copy this thread's elements into register fragment
cublasdx::copy_fragment<blas_alignment::c>(c_global_tensor, c_fragment_accumulator, partitioner);

// 16-byte (128-bit) alignment and suggested layouts help improve shared memory IO in GEMM
BLAS().execute(a_shared_tensor, b_shared_tensor, c_fragment_accumulator)
__syncthreads();

auto out_global_tensor = cublasdx::make_tensor(output, BLAS::get_layout_gmem_c());

// Partition out_global_tensor and copy this thread's elements into appropriate locations
cublasdx::copy_fragment<blas_alignment::c>(c_fragment_accumulator, out_global_tensor, partitioner);

内存管理#

  • 避免不必要地从/向全局内存读取/写入数据。

  • 确保全局内存读取/写入是合并的。

  • 使用 shared 内存或额外的寄存器来存储临时数据。

  • 将任何元素级预处理和后处理卸载到可以传递给 execute(…) 的转换仿函数，以避免访问共享内存。

高级#

进一步阅读#

参考文献#