Cuda Tensor Cores：NumBlocks 和 ThreadsPerBlock 的作用是什么？

Question

我想知道 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>>`: Incorrect
`wmma_matrix_mult<<1, 4>>`: Incorrect
`wmma_matrix_mult<<1, 8>>`: Incorrect
`wmma_matrix_mult<<1, 16>>`: Incorrect
`wmma_matrix_mult<<1, 32>>`: Correct

如果每个线程都执行相同的执行，为什么每个块的线程数很重要？正如你所看到的，我没有对threadIdx.x内核内部做任何事情。

Answer 1

张量核心操作发生在扭曲级别。wmma 中的 w表示这一点。参考文档：

这需要扭曲中所有线程的合作。

每个张量核心单元可以wmma::mma_sync在每个时钟周期从扭曲接受一次矩阵乘法运算（即 ）。

这意味着完整的扭曲（32 个线程）必须可用并参与，操作才有意义（即合法）。所有wmma::操作都是集体操作，这意味着整个扭曲预计将执行它们，并且对于正确使用是必要的。

如果您有多个扭曲参与（例如，线程块大小为 64 或 128 等），则您实际上要求完成多个操作，就像任何其他 CUDA 代码一样。

与任何其他 CUDA 代码一样，启动多个块的操作只是扩展正在完成的工作的一种方法，如果您想利用具有多个 SM 的 GPU 的资源，当然这是必要的。由于张量核心单元是每个 SM 的资源，因此有必要见证 CUDA GPU 为张量核心运算提供接近其满额定吞吐量的任何内容。

如果每个线程都执行相同的执行，为什么每个块的线程数很重要？

每个线程都没有做同样的事情。集体wmma::操作将代码隐藏在引擎盖下，根据线程所属的扭曲通道来专门化线程行为。例如，与任何其他扭曲通道中的任何线程相比，扭曲通道 0 中的线程将选择片段的不同元素进行关联（即加载、存储）。