为与其他语言一起使用而编译Python函数

Numba 可以将 Python 代码编译为 PTX 或 LTO-IR，以便 Python 函数可以被包含在用其他语言（如 C/C++）编写的 CUDA 代码中。它通常用于支持在库或应用程序上下文中编写的 Python 用户定义函数。

编译API可以在没有GPU的情况下使用，因为它不使用驱动程序函数，并且在过程中避免初始化CUDA。它通过以下函数调用：

numba.cuda.compile(pyfunc, sig, debug=False, lineinfo=False, device=True, fastmath=False, cc=None, opt=True, abi='c', abi_info=None, output='ptx')[源代码]

为给定的一组参数类型，将Python函数编译为PTX或LTO-IR。

参数:

  • pyfunc – 要编译的Python函数。

  • sig – 表示函数输入和输出类型的签名。如果这是一个不带返回类型的参数类型元组，则此函数推断的返回类型为返回类型。如果传递了包含返回类型的签名，编译后的代码将包括从推断的返回类型到指定返回类型的强制转换，并且此函数将返回指定的返回类型。

  • debug (bool) – 是否在编译代码中包含调试信息。

  • lineinfo (bool) – 是否在编译后的代码中包含源代码的行映射。通常这用于优化代码（因为调试模式会自动包含此功能），因此我们希望在LLVM IR中包含调试信息，但在最终输出中仅包含行映射。

  • device (bool) – 是否编译设备函数。

  • fastmath (bool) – 是否启用快速数学标志（ftz=1, prec_sqrt=0, prec_div=, 和 fma=1）

  • cc (tuple) – 要编译的计算能力，以元组 (MAJOR, MINOR) 形式表示。默认为 (5, 0)。

  • opt (bool) – 启用优化。默认为 True。

  • abi (str) – 编译函数的 ABI - 可以是 "numba" 或 "c"。请注意，Numba ABI 不被认为是稳定的。目前仅支持设备函数的 C ABI。

  • abi_info (dict) – 一组特定于 ABI 的选项。"c" ABI 支持一个选项，"abi_name"，用于提供包装函数的名称。"numba" ABI 没有选项。

  • output (str) – 要生成的输出类型，可以是 "ptx" 或 "ltoir"。

返回:

(代码, 推断返回类型): 编译后的代码和推断的返回类型

返回类型:

tuple

如果设备可用且需要为当前设备的计算能力编译代码（例如在使用 Numba 构建 JIT 编译工作流时），可以使用 compile_for_current_device 函数：

numba.cuda.compile_for_current_device(pyfunc, sig, debug=False, lineinfo=False, device=True, fastmath=False, opt=True, abi='c', abi_info=None, output='ptx')[源代码]

为当前设备的计算能力编译一个Python函数为PTX或LTO-IR，给定一个签名。这将调用 compile() 并传入当前设备合适的 cc 值。

大多数用户应使用上述两个函数；为了与现有用例保持向后兼容性，还提供了以下函数：

numba.cuda.compile_ptx(pyfunc, sig, debug=False, lineinfo=False, device=False, fastmath=False, cc=None, opt=True, abi='numba', abi_info=None)[源代码]

根据给定的签名将Python函数编译为PTX。参见 compile()。此函数的默认设置是编译一个符合Numba ABI的内核，而不是 compile() 的默认设置，即编译一个符合C ABI的设备函数。

numba.cuda.compile_ptx_for_current_device(pyfunc, sig, debug=False, lineinfo=False, device=False, fastmath=False, opt=True, abi='numba', abi_info=None)[源代码]

为当前设备的计算能力编译一个具有给定签名的Python函数到PTX。参见 compile_ptx()。

使用 C ABI

