Numba 函数的编组对象代码
Marshaling object code for a Numba function
我有一个可以通过 Numba 解决的问题：为查询服务器创建 Numpy ufuncs 以 (a) 将简单操作合并为单次数据传递，减少我的 #1 热点（内存带宽），以及 (b ) 将第三方 C 函数封装为动态的 ufunc，为查询系统的用户提供更多功能。
我有一个累加器节点，它拆分查询并收集结果和计算节点，实际上 运行 Numpy（网络中的不同计算机）。如果 Numba 编译发生在计算节点上，这将是重复的工作，因为它们针对相同的查询在不同的数据分区上工作——相同的查询意味着相同的 Numba 编译。此外，即使是最简单的 Numba 编译也需要 96 毫秒——只要 运行 在数百万个点上进行查询计算，这在计算节点上提供的时间更好。
所以我想在累积节点上一次进行Numba编译，然后将其发送到计算节点，这样它们就可以运行它。我可以保证两者具有相同的硬件，因此目标代码是兼容的。
我一直在 Numba API 中搜索此功能，但没有找到它（除了没有文档的 numba.serialize 模块；我不确定它的用途是什么） .解决方案可能不是 Numba 包的 "feature"，而是一种利用某人对 Numba and/or LLVM 的内部知识的技术。有谁知道如何获取目标代码、编组并重构它？如果有帮助，我可以在两台机器上都安装 Numba，我只是不能在目标机器上做任何太昂贵的事情。
好的，这是可能的，解决方案大量使用了 Numba 下的 llvmlite 库。
获取序列化函数
首先我们用 Numba 定义一些函数。
import numba
@numba.jit("f8(f8)", nopython=True)
def example(x):
return x + 1.1
我们可以通过
访问目标代码
cres = example.overloads.values()[0] # 0: first and only type signature
elfbytes = cres.library._compiled_object
如果你打印出 elfbytes，你会看到它是一个 ELF-encoded 字节数组（bytes 对象，而不是 str 如果你在 Python 3).如果您要编译共享库或可执行文件，这就是文件中的内容，因此它可以移植到具有相同体系结构、相同库等的任何机器上。
这个bundle里面有几个函数，你可以通过转储LLVM IR看到：
print(cres.library.get_llvm_str())
我们想要的名字是__main__.example.float64，我们可以在LLVM IR中看到它的类型签名：
define i32 @"__main__.example.float64"(double* noalias nocapture %retptr, { i8*, i32 }** noalias nocapture readnone %excinfo, i8* noalias nocapture readnone %env, double %arg.x) #0 {
entry:
%.14 = fadd double %arg.x, 1.100000e+00
store double %.14, double* %retptr, align 8
ret i32 0
}
记下以备将来参考：第一个参数是指向 double 的指针，它会被结果覆盖，第二个和第三个参数是永远不会被使用的指针，最后一个参数是输入double.
（另请注意，我们可以使用 [x.name for x in cres.library._final_module.functions] 以编程方式获取函数名称。Numba 实际使用的入口点是 cres.fndesc.mangled_name。）
我们将这个 ELF 和函数签名从进行所有编译的机器传输到进行所有计算的机器。
读回
现在在计算机上，我们将使用完全没有 Numba 的 llvmlite（在 this page 之后）。初始化它：
import llvmlite.binding as llvm
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter() # yes, even this one
创建 LLVM 执行引擎：
target = llvm.Target.from_default_triple()
target_machine = target.create_target_machine()
backing_mod = llvm.parse_assembly("")
engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
现在劫持它的缓存机制让它加载我们的 ELF，命名为 elfbytes:
def object_compiled_hook(ll_module, buf):
pass
def object_getbuffer_hook(ll_module):
return elfbytes
engine.set_object_cache(object_compiled_hook, object_getbuffer_hook)
像我们刚刚编译 IR 一样完成引擎，但实际上我们跳过了这一步。引擎将加载我们的 ELF，认为它来自其 disk-based 缓存。
engine.finalize_object()
我们现在应该在这个引擎的 space 中找到我们的功能。如果下面的returns 0L，有问题。它应该是一个函数指针。
func_ptr = engine.get_function_address("__main__.example.float64")
现在我们需要将 func_ptr 解释为我们可以调用的 ctypes 函数。我们必须手动设置签名。
import ctypes
pdouble = ctypes.c_double * 1
out = pdouble()
pointerType = ctypes.POINTER(None)
dummy1 = pointerType()
dummy2 = pointerType()
# restype first then argtypes...
cfunc = ctypes.CFUNCTYPE(ctypes.c_int32, pdouble, pointerType, pointerType, ctypes.c_double)(func_ptr)
现在我们可以调用它了：
cfunc(out, dummy1, dummy2, ctypes.c_double(3.14))
print(out[0])
# 4.24, which is 3.14 + 1.1. Yay!
更多并发症
如果 JITed 函数有数组输入（毕竟，您想对编译代码中的许多值进行紧密循环，而不是 Python），Numba 会生成识别 Numpy 数组的代码。这个调用约定非常复杂，包括异常对象的 pointers-to-pointers 和 Numpy 数组附带的所有元数据作为单独的参数。它不会生成一个入口点，您可以使用 Numpy 的 ctypes 接口。
但是，它确实提供了一个非常 high-level 的入口点，它以 Python *args, **kwds 作为参数并在内部解析它们。下面是你如何使用它。
首先找到名字以"cpython."开头的函数：
name = [x.name for x in cres.library._final_module.functions if x.name.startswith("cpython.")][0]
其中应该只有一个。然后，经过序列化和反序列化后，使用上面描述的方法得到它的函数指针：
func_ptr = engine.get_function_address(name)
并使用三个 PyObject* 参数和一个 PyObject* return 值进行转换。 （LLVM 认为这些是 i8*。）
class PyTypeObject(ctypes.Structure):
_fields_ = ("ob_refcnt", ctypes.c_int), ("ob_type", ctypes.c_void_p), ("ob_size", ctypes.c_int), ("tp_name", ctypes.c_char_p)
class PyObject(ctypes.Structure):
_fields_ = ("ob_refcnt", ctypes.c_int), ("ob_type", ctypes.POINTER(PyTypeObject))
PyObjectPtr = ctypes.POINTER(PyObject)
cpythonfcn = ctypes.CFUNCTYPE(PyObjectPtr, PyObjectPtr, PyObjectPtr, PyObjectPtr)(fcnptr)
这三个参数中的第一个是闭包（函数访问的全局变量），我假设我们不需要它。使用显式参数而不是闭包。我们可以利用 CPython 的 id() 实现 return 指针值这一事实来制作 PyObject 指针。
def wrapped(*args, **kwds):
closure = ()
return cpythonfcn(ctypes.cast(id(closure), PyObjectPtr), ctypes.cast(id(args), PyObjectPtr), ctypes.cast(id(kwds), PyObjectPtr))
现在函数可以调用为
wrapped(whatever_numpy_arguments, ...)
就像原来的 Numba 调度程序功能一样。
底线
毕竟，这值得吗？使用 Numba 进行 end-to-end 编译——最简单的方法——这个简单的函数需要 50 毫秒。要求 -O3 而不是默认的 -O2，我可以使它慢 40%。
然而，拼接 pre-compiled ELF 文件需要 0.5 毫秒：快了 100 倍。此外，编译时间会随着更复杂的函数而增加，但 splicing-in 过程对于任何函数都应始终花费 0.5 毫秒。
对于我的应用来说，这绝对是值得的。这意味着我一次可以在 10 MB 上执行计算，并且大部分时间都花在计算上（做实际工作），而不是编译（准备工作）。将其放大 100 倍，我将不得不一次对 1 GB 执行计算。由于一台机器被限制为 order-of 100 GB，并且必须在 order-of 100 个进程之间共享，我将面临更大的资源限制、负载平衡问题等危险，因为问题会太细化。
但对于其他应用来说，50 ms 不算什么。这完全取决于您的应用程序。
我有一个可以通过 Numba 解决的问题：为查询服务器创建 Numpy ufuncs 以 (a) 将简单操作合并为单次数据传递，减少我的 #1 热点（内存带宽），以及 (b ) 将第三方 C 函数封装为动态的 ufunc，为查询系统的用户提供更多功能。
我有一个累加器节点，它拆分查询并收集结果和计算节点，实际上 运行 Numpy（网络中的不同计算机）。如果 Numba 编译发生在计算节点上，这将是重复的工作，因为它们针对相同的查询在不同的数据分区上工作——相同的查询意味着相同的 Numba 编译。此外，即使是最简单的 Numba 编译也需要 96 毫秒——只要 运行 在数百万个点上进行查询计算，这在计算节点上提供的时间更好。
所以我想在累积节点上一次进行Numba编译，然后将其发送到计算节点，这样它们就可以运行它。我可以保证两者具有相同的硬件，因此目标代码是兼容的。
我一直在 Numba API 中搜索此功能，但没有找到它（除了没有文档的 numba.serialize 模块；我不确定它的用途是什么） .解决方案可能不是 Numba 包的 "feature"，而是一种利用某人对 Numba and/or LLVM 的内部知识的技术。有谁知道如何获取目标代码、编组并重构它？如果有帮助，我可以在两台机器上都安装 Numba，我只是不能在目标机器上做任何太昂贵的事情。
好的，这是可能的，解决方案大量使用了 Numba 下的 llvmlite 库。
获取序列化函数
首先我们用 Numba 定义一些函数。
import numba
@numba.jit("f8(f8)", nopython=True)
def example(x):
return x + 1.1
我们可以通过
访问目标代码cres = example.overloads.values()[0] # 0: first and only type signature
elfbytes = cres.library._compiled_object
如果你打印出 elfbytes，你会看到它是一个 ELF-encoded 字节数组（bytes 对象，而不是 str 如果你在 Python 3).如果您要编译共享库或可执行文件，这就是文件中的内容，因此它可以移植到具有相同体系结构、相同库等的任何机器上。
这个bundle里面有几个函数，你可以通过转储LLVM IR看到：
print(cres.library.get_llvm_str())
我们想要的名字是__main__.example.float64，我们可以在LLVM IR中看到它的类型签名：
define i32 @"__main__.example.float64"(double* noalias nocapture %retptr, { i8*, i32 }** noalias nocapture readnone %excinfo, i8* noalias nocapture readnone %env, double %arg.x) #0 {
entry:
%.14 = fadd double %arg.x, 1.100000e+00
store double %.14, double* %retptr, align 8
ret i32 0
}
记下以备将来参考：第一个参数是指向 double 的指针，它会被结果覆盖，第二个和第三个参数是永远不会被使用的指针，最后一个参数是输入double.
（另请注意，我们可以使用 [x.name for x in cres.library._final_module.functions] 以编程方式获取函数名称。Numba 实际使用的入口点是 cres.fndesc.mangled_name。）
我们将这个 ELF 和函数签名从进行所有编译的机器传输到进行所有计算的机器。
读回
现在在计算机上，我们将使用完全没有 Numba 的 llvmlite（在 this page 之后）。初始化它：
import llvmlite.binding as llvm
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter() # yes, even this one
创建 LLVM 执行引擎：
target = llvm.Target.from_default_triple()
target_machine = target.create_target_machine()
backing_mod = llvm.parse_assembly("")
engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
现在劫持它的缓存机制让它加载我们的 ELF，命名为 elfbytes:
def object_compiled_hook(ll_module, buf):
pass
def object_getbuffer_hook(ll_module):
return elfbytes
engine.set_object_cache(object_compiled_hook, object_getbuffer_hook)
像我们刚刚编译 IR 一样完成引擎，但实际上我们跳过了这一步。引擎将加载我们的 ELF，认为它来自其 disk-based 缓存。
engine.finalize_object()
我们现在应该在这个引擎的 space 中找到我们的功能。如果下面的returns 0L，有问题。它应该是一个函数指针。
func_ptr = engine.get_function_address("__main__.example.float64")
现在我们需要将 func_ptr 解释为我们可以调用的 ctypes 函数。我们必须手动设置签名。
import ctypes
pdouble = ctypes.c_double * 1
out = pdouble()
pointerType = ctypes.POINTER(None)
dummy1 = pointerType()
dummy2 = pointerType()
# restype first then argtypes...
cfunc = ctypes.CFUNCTYPE(ctypes.c_int32, pdouble, pointerType, pointerType, ctypes.c_double)(func_ptr)
现在我们可以调用它了：
cfunc(out, dummy1, dummy2, ctypes.c_double(3.14))
print(out[0])
# 4.24, which is 3.14 + 1.1. Yay!
更多并发症
如果 JITed 函数有数组输入（毕竟，您想对编译代码中的许多值进行紧密循环，而不是 Python），Numba 会生成识别 Numpy 数组的代码。这个调用约定非常复杂，包括异常对象的 pointers-to-pointers 和 Numpy 数组附带的所有元数据作为单独的参数。它不会生成一个入口点，您可以使用 Numpy 的 ctypes 接口。
但是，它确实提供了一个非常 high-level 的入口点，它以 Python *args, **kwds 作为参数并在内部解析它们。下面是你如何使用它。
首先找到名字以"cpython."开头的函数：
name = [x.name for x in cres.library._final_module.functions if x.name.startswith("cpython.")][0]
其中应该只有一个。然后，经过序列化和反序列化后，使用上面描述的方法得到它的函数指针：
func_ptr = engine.get_function_address(name)
并使用三个 PyObject* 参数和一个 PyObject* return 值进行转换。 （LLVM 认为这些是 i8*。）
class PyTypeObject(ctypes.Structure):
_fields_ = ("ob_refcnt", ctypes.c_int), ("ob_type", ctypes.c_void_p), ("ob_size", ctypes.c_int), ("tp_name", ctypes.c_char_p)
class PyObject(ctypes.Structure):
_fields_ = ("ob_refcnt", ctypes.c_int), ("ob_type", ctypes.POINTER(PyTypeObject))
PyObjectPtr = ctypes.POINTER(PyObject)
cpythonfcn = ctypes.CFUNCTYPE(PyObjectPtr, PyObjectPtr, PyObjectPtr, PyObjectPtr)(fcnptr)
这三个参数中的第一个是闭包（函数访问的全局变量），我假设我们不需要它。使用显式参数而不是闭包。我们可以利用 CPython 的 id() 实现 return 指针值这一事实来制作 PyObject 指针。
def wrapped(*args, **kwds):
closure = ()
return cpythonfcn(ctypes.cast(id(closure), PyObjectPtr), ctypes.cast(id(args), PyObjectPtr), ctypes.cast(id(kwds), PyObjectPtr))
现在函数可以调用为
wrapped(whatever_numpy_arguments, ...)
就像原来的 Numba 调度程序功能一样。
底线
毕竟，这值得吗？使用 Numba 进行 end-to-end 编译——最简单的方法——这个简单的函数需要 50 毫秒。要求 -O3 而不是默认的 -O2，我可以使它慢 40%。
然而，拼接 pre-compiled ELF 文件需要 0.5 毫秒：快了 100 倍。此外，编译时间会随着更复杂的函数而增加，但 splicing-in 过程对于任何函数都应始终花费 0.5 毫秒。
对于我的应用来说，这绝对是值得的。这意味着我一次可以在 10 MB 上执行计算，并且大部分时间都花在计算上（做实际工作），而不是编译（准备工作）。将其放大 100 倍，我将不得不一次对 1 GB 执行计算。由于一台机器被限制为 order-of 100 GB，并且必须在 order-of 100 个进程之间共享，我将面临更大的资源限制、负载平衡问题等危险，因为问题会太细化。
但对于其他应用来说，50 ms 不算什么。这完全取决于您的应用程序。