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

Question

如何在cupy.RawKernel或cupy.RawModule中使用WMMA:：load_matrix_sync等WMMA函数？ 有人能举个简单的例子吗

Answer 1

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

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

让我们看一下编程指南中给出的kernel example。要将其放在RawKernel中，我们需要将其作为字符串提供。为了支持内核C样式的使用，我将内核代码分解为可以使用C++的^ {< CD3> }函数，同时使用C样式链接导出内核入口点（^ {CD4}}）。示例代码执行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，否则在安装了CUDA10.2和特斯拉V100GPU的机器上运行。我提供的RawKernel{}对于这个演示是不必要的，您可以完全忽略这个参数

这段代码的目的是演示一个示例方法。我并不是说代码是无缺陷的或适合任何特定用途。使用它的风险自负。特别是，如果这段代码的任何方面发生了变化，我都不希望它能正常工作。我并不是说它是一个通用/灵活/可扩展的矩阵乘法例程