如何让cuda在3d网格中正常运行？

Question

我将一个小的cuda函数编码如下：

e=8
@cuda.jit
def my_kernel(e):
    i,j = cuda.grid(2)

    if i <= e and j<=e:
        i**2+j**2

# Host code   

threadsperblock = 16,16
blockspergrid = 32,32
my_kernel[blockspergrid, threadsperblock](e)
print('ok')

现在，如果我把它变成3d形式，它会给我一个错误：

e=8
@cuda.jit
def my_kernel(e):
    i,j,k = cuda.grid(3)

    if i <= e and j<=e and k<=e:
        i**2+j**2+k**2

# Host code   

threadsperblock = 16,16,16
blockspergrid = 32,32,32
my_kernel[blockspergrid, threadsperblock](e)
print('ok')

error:
[1] Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE

如何解决这个问题？ 非常感谢。