How to extend Numba so it supports pybind11-wrapped C++ struct?

Hi!

First, a bit more on my use case: I have a C++ application embedding the Python interpreter through pybind11, and calling specific Python functions with structs containing data from the application as parameters.

In this context, I am trying to call @njit-decorated functions from the functions called by the application (this is a first step, as I would like to directly call @cfunc-decorated functions later on), and I have trouble passing the Python object wrapping my structs to these functions as Numba does not recognize those.

To solve this, I’m trying to register these structs with the extending API, following the Interval example, but the furthest I’ve gone is getting stack overflows on the compilation. Anyway, I’m not even sure this is the correct way to go about this as my objects are not native Python objects. Ideally I would like to teach Numba how to directly access the various fields of my structs but I’m not sure which direction to take to do this.

Any help would be greatly appreciated!
Cheers

2 Likes

I would love this feature too. I am linking c-source code to pytorch and jax.

I am trying to wrap the pybind11 functions I use for pytorch with the jax header in a python script. To do this, I use numba’s cfunc to wrap my original pybind11 interface. An example is shown here. I am unfamiliar with this low-level detail but I hope this post demonstrates another use case for this feature.

Potentially using FFI to read the pybind11 so file will work soon if this PR does what I think it does.

I did a talk at GTC 2022 walking through how to implement exactly this use case - providing access to structures from C/C++ libraries inside Numba-jitted functions: Enabling Python User-Defined Functions in Accelerated Applications with Numba | NVIDIA On-Demand

The talk walks through doing this for a CUDA library in the CUDA target, but the principles are the same for the CPU target and host C/C++.