Is it possible to call a cuda.jit function which indirectly calls another cuda.jit function?

user_8529846 0
Viewed 697

I need to be able to call a GPU function which itself calls another GPU function indirectly:

from numba import cuda, jit
import numpy as np

# GPU function
@cuda.jit(device = True)
def euclidean_distance_gpu(input_vec, weight, diffs):
  i = cuda.grid(1)
  if i < input_vec.shape[0]:
    diffs[i] = (input_vec[i] - weight[i]) ** 2

@jit
# CPU function
def euclidean_distance_cpu(diffs):
  diffs_sum = np.sum(diffs)
  euclidean_distance = np.sqrt(diffs_sum)

  return euclidean_distance

@jit
# CPU function
def euclidean_distance(input_vec, weight):
  euclidean_distance_gpu[1, 5](input_vec, weight, diffs)
  
  return euclidean_distance_cpu(diffs)

@cuda.jit
# GPU function
def compare(input_vec, categories, diffs):
  i = cuda.grid(1)
  if i < categories.shape[0]:
    euclidean_dist = 0
    euclidean_dist = euclidean_distance(input_vec, categories[i])
    diffs[i] = euclidean_dist

vec1 = np.array([1, 2, 3, 4, 5])
c1 = np.array([2, 3, 4, 5, 6])
c2 = np.array([3, 4, 5, 6, 7])
c = np.array([c1, c2])
diffs = np.array([0, 0])
compare(vec1, c, diffs)

In this case, I need to call the compare() function, which itself calls euclidean_distance_gpu() through euclidean_distance(), and compare() and euclidean_distance_gpu() are functions meant to use a GPU.

From what I understand, both functions need to be decorated with @cuda.jit, with euclidean_distance_gpu() being decorated with @cuda.jit(device = True). However, when I call compare() later on, I don't see how this can be called without raising an error, since it first has to go through a CPU function (euclidean_distance() decorated with @jit).

My understanding is that you can only call a cuda.jit function from another cuda.jit function - is this correct? Let's say I made euclidean_distance() into a cuda.jit function as well. Is there a way to make this work correctly, through all those layers of function calls?

I'm quite new to jit - is there anything I can do here? Note these functions are actually more complex than shown so I'd like an actual solution rather than just in-lining the function.

numba jit cuda python
1 Answers

There is a lot of misconceptions in the words and code in this question. Put simply:

  1. Numba kernels cannot launch other Numba kernels (from the documentation)

... newer CUDA devices support device-side kernel launching; this feature is called dynamic parallelism but Numba does not support it currently

  1. Numba kernels can call Numba device functions, and Numba device functions can call other Numba device functions (link). And I am pretty sure that device functions work by lowering and inline expansion, i.e. Numba device functions do not use the CUDA ABI

  2. No CUDA code, Numba or otherwise can run code on the host CPU

If you look at the code in your question, it should be immediately obvious that basically all of it is illegal when viewed through the lens of items 1, 2, and 3 above.

I would suggest an alternative design pattern, but it is so unobvious what the code is trying to actually do that I cannot do so. There are too many strange contradictions in different parts of the code for that.

user_6818650
Related