从Python内核调用外部函数

Python 内核可以调用用其他语言编写的设备函数。直接支持 CUDA C/C++、PTX 和二进制对象（cubins、fat binaries 等）；其他语言的源代码必须首先编译为 PTX。Python 内核调用外部设备函数的组成部分包括：

设备函数在外语中的实现（例如 CUDA C）。
Python 中设备函数的声明。
一个与外部函数链接并调用的内核。

设备功能 ABI

Numba 的调用设备函数的 ABI 在 C/C++ 中定义了以下原型：

extern "C"
__device__ int
function(
  T* return_value,
  ...
);

原型的组成部分如下：

extern "C" 用于防止名称混淆，以便在Python中更容易声明函数。它可以被移除，但在Python中声明函数时必须使用混淆后的名称。
__device__ 是定义函数为设备函数的必需项。
返回值始终是 int 类型，用于指示是否发生了 Python 异常。由于 Python 异常不会在外部函数中发生，因此调用者应始终将其设置为 0。
第一个参数是一个指向类型为 T 的返回值的指针，该返回值在本地地址空间 [1] 中分配，并由调用者传入。如果函数返回一个值，被调用者应设置指针所指向的对象以存储返回值。
后续参数应与从Python内核传递给函数的参数类型和顺序相匹配。

用其他语言编写的函数必须编译成符合此原型规范的PTX。

接受两个浮点数并返回一个浮点数的函数将具有以下原型：

extern "C"
__device__ int
mul_f32_f32(
  float* return_value,
  float x,
  float y
);

注释

Python 中的声明

要在 Python 中声明一个外部设备函数，请使用 declare_device():

numba.cuda.declare_device(name, sig)[源代码]

声明一个外部函数的签名。返回一个描述符，该描述符可用于从Python内核调用该函数。

参数:

name (str) – 外部函数的名称。
sig – 函数的 Numba 签名。

返回的描述符名称不必与外部函数的名称匹配。例如，当：

mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)')

一旦声明，在核函数中调用 mul(a, b) 将被翻译为编译代码中的 mul_f32_f32(a, b) 调用。

传递指针

Numba 的调用约定要求为数组参数传递多个值。这些包括数据指针以及形状、步幅和其他信息。这与大多数 C/C++ 函数的预期不兼容，后者通常只期望数据指针。为了使 C 设备代码和 Python 内核之间的调用约定保持一致，有必要使用 C 指针类型声明数组参数。

例如，具有以下原型的函数：

numba/cuda/tests/doc_examples/ffi/functions.cu

extern "C"
__device__ int
sum_reduce(
  float* return_value,
  float* array,
  int n
);

将声明如下：

来自 numba/cuda/tests/doc_examples/test_ffi.py 中的 test_ex_from_buffer

signature = 'float32(CPointer(float32), int32)'
sum_reduce = cuda.declare_device('sum_reduce', signature)

要获取用于传递给外部函数的数组数据指针，请使用 cffi.FFI 实例的 from_buffer() 方法。例如，使用 sum_reduce 函数的内核可以定义为：

来自 numba/cuda/tests/doc_examples/test_ffi.py 中的 test_ex_from_buffer

import cffi
ffi = cffi.FFI()

@cuda.jit(link=[functions_cu])
def reduction_caller(result, array):
    array_ptr = ffi.from_buffer(array)
    result[()] = sum_reduce(array_ptr, len(array))

其中 result 和 array 都是 float32 数据的数组。

链接和调用函数

link 关键字参数在 @cuda.jit 装饰器中接受一个由绝对路径或相对于当前工作目录的路径指定的文件名列表。名称以 .cu 结尾的文件将使用 NVIDIA 运行时编译器 (NVRTC) 编译并作为 PTX 链接到内核中；其他文件将直接传递给 CUDA 链接器。

例如，以下内核调用 functions.cu 文件中实现的 mul_f32_f32() 函数，该函数在上文中声明为 mul()：

@cuda.jit(link=['functions.cu'])
def multiply_vectors(r, x, y):
    i = cuda.grid(1)

    if i < len(r):
        r[i] = mul(x[i], y[i])

C/C++ 支持

通过使用NVRTC，提供了对CUDA C/C++代码编译和链接的支持，但需注意以下事项：

只有在使用 NVIDIA 绑定的情况下才可用。请参阅 NUMBA_CUDA_USE_NVIDIA_BINDING。
必须为已安装的 NVIDIA CUDA 绑定提供适合版本的 NVRTC 库。
CUDA 包含路径在 Linux 上默认假设为 /usr/local/cuda/include，在 Windows 上为 $env:CUDA_PATH\include。可以使用环境变量 NUMBA_CUDA_INCLUDE_PATH 进行修改。
CUDA 包含目录将在包含路径上提供给 NVRTC；不支持额外的包含。

完整示例

此示例演示了调用用 CUDA C 编写的外部函数，以将两个数组中的数字对相乘。

外部函数编写如下：

numba/cuda/tests/doc_examples/ffi/functions.cu

// Foreign function example: multiplication of a pair of floats

extern "C" __device__ int
mul_f32_f32(
  float* return_value,
  float x,
  float y)
{
  // Compute result and store in caller-provided slot
  *return_value = x * y;

  // Signal that no Python exception occurred
  return 0;
}

Python 代码和内核是：

来自 numba/cuda/tests/doc_examples/test_ffi.py 中的 test_ex_linking_cu

from numba import cuda
import numpy as np
import os

# Declaration of the foreign function
mul = cuda.declare_device('mul_f32_f32', 'float32(float32, float32)')

# Path to the source containing the foreign function
# (here assumed to be in a subdirectory called "ffi")
basedir = os.path.dirname(os.path.abspath(__file__))
functions_cu = os.path.join(basedir, 'ffi', 'functions.cu')

# Kernel that links in functions.cu and calls mul
@cuda.jit(link=[functions_cu])
def multiply_vectors(r, x, y):
    i = cuda.grid(1)

    if i < len(r):
        r[i] = mul(x[i], y[i])

# Generate random data
N = 32
np.random.seed(1)
x = np.random.rand(N).astype(np.float32)
y = np.random.rand(N).astype(np.float32)
r = np.zeros_like(x)

# Run the kernel
multiply_vectors[1, 32](r, x, y)

# Sanity check - ensure the results match those expected
np.testing.assert_array_equal(r, x * y)

备注

上面的例子是为了说明一个外部函数调用而设计的，它非常简单——由于网格较小和外部函数的工作量较轻，因此不会期望它具有特别高的性能。