是否可以从C/ c++中调用numba jited函数?

Is it possible to call numba jited function from C/C++?

本文关键字:numba 调用 jited 函数 c++ 是否      更新时间:2023-10-16

我想在python中接口c++库,它将函数指针作为参数。我认为有可能在C中使用PyEval_CallObject进行调用,但要进一步进行,我需要正确的签名(输入和输出参数的类型)。是否有可能从python返回带有指定签名的回调?

我也有点担心性能,所以我也看了python numba项目,它编译python函数。我很感兴趣,如果它可以在C/c++中访问,以提高性能。

Numba的jit函数在内存中编译,而不是在文件中编译。该函数具有常规的Python接口,它是一个薄包装器,将Python参数转换为C类型,并将丢弃的机器码调用为常规的C/汇编函数。

Numba动态地按需编译(jit)函数,这意味着它在第一次调用时编译。此外,Numba预编译函数与不同的参数类型,即,如果你传递一个类型Numba编译一个实例的机器代码,如果你传递另一个类型作为参数Numba编译另一个实例。编译后的代码被缓存,相同类型的参数在第二次调用时使用缓存版本的机器码。

Numba的函数是常规的python函数。但是Numba对ctypes库非常友好,所以它支持转换到ctypes的CFUNCTYPE,这允许使用给定的固定C类型参数直接访问机器码。

为了只指定一种类型的参数,你必须这样写

@numba.njit(numba.int64(numba.int64, numba.int64))

。向@njit装饰器传递固定类型的参数

顺便说一句,您可以使用@jit(优化程度较低,但可以接受几乎任何Python代码)或@njit(优化程度较高,但需要严格的代码编写规则)。

下面我为您创建了一个c++代码示例,它从c++调用Numba的jit函数。我展示了如何以三种方式调用它—1)使用ctypes包装器作为C函数。2)使用PyObject_Call()作为Python函数。3)通过@numba.cfunc装饰器和.address属性,阅读这里(@max9111建议的第三个)。

如果只有C类型时需要速度和最小的调用开销,那么选择@cfunc的解决方案,这是最有效(最快)的解决方案(也是最短的实现),根据@max9111。基本上,.address属性只是给出了C函数机器码的直接地址,没有任何开销和准备。如果你有Python对象作为参数和/或者你想传递不同Python类型的对象和/或者如果你有@jit(不是@njit)和/或者你不在乎速度,请使用pyfunc方法。

重要提示! !我的代码没有清除任何东西,您应该通过Py_XDECREF()删除所有创建的对所有对象的引用。另外,我可能没有做所有的错误检查,你必须检查每个可能返回错误的Python C API函数。我做这两个简化是为了使代码更短,更容易理解,只是为了显示示例的要点。

上网试试!

#include <stdexcept>
#include <string>
#include <iostream>
#include <cstdint>
#include <Python.h>
#define ASSERT(cond) { if (!(cond)) throw std::runtime_error( 
    "Assertion (" #cond ") failed at line " + std::to_string(__LINE__) + "!"); }
#define CH(cond) [&]{ auto r = (cond); if (PyErr_Occurred()) { PyErr_Print(); std::cerr << std::flush; 
    throw std::runtime_error("PyAssertion (" #cond ") failed at line " + std::to_string(__LINE__) + "!"); } return r; }()
int main() {
    try {
        std::cout << "Wait..." << std::endl;
        Py_Initialize();
        auto globals = CH(PyDict_New()), locals = CH(PyDict_New());
        CH(PyRun_String(R"(
import numba as nb, ctypes
@nb.njit(nb.int64(nb.int64, nb.int64))
def mul1(a, b):
    return a * b
@nb.cfunc(nb.int64(nb.int64, nb.int64))
def mul2(a, b):
    return a * b
cmul1 = ctypes.CFUNCTYPE(
    ctypes.c_int64, ctypes.c_int64, ctypes.c_int64)(mul1)
addr1 = ctypes.cast(cmul1, ctypes.c_void_p).value
addr2 = mul2.address
        )", Py_file_input, globals, locals));
        //std::cout << CH(PyUnicode_AsUTF8AndSize(CH(PyObject_Str(locals)), nullptr)) << std::endl;
        auto cfunc1 = (int64_t (*)(int64_t, int64_t))
            CH(PyLong_AsUnsignedLongLong(CH(PyDict_GetItemString(locals, "addr1"))));
        auto cfunc2 = (int64_t (*)(int64_t, int64_t))
            CH(PyLong_AsUnsignedLongLong(CH(PyDict_GetItemString(locals, "addr2"))));
        std::cout << "pyfunc: 3 * 5 = " << CH(PyUnicode_AsUTF8AndSize(
            CH(PyObject_Str(CH(PyObject_Call(
                CH(PyDict_GetItemString(locals, "mul1")), PyTuple_Pack(
                    2, PyLong_FromLongLong(3), PyLong_FromLongLong(5)), nullptr /* named args */
            )))), nullptr)) << std::endl << std::flush;
        std::cout << "cfunc1 (0x" << std::hex << uint64_t(cfunc1) << std::dec
            << "): 3 * 5 = " << cfunc1(3, 5) << std::endl << std::flush;
        std::cout << "cfunc2 (0x" << std::hex << uint64_t(cfunc2) << std::dec
            << "): 3 * 5 = " << cfunc2(3, 5) << std::endl << std::flush;
        ASSERT(Py_FinalizeEx() == 0);
        return 0;
    } catch (std::exception const & ex) {
        std::cerr << "Exception: " << ex.what() << std::endl << std::flush;
        return -1;
    }
}
输出:

Wait...
pyfunc: 3 * 5 = 15
cfunc1 (0x7f5db072f080): 3 * 5 = 15
cfunc2 (0x7f5db05b2010): 3 * 5 = 15