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

Question

I would like to interface C++ library in python which takes function pointer as argument. I look that it is possible to make call in C with PyEval_CallObject, but to proceed it further I need correct signature (types of input and output arguments). Is it possible to return callback from python with specified signature?

Also I am a bit performance woried, so I also have looked on python numba project which does compilation of python function. I am quite interested if it can be accessed in C/C++ so to improve performance.

Answer 1

Numba's jited function is compiled in-memory, not in file. This function has regular Pythonic interface, which is a thin wrapper that converts Python arguments to C types and calls jited machine code as regular C/Assembler function.

Numba compiles (jits) function on the fly and on-demand, meaning that it compiles on very first call. Also Numba precompiles function with different arguments types, i.e. if you pass one type Numba compiles one instance of machine code, if you pass another type as argument Numba compiles another instance. Compiled code is cached, same types of arguments on second call use cached version of machine code.

Numba's functions are regular Pythonic functions. But Numba is very friendly to ctypes library, so it supports conversion to ctypes' CFUNCTYPE which allows access directly to machine code with given fixed C types of arguments.

In order to specify just one type of arguments you have to write something like

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

i.e. pass fixed types of arguments to @njit decorator.

BTW, you can use @jit (less optimized, but can accept almost any Python code), or @njit (more optimized, but needs strict rules for code writing).

Below I created example of C++ code for you, which calls Numba's jited function from C++. I showed how to call it in three ways - 1) as C function using ctypes wrapper. 2) as Python function using PyObject_Call(). 3) through @numba.cfunc decorator and .address property, read about it here (3rd suggested by @max9111).

If you need speed and smallest CALL overhead when having just C types then choose solution with @cfunc, it is the most efficient (fastest) solution (also shortest to implement), according to @max9111. Basically .address property just gives direct address of C function's machine code without any overhead and preparation. Use pyfunc method if you have Python objects as arguments and/or you want to pass objects of different Python types and/or if you have @jit (not @njit) and/or you don't care about speed.

Important Note!!! My code doesn't cleanup anything, you should delete all created references to all objects through Py_XDECREF(). Also I might be doing not all error checking, you have to error-check every Python C API function that may return error. I did this two simplifications to make code shorter and more understandable, just to show main points of example.

Try it online!

#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;
    }
}

Output:

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