如何在 Python 中将支持的 numpy 和数学函数与 CUDA 一起使用？答案

【问题标题】：How to use supported numpy and math functions with CUDA in Python?如何在 Python 中将支持的 numpy 和数学函数与 CUDA 一起使用？
【发布时间】：2021-02-20 16:54:56
【问题描述】：

根据numba 0.51.2documentation，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

我是不是用错了函数？

【问题讨论】：

好吧，arctan2 不在您链接的页面上支持的功能列表中。当我查看该页面时，您对支持的文字描述也不是我发现的。我没有看到支持“大多数”math 功能，我看到提供了支持的功能列表。如果您在 SO 上查看numba 标签，您会发现math 用法的示例，例如here。
@RobertCrovella 我在调试math 和numpy 错误的过程中写了这个问题，并让情况有点混乱。为了澄清这个问题，我将问题指向支持的math 函数，特别是math.atan(y, x)。该链接提供了支持的功能。

标签： python numpy cuda numba

【解决方案1】：

问题根源的提示在这里：

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

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

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

float64 = atan2(int64, int64)

一种可能的解决方案是将您的 atan2 输入参数转换为匹配 numba 似乎想要从该函数返回的类型，这显然是 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)

【讨论】：