如何在 Cupy 内核中使用 WMMA 函数？

Question

如何在cupy.RawKernel或cupy.RawModule中使用wmma::load_matrix_sync等WMMA函数？ 有人可以提供一个最小的例子吗？

Answer 1

我们可以结合cupy RawKernel 和wmma programming 的信息来提供大部分所需的材料。我不打算提供有关 wmma 编程的教程，还有其他资源，例如 this blog 和 cutlass template library。

请注意，wmma 函数需要 7.0 或更高的计算能力。您必须在 Volta、Turing 或 Ampere GPU 上运行。

让我们以编程指南中给出的kernel example 为例。要将其放入RawKernel，我们需要将其作为字符串提供。为了支持内核 C 风格的使用，我将内核代码分解为一个可以使用 C++ 的__device__ 函数，同时使用 C 风格链接导出内核入口点（wmma_ker）。示例代码执行 16x16 矩阵乘法（使用单个扭曲）。这是一个有效的例子：

# cat t24.py
import numpy
import cupy as cp
ddim = 16
bdim = 32
gdim = 1
a = cp.ones(ddim*ddim, dtype=cp.float16)
b = cp.ones(ddim*ddim, dtype=cp.float16)
c = cp.zeros(ddim*ddim, dtype=cp.float32)
wmma_ker = cp.RawKernel(r'''
  #include <mma.h>
  using namespace nvcuda;
  __device__ void wmma_ker_dev(half *a, half *b, float *c) {
  // Declare the fragments
    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::col_major> a_frag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> c_frag;

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

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

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

  // Store the output
    wmma::store_matrix_sync(c, c_frag, 16, wmma::mem_row_major);
  }
  extern "C" {
    __global__ void wmma_ker(half *a, half *b, float *c) {
          wmma_ker_dev(a,b,c);
    }
  }
 ''', 'wmma_ker', options=("-restrict","-lineinfo"))
wmma_ker((gdim,1), (bdim,1), (a,b,c))  # grid, block and arguments
r_o = cp.asnumpy(c)
print(r_o)
# cuda-memcheck python t24.py
========= CUDA-MEMCHECK
[16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16. 16.
 16. 16. 16. 16.]
========= ERROR SUMMARY: 0 errors
#

我使用pip install cupy-cuda102为此设置了cupy，否则在安装了CUDA 10.2和Tesla V100 GPU的机器上运行。我提供的 RawKernel options 对于此演示来说是不必要的，您可以完全省略该参数。

这段代码的目的是演示一个示例方法。我并不是说代码没有缺陷或适用于任何特定目的。需要您自担风险使用它。特别是，如果更改了它的任何方面，我不希望此代码正常工作。我并不是说它是一个通用/灵活/可扩展的矩阵乘法例程。