是否可以调用一个间接调用另一个cuda.jit函数的cuda.jit函数？

Question

我需要能够调用一个GPU函数，该函数本身间接地调用另一个GPU函数：

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)

在这种情况下，我需要调用 compare() 函数，该函数本身通过 euclidean_distance() 和 euclidean_distance_gpu() 调用 euclidean_distance_gpu()，而且 compare() 和 euclidean_distance_gpu() 是要使用 GPU 的函数。

据我所知，这两个函数都需要使用 @cuda.jit 进行装饰，其中 euclidean_distance_gpu() 需要使用 @cuda.jit(device = True) 进行装饰。但是，当我稍后调用 compare() 时，我不知道如何在不出错的情况下进行调用，因为它首先必须经过一个 CPU 函数 (euclidean_distance() 被 @jit 装饰)。

我的理解是，您只能从另一个 cuda.jit 函数中调用 cuda.jit 函数 - 这是正确的吗？假设我也将 euclidean_distance() 转换为 cuda.jit 函数。是否有一种方法可以通过所有这些层级的函数调用来正确运行此操作？

我对 jit 还相当新 - 在这里我能做些什么吗？请注意，这些函数实际上比所示的更复杂，因此我希望得到一个实际的解决方案，而不仅仅是将函数嵌入行中。

- Matthew So

---

如果您在代码中包含（或想要使用）装饰器，那么这个问题会变得更简单、更容易理解。 - talonmies

Answer 1

这个问题中的词语和代码存在很多误解，简单来说：

1. Numba内核无法启动其他Numba内核（来自文档）。      ... 较新的CUDA设备支持设备端内核启动；该功能称为动态并行性，但目前Numba不支持它。

2.Numba内核可以调用Numba设备函数，Numba设备函数也可以调用其他Numba设备函数 (链接)。我很确定设备函数是通过降低和内联扩展实现的，即Numba设备函数不使用CUDA ABI。

3.没有CUDA代码或其他代码可以在主机CPU上运行。

如果您查看您问题中的代码，根据以上三点中的角度，基本上所有代码都是非法的。我建议一种替代设计模式，但由于代码实际上正在尝试做什么非常不明显，因此我无法这样做。不同部分的代码中存在太多奇怪的矛盾。

- talonmies