避免 Numba CUDA Jit 竞争条件答案

【问题标题】：Avoiding Numba CUDA Jit Race Condition避免 Numba CUDA Jit 竞争条件
【发布时间】：2020-01-16 08:17:41
【问题描述】：

我有一个简单的例子：

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 for 循环移动到内核中，但看起来我正在进入竞争条件。这是我目前拥有的：

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.]

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

【问题讨论】：

从您的第一个示例到第二个示例，大概您正在演示将 5 个 max_i 循环带入内核，您进行了两项更改：1.您将网格增加了 5 倍，2.您在内核中添加了一个循环。从概念上讲，您只需要其中一项更改。但这里更大的问题是，将该循环排除在内核之外意味着您在每次内核调用时都有一个网格范围的同步，我认为这对于您的算法是必要的。内核内部的循环并没有消除对它的需要。原子不会解决这个问题。你需要一个内核网格同步，而 numba cuda 没有这个
是的，你的观察是正确的。我添加了第二个维度，我认为只有在所有 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 同时执行，现在我有相同的竞争条件，但对于赔率和偶数。我有什么选择？
在我看来，明智的选择是您在上一个问题中的代码实现。我没有费心解决这个问题，因为您没有提供测试用例，不，我不会尝试对内核进行逆向工程，以找出如何用足够的代码包装它来制作测试用例。您可能已经拥有它并决定不提供它。所以我决定不关注。没有测试用例 IMO 的性能问题“不清楚且无用”。
如果您真的非常想将网格同步推送到内核中，那么正确的方法在我所知道的任何 cuda python 实现中都没有公开（还）。您可以将内核构建到 CUDA C++ 中的库中，并使用 python ctypes 从 python 调用它。可能还有其他方法。但如果不了解问题的范围和可能的一些分析，我不会做所有这些工作。仅使用内核无法做到这一点。
因此，考虑到在内核外部使用 i for 循环会产生正确的答案，但会导致许多缓慢/昂贵的内核调用，是否有更好的方法从Python i for 循环？尝试将 for 循环推入内核的目的是为了加快速度，我觉得自己离我们如此之近又如此遥远。

标签： python numpy cuda numba

【解决方案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++ does。无论如何，运行由单个块组成的 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 值（使用合适的块和网格大小算术，此处未描述，但通常遵循示例中的内容）。

【讨论】：