是否可以调用间接调用另一个 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_gpu() 到euclidean_distance()，而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 很陌生 - 有什么我可以在这里做的吗？请注意，这些函数实际上比显示的更复杂，所以我想要一个实际的解决方案，而不仅仅是内联函数。

Answer 1

这个问题的文字和代码有很多误解。简单地说：

1. Numba 内核无法启动其他 Numba 内核（来自 documentation）

... 较新的 CUDA 设备支持设备端内核启动；此功能称为动态并行，但 Numba 目前不支持它

1. Numba 内核可以调用 Numba 设备函数，而 Numba 设备函数可以调用其他 Numba 设备函数 (link)。而且我很确定设备功能通过降低和内联扩展来工作，即 Numba 设备功能不使用 CUDA ABI

2. 没有 CUDA 代码、Numba 或其他方式可以在主机 CPU 上运行代码

如果您查看问题中的代码，应该立即明显看出，从上面第 1、2 和 3 项的角度来看，基本上所有代码都是非法的。

我会建议另一种设计模式，但代码实际上试图做什么并不明显，我无法这样做。为此，代码的不同部分有太多奇怪的矛盾。