在cupy中使用元素内核对条目求和的问题

Posted 2023-02-16

【中文标题】在cupy中使用元素内核对条目求和的问题

【英文标题】：Problem with summing entries using the elementwise kernel in cupy

【发布时间】：2019-09-27 13:51:35

【问题描述】：

在第一个代码示例 (kernel_conv) 中，我编写了一个简单的卷积，它与预期的结果 [1,1,2,1,1] 一起工作。

然后我使用元素内核来对向量的所有条目求和。但是，如果我运行第二个示例 (kernel_sum)，我会得到结果 [3,0,0]，但会期望 [6,0,0]。

这两个例子有什么区别？为什么第一个示例中的变量 y 更新了，而第二个示例中的变量似乎被覆盖了？

import numpy as np 
import cupy as cp 

kernel_conv = cp.ElementwiseKernel(
    'raw float32 x', 'raw float32 y',
    ''' int idx = i*2 + 1;
        for(size_t j=0;j<3;j++)
          y[idx - 1 + j] += x[j];
        
    ''', 'conv')

x = cp.asarray(np.array([1,1,1]),dtype=np.float32)
y = cp.zeros((5,),dtype=np.float32)
z = kernel_conv(x,y,size=2)
print(z)

kernel_sum = cp.ElementwiseKernel(
  'raw float32 x', 'raw float32 y',
  ''' 
      y[0] += x[i]
  ''', 'summe')

x = cp.asarray(np.array([1, 2, 3]), dtype=np.float32)
y = cp.zeros((3,),dtype=np.float32)
z = kernel_sum(x,y,size=3)
print(z)

【参考方案1】：

kernel_sum 的错误结果是由于数据争用造成的。在这种情况下，3 个线程尝试同时写入全局内存的同一地址（y[0]）。为避免数据竞争，您应该 1) 使用 atomicAdd 或 2) 使用 cupy.ReductionKernel 进行缩减。

实际上，kernel_conv 也有数据竞争。运行y[2] += x[2] 的第一个线程可能与运行y[2] += x[0] 的第二个线程冲突。由于第一个在实际执行中略有滞后，所以结果不受影响，但这是一个时序问题，一般不能保证*。要更正它，您可以在此处再次使用 atomicAdd，或者您也可以更改通过多个线程运行计算的方式（例如，启动 5 个线程，每个线程计算 y 的不同元素）。

* 确实，在上面kernel_conv 的例子中，我猜当所有线程都运行在同一个warp 中时，正确性是有保证的，即线程数不大于32。这是因为所有线程在同一个warp中同步运行，直到控制流发散。如果线程数设置为较大的值，则可能会在 warp 边界处发生数据竞争。