标签: cuda-wmma

张量核心可以通过 CUDA 中的 WMMA 接口以编程方式访问（请参阅https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#wmma和https://developer.nvidia.com/博客/programming-tensor-cores-cuda-9/）。最近，在 Ampere 一代卡中，Nvidia 宣布能够使用稀疏矩阵执行稀疏张量运算，如下所示：https ://developer.nvidia.com/blog/acceleating-inference-with-sparsity-using-ampere-和-tensorrt/

所呈现的格式似乎采用元素对及其在四个元素段（2 位索引）内的顺序。然而，查看wmma 文档，我找不到任何提及这一点的内容，也找不到如何访问这些特殊的张量核心操作。AFAICT 的该功能的公告页面也没有说明这一点。

如何访问 cuda 中的稀疏张量核心功能？

我想知道 NumBlocks 和 ThreadsPerBlock 对这个简单的矩阵乘法例程的影响是什么

__global__ void wmma_matrix_mult(half *a, half *b, half *out) {

   // Declare the fragments
   wmma::fragment<wmma::matrix_a, M, N, K, half, wmma::row_major> a_frag;
   wmma::fragment<wmma::matrix_b, M, N, K, half, wmma::row_major> b_frag;
   wmma::fragment<wmma::accumulator, M, N, K, half> c_frag;

   // Initialize the output to zero
   wmma::fill_fragment(c_frag, 0.0f);

   // Load the inputs
   wmma::load_matrix_sync(a_frag, a, N);
   wmma::load_matrix_sync(b_frag, b, N);

   // Perform the matrix multiplication
   wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

   // Store the output
   wmma::store_matrix_sync(out, c_frag, N, wmma::mem_row_major);
}

呼唤

`wmma_matrix_mult<<1, 1>>`: Incorrect
`wmma_matrix_mult<<1, 2>>`: …