在 cupy 中将结构传递给原始内核

Question

我的 CUDA 内核采用 float3、int2 等结构作为参数。我似乎无法通过 cupy rawkernel 接口正确地将参数传递给这些内核。我曾尝试为 float3 参数传递一个包含 3 个浮点数的 1d cupy 数组，但该参数在内核中未被正确解释。我尝试传递一个 ctypes 结构，但返回了一个不受支持的类型错误。是否可以将自定义结构发送到 cupy 中的原始内核？如果可以，怎么做？

我尝试使用 ctype 结构如下：

class float3(ctypes.Structure): 
    fields = [ ("X", c_float), ("Y", c_float), ("Z", c_float)] 

from cupy.cuda.function import CPointer 

class CFloat3(CPointer): 
    def __init__(self, v): super().__init__(ctypes.addressof(v)) 
        self.val = v 

val= float3(1.5, 3, 5) 
cval= CFloat3(val)

这绕过了 cupy 的类型检查，但仍然没有正确地将值传递给内核。如果您查看 cupy 源代码中的功能模块，它似乎应该可以工作。它只是传递结构的指针。我还尝试了 id(v) 和 ctypes.POINTER(float3)(v) 而不是 ctypes.addressof 来获取结构的地址，但这也不起作用。

我可以通过编写接受数组作为输入的内核包装器来解决这个问题，然后将数组转换为结构以调用我的常规内核。虽然这对我来说很难看。如果不能做到这一点，那么不提供将结构传递给内核的能力似乎是一个很大的疏忽。

Answer 1

我同意这个评论；在一般情况下，我无法找到一种方法来完成这项工作。

通过重新调整 np.complexXX 数据类型的用途，或许可以使用一种 hacky 方法来处理 float2 或 double2。这是一个例子：

$ cat t19.py
import numpy as np
import cupy
ddim = 64
bdim = 32
d = np.complex64(1+2j)
i = cupy.ones((ddim*3), dtype=cupy.float32).reshape(ddim, 3)
o = cupy.zeros((ddim*3), dtype = cupy.float32).reshape(ddim, 3)
my_test = cupy.RawKernel(r'''
  extern "C" __global__
  void my_test(const float2 d, const  float3 * __restrict__  i, float3 * __restrict__ o, int dim) {
  int x = blockDim.x * blockIdx.x + threadIdx.x;
  if (x < dim){
    float3 temp = i[x];
    temp.x += d.x;
    temp.y += d.y;
    temp.z += d.x;
    o[x] = temp;}
  }
 ''', 'my_test')
gdim = ddim//bdim + 1
my_test((gdim,1), (bdim,1), (d, i,o,ddim))  # grid, block and arguments
r_o = cupy.asnumpy(o)
print(r_o)
$ python t19.py
[[2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]
 [2. 3. 2.]]
$

我在 numpy 结构化数据类型方面运气不佳，这似乎是合乎逻辑的路径。

Answer 2

感谢您的提问。

解决 float2、float3 类型问题的一种（hackish）方法是在内核中转换 cupy 数组指针（但是，这并不真正推荐）：

import cupy
add_kernel = cupy.RawKernel(r'''
    extern "C" __global__
    void my_add(const float* x1, float* y) {
        int tid = blockDim.x * blockIdx.x + threadIdx.x;
        float3* xf3 = (float3*) x1;
        y[tid] = xf3->x + xf3->y + xf3->z;
    }
    ''', 'my_add')
x1 = cupy.array([1, 2, 3], dtype='float32')
y = cupy.array([0], dtype='float32')
add_kernel((1,), (1,), (x1, y))

但是，CuPy 不支持结构化数组，因此无法将 cupy 数组映射到 cuda 内核中的用户定义结构。

在 cupy 中将结构传递给原始内核

Passing structure to raw kernel in cupy

python

cuda

chainer

cupy