CuPy 数据上的“无法确定 Numba 类型”错误

Question

我正在尝试使用 Numba 代码处理一些由 CuPy 逻辑生成的数据。但得到不同的错误。以简化示例为例

import cupy as cp
import numba
from numba import void, int32, int64, float32, float64
import numpy as np

@numba.jit
def numba_test(a, b, n, m):
    for i in range(n):
        for j in range(m):
            a[i, j] += b[i, j]

a = cp.zeros((100, 10), dtype=np.float32)
b = cp.zeros((100, 10), dtype=np.float32)
numba_test(a, b, 100, 10)

我来了

<ipython-input-1-d30521cb61ff>:6: NumbaWarning: 
Compilation is falling back to object mode WITH looplifting enabled because Function "numba_test" failed type inference due to: non-precise type pyobject
[1] During: typing of argument at <ipython-input-1-d30521cb61ff> (8)

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  @numba.jit
<ipython-input-1-d30521cb61ff>:6: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "numba_test" failed type inference due to: cannot determine Numba type of <class 'numba.dispatcher.LiftedLoop'>

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  @numba.jit
/usr/local/lib/python3.6/dist-packages/numba/object_mode_passes.py:178: NumbaWarning: Function "numba_test" was compiled in object mode without forceobj=True, but has lifted loops.

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  state.func_ir.loc))
/usr/local/lib/python3.6/dist-packages/numba/object_mode_passes.py:188: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  state.func_ir.loc))
<ipython-input-1-d30521cb61ff>:6: NumbaWarning: 
Compilation is falling back to object mode WITHOUT looplifting enabled because Function "numba_test" failed type inference due to: non-precise type pyobject
[1] During: typing of argument at <ipython-input-1-d30521cb61ff> (8)

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  @numba.jit
/usr/local/lib/python3.6/dist-packages/numba/object_mode_passes.py:178: NumbaWarning: Function "numba_test" was compiled in object mode without forceobj=True.

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  state.func_ir.loc))
/usr/local/lib/python3.6/dist-packages/numba/object_mode_passes.py:188: NumbaDeprecationWarning: 
Fall-back from the nopython compilation path to the object mode compilation path has been detected, this is deprecated behaviour.

For more information visit http://numba.pydata.org/numba-doc/latest/reference/deprecation.html#deprecation-of-object-mode-fall-back-behaviour-when-using-jit

File "<ipython-input-1-d30521cb61ff>", line 8:
def numba_test(a, b, n, m):
    for i in range(n):
    ^

  state.func_ir.loc))

我在 Google Colab 上运行代码，之前能够运行一些 CuPy 代码和相当多的 Numba 代码。

在@numba.jit 之后添加 (void(float64[:, :], float64[:, :], int64, int64)) 没有帮助，会产生类似于无效参数类型 pyObject 的错误。用 cp.asarray(a) 替换 a 和 b 也无济于事。如何运行？

开发者写他们提供必要的接口https://docs.cupy.dev/en/stable/reference/interoperability.html

@cuda.jit
def add(x, out):
        start = cuda.grid(1)
        stride = cuda.gridsize(1)
        for i in range(start, x.shape[0], stride):
                out[i] = x[i] + 2

a = cupy.arange(10)
out = cupy.zeros_like(a)

add[1, 32](a, out)

作品

Answer 1

您正在调用 CPU JIT 编译器。您应该按照您链接的说明页面进行操作

from numba import cuda

@cuda.jit
def ...

使用 CUDA JIT 编译器。