Allow mixing types (C or Python) within single jit-function

[polkovnikov]

Feature request

As far as I know (maybe I'm wrong) Numba automatically decides what types to use for variables.

In njit-functions (nopython=True) all types are C static types besides, only inside numba.objmode() block Python types are allowed.

In jit-function (nopython=False) all types are Python types.

Basically there is a strict distinction - 1) pure-C mode (C-only types) 2) pure-Python mode (Python-only types).

If I'm correct then Numba uses Python C API to create JIT code which is compiled by Clang/LLVM.

It means that there is no reason why this C code can't mix pure-C and pure-Python types within neighboring expressions of single function.

In Cython there is cdef which allows users to manually tell that this type is C-only. Would be great if Numba can mix C/Python types together too, like Cython does.

Also like in Cython Numba should allow to specify manually any type of local variable. As I know this is already allowed by passing locals = {'x': SomeType} to @jit decorator. So if C/Python types are mixed then this locals = argument should allow specifying both C and Python types.

Why is it needed? I often just add @njit/@jit decorator to ready-made Python functions and Numba perfectly JITs and optimizes on the fly. That is very nice. For example I specified @jit(nopython = False, objmode = True), it compiles pure object-mode code, but now I want to have some optimizations to this code without migrating fully to @njit. I want most of variables to stay Python objects as they are classes, but now I want to optimize some very computationally intense variables (like loop variables or some mathematical accumulators) and for that I want to specify that they should be C types like locals = {'k': 'c_uint64'}.

Basically by this feature request I wanted to be able to have some intermediate states of code between pure Python and pure C.

In @njit it is currently allowed to have with blocks of objmode, which partially solves my request. But this with-blocks don't allow (without extra tricks) for Python objects to leave a block. There is as I remember numba_passthrough library that allows such tricks as Python objects leaving objmode blocks, but this solution is very un-intuitive, would be nice that Numba can support this natively. And even more - allow mix of any C/Python types within same code, like Cython allows.

[stuartarchibald]

Thanks for opening this feature request.

As far as I know (maybe I'm wrong) Numba automatically decides what types to use for variables.

By default, yes.

In njit-functions (nopython=True) all types are C static types besides, only inside numba.objmode() block Python types are allowed.

Partly true. Numba has it's own type system which is a hybrid of NumPy types and Python types, there's no real connection to C lineage languages other than sometimes the types would be inferred similarly and some of the type names/concepts are shared.

In jit-function (nopython=False) all types are Python types.

All types are the numba.types.pyobject type, this reflects that everything in python is an object.

Basically there is a strict distinction - 1) pure-C mode (C-only types) 2) pure-Python mode (Python-only types).

There's actually a 3rd mode, which is a coarse grained hybrid of the two, see: https://numba.readthedocs.io/en/stable/glossary.html#term-loop-lifting

If I'm correct then Numba uses Python C API to create JIT code which is compiled by Clang/LLVM.

I'm not sure to what this refers. Numba generates LLVM IR in all cases, there's no C lineage languages or code generate to an intermediate language involved.

It means that there is no reason why this C code can't mix pure-C and pure-Python types within neighboring expressions of single function.

That's close to the 3rd mode from above. It usually leads to hard to explain performance characteristics. Essentially, Numba can detect loops that can be compiled in "nopython" mode and leave everything else as calls to the CPython CPI. The problem with this is that it's hard to feed back to users what's happened and why. This is why the recommended use of the JIT compiler is to always specify @jit(nopython=True) so as to compiled without Python calls or to fail compiling. It gives much better performance and a more consistent user experience.

In Cython there is cdef which allows users to manually tell that this type is C-only. Would be great if Numba can mix C/Python types together too, like Cython does.

Many years ago Numba experimented with this, it's quite hard to do and IIRC you don't really get that much performance gain and it's again hard to explain what went on.

Also like in Cython Numba should allow to specify manually any type of local variable. As I know this is already allowed by passing locals = {'x': SomeType} to @jit decorator. So if C/Python types are mixed then this locals = argument should allow specifying both C and Python types.

As noted above, there are no C types, and no compilation mode that would permit this fine grained interleaving of specific nopython types with pyobject. It's just Numba's type system types from numba import types (they are all under that module).

Why is it needed? I often just add @njit/@jit decorator to ready-made Python functions and Numba perfectly JITs and optimizes on the fly. That is very nice. For example I specified @jit(nopython = False, objmode = True), it compiles pure object-mode code, but now I want to have some optimizations to this code without migrating fully to @njit. I want most of variables to stay Python objects as they are classes, but now I want to optimize some very computationally intense variables (like loop variables or some mathematical accumulators) and for that I want to specify that they should be C types like locals = {'k': 'c_uint64'}.

It sounds like refactoring your code to pull out the numerically intense calculations into separate functions would be the way to go. This means you can @njit them and also micro-benchmark them. It's also strongly recommended to not use locals unless you really need to and are content with the consequences, it can lead to strange things happening in the type inference mechanism. For most use cases, Numba can just work out what the type of things should be based on the input values.

Basically by this feature request I wanted to be able to have some intermediate states of code between pure Python and pure C.

In @njit it is currently allowed to have with blocks of objmode, which partially solves my request. But this with-blocks don't allow (without extra tricks) for Python objects to leave a block. There is as I remember numba_passthrough library that allows such tricks as Python objects leaving objmode blocks, but this solution is very un-intuitive, would be nice that Numba can support this natively. And even more - allow mix of any C/Python types within same code, like Cython allows.

I'm not sure that Numba is a great fit for your use case if you are wanting to spend a lot of time in object mode, or have fine grained interleaving of object mode/nopython mode. As noted above, refactoring the algorithm is probably the way to go.

Hope this helps.

[github-actions]

This issue is marked as stale as it has had no activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with any updates and confirm that this issue still needs to be addressed.

[polkovnikov]

@stuartarchibald As you marked this issue as question and Bot marked with stale, seems that this issue is not an issue, but more like clarification question, then I'm closing it. Can be reopened later if needed.

[stuartarchibald]

Great, thanks @polkovnikov. Also, for the future, https://numba.discourse.group/ is a good place to start discussions/ask questions. If something on there turns out to be an issue/feature request it can be transcribed into this tracker. Thanks!