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

It's very unclear what you're asking. Are you trying to get these variables to be represented as `const` values in the emitted C/machine code? If so, why? You cannot verify const-correctness at the Python level, so this might not make sense. That is, if these exist as class attributes, as in your example code, then anyone can modify them willy-nilly at any time in Python *before* and JIT-compiled function call is dispatched. Doing const-correctness at the C-level would not gain you anything, that I'm aware of (but would be very interested in counter-examples). — ely, Feb 20 '15 at 17:31
I want them to be recognized and handled not as python objects, but as integer constants. I also want the jit compiler to be able to scratch away the attribute lookup on the container object. — Michael, Feb 20 '15 at 17:44
NOTE: This question is heavily outdated now. numba now supports `Enum`s as well as `namedtuple`s, which both offer more elegant ways to answer this question. — Michael, Jun 09 '16 at 08:43

DavidW · Accepted Answer · 2015-03-14T17:10:23.750

3

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!

edited Mar 14 '15 at 17:10

answered Mar 14 '15 at 16:08

DavidW

29,336
6
55
86

Interesting approach! This should be generalizable into decorators for every constants set. Indeed hacky :) – Michael Mar 15 '15 at 09:08
I did consider decorators, but couldn't see how exactly to fo it. I'd be interested to see an implementation... – DavidW Mar 15 '15 at 09:37
I had thought about things like this, but in the end it didn't work out as I had thought: http://stackoverflow.com/questions/392349/modify-bound-variables-of-a-closure-in-python – Michael Mar 18 '15 at 15:18
One option that might work is to use Byteplay (https://wiki.python.org/moin/ByteplayDoc) to modify the function bytecode. Finding and replacing attribute lookups from your class of constants would probably not be too hard (they'd all follow the same pattern, I guess). Since (I think) Numba reads the bytecode, it would probably work. It does rely on what looks like an unmaintained library though. – DavidW Mar 18 '15 at 20:11

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

1 Answers1