为什么我的 CUDA 内核 (Numba) 在使用相同输入的连续调用中表现不同？

Question

我在 python 的 numba + cuda 中有一个菜鸟错误。 Numba 版本是 0.51，CUDA 版本是 10.2。当使用完全相同的输入重复调用时，下面的代码会给出非常不同的输出。

import numpy as np
from numba import cuda, jit

@cuda.jit()
def writeToArray(vec, array_in, array_out):
    ''' vec is a 3x1 vector, array_in is a 3D array, array_out is a 3D array of the shape of array in'''
    i,j,k = cuda.grid(3)
    value = array_in[i,j,k] * vec[0] + array_in[i,j,k] * vec[0] + array_in[i,j,k] * vec[0]
    cuda.atomic.max(array_out,(i,j,k), value)
    # cuda.synchronize()

def test():
    
    threadsperblock = (8,8,8)
    blockspergrid_x = ( 17 + threadsperblock[0]) // threadsperblock[0]
    blockspergrid_y = ( 21 + threadsperblock[1]) // threadsperblock[1]
    blockspergrid_z = ( 5 + threadsperblock[2]) // threadsperblock[2]
    blockspergrid = (blockspergrid_x, blockspergrid_y, blockspergrid_z)
    array_in = np.random.rand(17,21,5).astype(np.float_)
    vec = np.array([1.0, -1.0, 1.0]).astype(np.float_)
    d_array_in = cuda.to_device(array_in)
    d_vec = cuda.to_device(vec)
    while True:
        array_out_1 = -999.999*np.ones_like(array_in)
        array_out_2 = -999.999*np.ones_like(array_in)
        d_array_out_1= cuda.to_device(array_out_2)
        d_array_out_2 = cuda.to_device(array_out_2)
        writeToArray[blockspergrid, threadsperblock](d_vec, d_array_in, d_array_out_1)
        writeToArray[blockspergrid, threadsperblock](d_vec, d_array_in, d_array_out_2)
        array_out_1_host = d_array_out_1.copy_to_host()
        array_out_2_host = d_array_out_2.copy_to_host()
        assert(np.allclose(array_out_1_host, array_out_2_host))

if __name__ == "__main__":
    test()

这应该不会中断，但最终，在 while 循环迭代大约 10 次之后，断言失败了。我做错了什么？

Answer 1

您的内核代码正在进行非法的 out-of-bounds 访问。当您以这种方式调整网格大小时：

blockspergrid_x = ( 17 + threadsperblock[0]) // threadsperblock[0]
blockspergrid_y = ( 21 + threadsperblock[1]) // threadsperblock[1]
blockspergrid_z = ( 5 + threadsperblock[2]) // threadsperblock[2]

您运行创建“额外”线程的可能性。这些线程的 i、j、k 索引在输入数组的“形状”之外。您不希望这些线程起作用。通常的做法是在你的内核代码中加入一个“线程检查”：

$ cat t31.py
import numpy as np
from numba import cuda, jit

@cuda.jit()
def writeToArray(vec, array_in, array_out):
    ''' vec is a 3x1 vector, array_in is a 3D array, array_out is a 3D array of the shape of array in'''
    i,j,k = cuda.grid(3)
    if i < array_in.shape[0] and j < array_in.shape[1] and k < array_in.shape[2]:
        value = array_in[i,j,k] * vec[0] + array_in[i,j,k] * vec[0] + array_in[i,j,k] * vec[0]
        cuda.atomic.max(array_out,(i,j,k), value)
    # cuda.synchronize()

def test():

    threadsperblock = (8,8,8)
    blockspergrid_x = ( 17 + threadsperblock[0] -1) // threadsperblock[0]
    blockspergrid_y = ( 21 + threadsperblock[1] -1) // threadsperblock[1]
    blockspergrid_z = (  5 + threadsperblock[2] -1) // threadsperblock[2]
    blockspergrid = (blockspergrid_x, blockspergrid_y, blockspergrid_z)
    array_in = np.random.rand(17,21,5).astype(np.float_)
    vec = np.array([1.0, -1.0, 1.0]).astype(np.float_)
    d_array_in = cuda.to_device(array_in)
    d_vec = cuda.to_device(vec)
    i=0
    while i<20:
        array_out_1 = -999.999*np.ones_like(array_in)
        array_out_2 = -999.999*np.ones_like(array_in)
        d_array_out_1= cuda.to_device(array_out_2)
        d_array_out_2 = cuda.to_device(array_out_2)
        writeToArray[blockspergrid, threadsperblock](d_vec, d_array_in, d_array_out_1)
        writeToArray[blockspergrid, threadsperblock](d_vec, d_array_in, d_array_out_2)
        array_out_1_host = d_array_out_1.copy_to_host()
        array_out_2_host = d_array_out_2.copy_to_host()
        assert(np.allclose(array_out_1_host, array_out_2_host))
        i+=1
        print(i)

if __name__ == "__main__":
    test()
$ cuda-memcheck python t31.py
========= CUDA-MEMCHECK
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
========= ERROR SUMMARY: 0 errors
$

注意上面cuda-memcheck的用法。您可以以与原始代码类似的方式使用它来查看 out-of-bounds 错误消息。像这样的越界访问将导致 CUDA 上下文失败，这意味着每次内核启动（以及所有其他 CUDA activity）在第一次失败后，什么都不做。

我并不是建议在所有可能的情况下测试精确的浮点相等性，但它似乎适用于这种特殊情况。您可能希望阅读 this.

为什么我的 CUDA 内核 (Numba) 在使用相同输入的连续调用中表现不同？

Why is my CUDA kernel (Numba) behaving differently on successive calls with same inputs?

cuda

gpu

race-condition

python-3.x

numba