How To Use Supported Numpy And Math Functions With Cuda In Python?

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):

Solution 1:

The hint to the source of the problem is here:

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

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 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)

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, float64import 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)