如何编写一个cupy用户定义的核函数来计算分段和

Posted 2023-02-16

【中文标题】如何编写一个cupy用户定义的核函数来计算分段和

【英文标题】：How to write a cupy user-defined kernel function to calculate the segmented sum

【发布时间】：2020-11-30 15:02:59

【问题描述】：

我现在用下面这个函数，但是觉得不行，但是看不懂cupy内核定义的描述。当涉及大量数据时，此功能非常占用内存且非常耗时。

def cupy_sum(self, bins):
    bidx = cupy.cumsum(bins) -1,
    return cupy.diff(cupy.r_[0, cupy.cumsum(self)[bidx]])

参考其他例子，写如下代码，不知道有没有问题。

sum_section_kernel = cp.ElementwiseKernel(
'raw T bins, raw T dats',
'float32 out',
'''
T bin_f = bins[i  ];
T bin_l = bins[i+1];
T biv   = 0;
for(size_t j=bin_f; j<bin_l; j++)
  biv += dats[j];
        
out = biv; 
''',
'summe')

a = cp.array([4, 3, 5], dtype=cp.float32)
b = cp.array([1, 1, 1.1, 1,   2, 2, 2,   3, 3, 3, 3, 3], dtype=cp.float32)
y = cp.empty(3, dtype=cp.float32)

a = cp.r_[0,a.cumsum()]
out = sum_section_kernel(a, b, y)
print(out)

> [ 4.100  6.000  15.000]

【参考方案1】：

上面已经放了例子了，速度没有提升，不过我觉得还是有节省内存的好处的。