如何在pyCUDA内核中生成随机数?

How to generate random number inside pyCUDA kernel?

我正在使用 pyCUDA 进行 CUDA 编程。我需要在内核函数中使用随机数。 CURAND 库在其中不起作用 (pyCUDA)。因为GPU的工作量很大,在CPU里面生成随机数再传给GPU是行不通的,反而瓦解了使用GPU的动机。

补充问题:

  1. 有没有办法使用 1 个块和 1 个线程在 GPU 上分配内存。
  2. 我使用了不止一个内核。我需要使用多个 SourceModule 块吗?

无论您在问题中断言什么,PyCUDA 都对 CUrand 提供了相当全面的支持。 GPUArray 模块有一个直接接口,可以使用主机端 API 填充设备内存(请注意,在这种情况下,GPU 上的随机生成器 运行)。

也完全可以在 PyCUDA 内核代码中使用来自 CUrand 的设备端 API。在这个用例中,最棘手的部分是为线程生成器状态分配内存。共有三种选择——在代码中静态分配、动态使用主机内存端分配和动态使用设备端内存分配。以下(经过非常轻微测试的)示例说明了后者,正如您在问题中询问的那样:

import numpy as np
import pycuda.autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray

code = """
    #include <curand_kernel.h>

    const int nstates = %(NGENERATORS)s;
    __device__ curandState_t* states[nstates];

    __global__ void initkernel(int seed)
    {
        int tidx = threadIdx.x + blockIdx.x * blockDim.x;

        if (tidx < nstates) {
            curandState_t* s = new curandState_t;
            if (s != 0) {
                curand_init(seed, tidx, 0, s);
            }

            states[tidx] = s;
        }
    }

    __global__ void randfillkernel(float *values, int N)
    {
        int tidx = threadIdx.x + blockIdx.x * blockDim.x;

        if (tidx < nstates) {
            curandState_t s = *states[tidx];
            for(int i=tidx; i < N; i += blockDim.x * gridDim.x) {
                values[i] = curand_uniform(&s);
            }
            *states[tidx] = s;
        }
    }
"""

N = 1024
mod = SourceModule(code % { "NGENERATORS" : N }, no_extern_c=True, arch="sm_52")
init_func = mod.get_function("_Z10initkerneli")
fill_func = mod.get_function("_Z14randfillkernelPfi")

seed = np.int32(123456789)
nvalues = 10 * N
init_func(seed, block=(N,1,1), grid=(1,1,1))
gdata = gpuarray.zeros(nvalues, dtype=np.float32)
fill_func(gdata, np.int32(nvalues), block=(N,1,1), grid=(1,1,1))

这里有一个初始化内核,它需要 运行 一次来为生成器状态分配内存并用种子初始化它们,然后是一个使用这些状态的内核。如果你想 运行 很多线程,你需要注意 malloc 堆大小限制,但这些可以通过 PyCUDA 驱动程序 API 接口进行操作。

接受的答案有一个问题。 我们在那里有一个名称修改,有点讨厌(这些 _Z10initkerneli_Z14randfillkernelPfi)。 为避免这种情况,我们可以手动将代码包装在 extern "C" {...} 子句中。

code = """
    #include <curand_kernel.h>

    const int nstates = %(NGENERATORS)s;
    __device__ curandState_t* states[nstates];
    extern "C" {

    __global__ void initkernel(int seed)
    { .... }

    __global__ void randfillkernel(float *values, int N)
    { .... }
    }
"""

那么代码还是编译成no_extern_c=True:

mod = SourceModule(code % { "NGENERATORS" : N }, no_extern_c=True)

这应该适用于

init_func = mod.get_function("initkernel")
fill_func = mod.get_function("randfillkernel")

希望对您有所帮助。