如何在 Python 中使用 CUDA 支持的 numpy 和数学函数？

Question

根据 numba 0.51.2 documentation，CUDA Python 支持几个 math 函数。但是，它在以下内核函数中不起作用：

@cuda.jit
def find_angle(angles):
    i, j = cuda.grid(2)
    if i < angles.shape[0] and j < angles.shape[1]:
        angles[i][j] = math.atan2(j, i)

输出：

numba.core.errors.LoweringError: Failed in nopython mode pipeline (step: nopython mode backend)
No definition for lowering <built-in function atan2>(int64, int64) -> float64

我是不是用错了函数？

Answer 1

问题根源的提示在这里：

No definition for lowering <built-in function atan2>(int64, int64) -> float64

由 cuda.grid() 编辑的参数 return（即您传递给 atan2 的 i、j）是整数值，因为它们是相关的索引。

numba 找不到它可以使用的 atan2 版本，该版本接受两个整数参数和 returns 一个浮点值：

float64 = atan2(int64, int64)

一个可能的解决方案是转换你的 atan2 输入参数以匹配 numba 似乎想要从该函数 return 的类型，这显然是 float64:

from numba import cuda, float64
import numpy
import math


@cuda.jit
def find_angle(angles):
    i, j = cuda.grid(2)
    if i < angles.shape[0] and j < angles.shape[1]:
        angles[i][j] = math.atan2(float64(j), float64(i))

block_x = 32
block_y = 32
block = (block_x, block_y)
x = 256
y = 256
grid = (x//block_x, y//block_y) # not for arbitrary x and y

angles = numpy.ones((x, y), numpy.float64)
find_angle[grid, block](angles)

如何在 Python 中使用 CUDA 支持的 numpy 和数学函数？

How to use supported numpy and math functions with CUDA in Python?

python

cuda

numpy

numba