Numba 内部使用自己的 ABI - 这如 设备功能 ABI 中所述，没有 extern "C" 修饰符。调用 Numba ABI 设备函数需要解决三个问题：

  • 函数名将根据 Numba 的 ABI 规则进行混淆 - 这些规则基于 Itanium C++ ABI 规则，但在此基础上有所扩展。

  • Python 的返回值预计会被存储到第一个参数传入的指针值中。

  • 编译函数的返回值将包含一个状态码，而不是函数的返回值。对于在Numba之外使用Numba编译的函数，这通常可以忽略。

解决所有这些问题的一个简单方法是使用 C ABI 编译设备函数。这将导致以下结果：

  • 编译代码中设备函数的名称可以被控制。默认情况下，它将与Python中的函数名称匹配，因此很容易确定。这是函数的 __name__，而不是 __qualname__，因为 __qualname__ 编码了额外的范围信息，这会使函数名称难以预测，并且在很多情况下，在C语言中是非法的标识符。

  • Python代码的返回值被放置在编译函数的返回值中。

  • 状态码被忽略 / 未报告，因此不需要处理它们。

如果需要指定编译函数的名称，可以通过在 abi_info 字典中传递名称来控制，该名称位于键 'abi_name' 下。

使用 compile() 和 compile_for_current_device() 函数时，默认使用 C ABI 进行编译。为了保持与现有用例的兼容性，compile_ptx() 和 compile_ptx_for_current_device() 函数默认使用 Numba ABI。

C 和 Numba ABI 示例

以下函数：

def add(x, y):
    return x + y

使用 Numba ABI 编译，例如：

ptx, resty = cuda.compile_ptx(add, int32(int32, int32), device=True)

在PTX中的结果，其中函数原型为：

.visible .func  (.param .b32 func_retval0) _ZN8__main__3addB2v1B94cw51cXTLSUwv1sCUt9Uw1VEw0NRRQPKzLTg4gaGKFsG2oMQGEYakJSQB1PQBk0Bynm21OiwU1a0UoLGhDpQE8oxrNQE_3dEii(
    .param .b64 _ZN8__main__3addB2v1B94cw51cXTLSUwv1sCUt9Uw1VEw0NRRQPKzLTg4gaGKFsG2oMQGEYakJSQB1PQBk0Bynm21OiwU1a0UoLGhDpQE8oxrNQE_3dEii_param_0,
    .param .b32 _ZN8__main__3addB2v1B94cw51cXTLSUwv1sCUt9Uw1VEw0NRRQPKzLTg4gaGKFsG2oMQGEYakJSQB1PQBk0Bynm21OiwU1a0UoLGhDpQE8oxrNQE_3dEii_param_1,
    .param .b32 _ZN8__main__3addB2v1B94cw51cXTLSUwv1sCUt9Uw1VEw0NRRQPKzLTg4gaGKFsG2oMQGEYakJSQB1PQBk0Bynm21OiwU1a0UoLGhDpQE8oxrNQE_3dEii_param_2
 )

请注意，有三个参数，用于返回值的指针、x 和 y。名称在 Numba 内部以一种难以预测的方式进行了处理。

使用以下命令编译为 C ABI：

ptx, resty = cuda.compile_ptx(add, int32(int32, int32), device=True, abi="c")

而是产生以下 PTX 原型：

.visible .func  (.param .b32 func_retval0) add(
    .param .b32 add_param_0,
    .param .b32 add_param_1
)

函数名与 Python 源代码中的函数名匹配，并且正好有两个参数，分别用于 x 和 y。函数的结果直接放在返回值中：

add.s32      %r3, %r2, %r1;
st.param.b32         [func_retval0+0], %r3;

为了区分编译后的 add() 函数的不同变体，以下示例在 abi_info 字典中指定了其 ABI 名称：

ptx, resty = cuda.compile_ptx(add, float32(float32, float32), device=True,
                              abi="c", abi_info={"abi_name": "add_f32"})

生成PTX原型：

.visible .func  (.param .b32 func_retval0) add_f32(
    .param .b32 add_f32_param_0,
    .param .b32 add_f32_param_1
)

这不会与其他名称的定义冲突（例如上面的 int32 变体）。