避免Numba CUDA Jit竞争条件

Question

我有一个简单的例子：

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(i, y, z):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)

    for j in range(start, y.shape[0], stride):
        # Note that these aren't my real functions but they demo the point
        if i < j:
            y[j, 0] = i
            z[j, 0] = i + j
        if i == j:
            y[j, 1] = i
            z[j, 1] = i * j
        if i > j:
            y[j, 2] = i
            z[j, 2] = j


if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = 10
    blocks_per_grid = math.ceil(y.shape[0]/threads_per_block[1])

    for i in range(max_i):
        func[blocks_per_grid, threads_per_block](i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

输出应该是这样的：

[[1. 0. 4.]
 [0. 1. 4.]
 [1. 2. 4.]
 [2. 3. 4.]
 [3. 4. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]]

然而，当max_i很大时，大部分时间都花费在调用CUDA内核上，我希望尽可能使这个内核更快。因此，我尝试着将max_i的循环移动到内核中，但似乎会遇到竞争条件。这是我目前拥有的：

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(max_i, y, z):
    a, b = cuda.grid(2)
    a_stride, b_stride = cuda.gridsize(2)

    for i in range(a, max_i, a_stride):
        for j in range(b, y.shape[0], b_stride):
            if i < j:
                y[j, 0] = i
                z[j, 0] = i + j
            if i == j:
                y[j, 1] = i
                z[j, 1] = i * j
            if i > j:
                y[j, 2] = i
                z[j, 2] = j

if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = (1, 10)
    blocks_per_grid = (max_i, math.ceil(y.shape[0]/threads_per_block[1]))

    func[blocks_per_grid, threads_per_block](max_i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

这个（不正确的）输出看起来像：

[[1. 0. 4.]
 [0. 1. 4.]
 [1. 2. 4.]
 [1. 3. 4.]  # Should be [2. 3. 4.]
 [3. 4. 1.]
 [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [3. 1. 1.]  # Should be [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]  # Should be [4. 1. 1.]
 [0. 1. 1.]]  # Should be [4. 1. 1.]

如上所述，如何使用单个内核获得正确的答案，同时使这个内核尽可能快（即避免原子操作）？

- slaw

---

从您的第一个示例到第二个示例，您似乎演示了将5个max_i循环引入内核，您进行了两个更改：1.将网格增加了5倍，2.在内核中添加了一个循环。从概念上讲，您只需要其中一个更改。但更大的问题是，将该循环保持在内核之外意味着每次内核调用时都需要进行整个网格的同步，我认为这对于您的算法是必要的。内核中的循环并不能消除这种需求。原子操作也无法解决这个问题。您需要一个内核内的网格同步，而numba cuda没有这个功能。 - Robert Crovella

是的，你的观察是正确的。我添加了第二个维度，认为我可以强制执行，比如说，只有在所有的i=1开始执行（或同时执行）后，才能执行i=2。然后，我会将i=1的结果（以及所有的i=odd）写入一个单独的y_odd和z_odd中，而i=2（以及所有的i=even）将写入一个y_even和z_even中。但是我意识到，如果总线程数大于2 * x.shape[0]，那么这将是一个问题，因为对于i=3来说，有足够的线程与i=1同时执行，现在我有了相同的竞争条件，但是是针对奇数和偶数。我的选择是什么？ - slaw

我认为明智的选择是你在之前问题中提出的代码实现。由于你没有提供测试用例，所以我没有费心去解决那个问题。不，我不会试图反向工程内核来弄清楚如何用足够的代码包装它以使其成为一个测试用例。你很可能已经有了这个，并决定不提供它。所以我决定不去关注这个问题。在我看来，没有测试用例的性能问题是“不清楚且无用的”。 - Robert Crovella

如果你真的非常想将网格同步推入内核，我所知道的任何CUDA Python实现中都没有暴露正确的方法（尚未）。你可以在CUDA C++中将你的内核构建成库，并使用Python ctypes从Python中调用它。可能还有其他方法。但是，在不了解问题范围和可能的一些分析之前，我不会在没有理解的情况下进行所有这些工作。仅仅一个内核是无法做到这一点的。 - Robert Crovella

那么，考虑到将i循环放在内核之外可以产生正确的答案，但会导致许多缓慢/昂贵的内核调用，有没有更好的方法从Python i循环中多次调用内核？尝试将for循环推入内核的目的是为了加速处理速度，我感觉离成功已经很接近了，但又似乎遥不可及。 - slaw

关于我之前的问题，您是完全正确的，我只能请求您的耐心和谅解。这个CUDA东西对我来说是全新的，所以我正在我的有限理解的边缘操作。因此，甚至很难确定我应该问什么。希望这个问题更好一些。我发现内核的内容并不重要（我能够注释掉内核的主体，并且使用大的max_i得到相同的时间结果），因为98-99％的时间实际上是在调用内核。感谢您抽出时间帮助我解决这个问题！ - slaw

Answer 1

如上所述，如何在使用单个内核的情况下获得正确答案并使该内核尽可能快（即避免原子操作）？

对于您在此处展示的数据大小，一个非常简单的方法就是将所有操作放在一个CUDA线程块中，在循环结束时进行块同步，并直接将max_i循环放入内核中：

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(max_i, y, z):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)
    for i in range(max_i):
        for j in range(start, y.shape[0], stride):
            if i < j:
                y[j, 0] = i
                z[j, 0] = i + j
            if i == j:
                y[j, 1] = i
                z[j, 1] = i * j
            if i > j:
                y[j, 2] = i
                z[j, 2] = j
        cuda.syncthreads()

if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = 1024
    blocks_per_grid = 1

    func[blocks_per_grid, threads_per_block](max_i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

这将适用于最多1024个值的n。然而在一般情况下，如果n大于1024，我们需要另一种方法。要扩展先前的方法，我们需要进行网格范围的同步（当我们超过一个单一块时），但我不知道有任何CUDA Python实现可以提供这个功能，尽管CUDA C ++会有。

通常，仅运行由单个块组成的CUDA代码并不是利用GPU性能的好方法。

相反，对于你所展示的函数，我们可以观察到只有一个值最终会出现在输出数组的每个位置，即使每个位置可能被写入多次。因此，我们的挑战在于确定正确的输出值（即在给定max_i循环数后将被写入该位置的最终值），并仅通过每个位置进行一次遍历。以下是生成仅y输出的示例：

from numba import cuda
import numpy as np
import math

@cuda.jit
def func(max_i, y, z):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)
    for j in range(start, y.shape[0], stride):
        if j > 0:
            y[j, 0] = min(max_i-1, j-1)
#            z[j, 0] = i+j
        if j < max_i:
            y[j, 1] = j
#            z[j, 1] = i * j
        if j < max_i-1:
            y[j, 2] = max_i-1
#            z[j, 2] = j

if __name__ == '__main__':
    n = 30
    y = np.ones((n, 3))
    z = np.ones((n, 3)) * -1
    device_y = cuda.to_device(y)
    device_z = cuda.to_device(z)
    max_i = 5
    threads_per_block = 1024
    blocks_per_grid = 1

    func[blocks_per_grid, threads_per_block](max_i, device_y, device_z)

    out = device_y.copy_to_host()
    print(out)

使用类似的方法可以生成z值，这种方法适用于大于1024的n值（使用适当的块和网格大小算法，未在此处描述，但通常遵循您示例中的方法）。

- Robert Crovella