Best way to handle multiple constants in a container with numba?

Question

EDIT: This question is heavily outdated! numba now supports Enum and namedtuple out of the box, which both provide a reasonable solution for grouping constants.

---

I'm doing some bitshifting in python and want to speed it up with numba. For that, I have lots of constant integer values, that I have to handle in a possibly well readable manner. I would like to group them together to enum-like objects, having all the constants within one namespace, accessible with the attribute-get operator. And of course I'd also like, that numba understands what's going on there, so that it can maintain high speeds with jit compilation. My first and most naive try on that looked like that:

class SomeConstantsContainer:
    SOME_NAME = 0x1
    SOME_OTHER_CONSTANT = 0x2
    AND_ANOTHER_CONSTANT = 0x4

Unfortunately, when I look at the annotation, it looks like numba is not understanding that the values are constant, and it always falls back to slow object access on python objects. This is what the annotation says about it:

#   $29.2 = global(SomeConstantsContainer: <class 'constants.SomeConstantContainer'>)  :: pyobject
#   $29.3 = getattr(attr=SOME_VARIABLE, value=$29.2)  :: pyobject

I know that I always could fall back to something like this:

from numpy import np
SOME_STUPID_CONSTANT = np.int64(0x1)
ANOTHER_STUPID_CONSTANT = np.int64(0x2)

In that case the jit compiler a) does not need to look up the attribute of the container and b) knows for sure, that it has to deal with a plain integer. It's just incredibly ugly to write like that. I could live with marking all constants explicitly as integers, or let the container do that. Nevertheless, I really would like to group the constants in containers for clarity, and the jit compiled version understanding the syntax and not to waste time on some slow python attribute lookup for every use of the constants. Any better ideas, how to make the second approach more like the first approach, but keeping high execution speeds? Is there some enum container, that numba understands, which I just missed?

Edit: Also using the new enum container is not of help:

@enum.unique
class SomeConstantsContainer(enum.IntEnum):
    SOME_NAME = 0x1
    SOME_OTHER_CONSTANT = 0x2
    AND_ANOTHER_CONSTANT = 0x4

This gives:

    #   $42.3 = global(SomeConstantsContainer: <enum 'SomeConstantsContainer'>)  :: pyobject
    #   $42.4 = getattr(attr=SOME_OTHER_CONSTANT, value=$42.3)  :: pyobject

Answer 1

A different approach, but one that still has the advantage of containing the variables would be to use captured variables in a wrapped function:

def make_numba_functions():
    SOME_NAME = 0x1
    SOME_OTHER_CONSTANT = 0x2
    AND_ANOTHER_CONSTANT = 0x4

    @jit
    def f1(x,y):
        useful code goes here

    @jit
    def f2(x,y,z):
        some more useful code goes here

    return f1, f2

f1,f2 = make_numba_functions()

You could obviously combine this with using a class as a namespace, and unpack the contents inside the outer function.

When I try this with a very basic example and use inspect_types I get

$0.1 = freevar(A: 10.0)  :: float64

(where A is just my constant). I suspect this is what you want. I did try to look at the generated assembler (os.environ['NUMBA_DUMP_ASSEMBLY']='1') but I'm not good enough at assembler to pick out the relevant line and confirm it does what you want. However - it does compile with nopython=True, which at least suggests it's working efficiently.

A bit of a hack, but one that works pretty well hopefully!