如何在Python中将受支持的numpy和数学函数与CUDA一起使用? [英] How to use supported numpy and math functions with CUDA in Python?

问题描述

根据 numba 0.51.2 文档，CUDA Python支持几个 math 函数.但是，它在以下内核函数中不起作用:

According to numba 0.51.2 documentation, CUDA Python supports several math functions. However, it doesn't work in the following kernel function:

@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

我使用该功能不正确吗?

Am I using the function incorrectly?

推荐答案

问题根源的提示在这里:

The hint to the source of the problem is here:

没有用于降低<内置函数atan2>(int64，int64)->的定义.float64

由 cuda.grid()返回的参数(即要传递给 atan2 i ， j >)是整数值，因为它们与索引相关.

The arguments returned by cuda.grid() (i.e. i, j which you are passing to atan2) are integer values because they are related to indexing.

numba找不到可以使用的带有两个整数参数并返回浮点值的 atan2 版本:

numba can't find a version of atan2 that it can use that takes two integer arguments and returns a floating-point value:

float64 = atan2(int64, int64)

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

One possible solution is to convert your atan2 input arguments to match the type that numba seems to want to return from that function, which is evidently 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中将受支持的numpy和数学函数与CUDA一起使用?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！