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

Question

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

Answer 1

我们可以结合 cupy

RawKernel

和 wmma 编程 的信息来提供大部分所需的材料。我不打算提供有关 wmma 编程的教程，还有其他资源，例如此博客和cutlass 模板库。

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

让我们以编程指南中给出的内核示例为例。要将其放入

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

对于本演示来说是不必要的，您可以完全省略该参数。

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