使用 numba 在容器中处理多个常量的最佳方法？

Question

编辑：这个问题已经严重过时了！ numba 现在支持开箱即用的Enum 和namedtuple，它们都为分组常量提供了合理的解决方案。

---

我正在 python 中进行一些位移，并希望使用 numba 加速它。为此，我有很多常量整数值，我必须以一种可能易于阅读的方式处理它们。我想将它们组合在一起成为类似枚举的对象，将所有常量都放在一个命名空间中，可以通过属性获取运算符进行访问。当然我也希望 numba 了解那里发生的事情，以便它可以通过 jit 编译保持高速。我的第一次也是最天真的尝试是这样的：

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

不幸的是，当我查看注释时，看起来 numba 不理解值是恒定的，并且它总是退回到对 python 对象的缓慢对象访问。这就是注释所说的：

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

我知道我总是可以退回到这样的事情：

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

在这种情况下，jit 编译器 a) 不需要查找容器的属性 b) 肯定知道它必须处理一个纯整数。这样写简直丑到令人难以置信。我可以忍受将所有常量显式标记为整数，或者让容器这样做。尽管如此，为了清楚起见，我真的很想将常量分组在容器中，并且 jit 编译的版本理解语法，而不是浪费时间在每次使用常量时进行一些缓慢的 python 属性查找。任何更好的想法，如何使第二种方法更像第一种方法，但保持高执行速度？是否有一些 numba 理解的枚举容器，我只是错过了？

编辑： 同样使用新的enum 容器也无济于事：

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

这给出了：

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

Answer 1

另一种方法，但仍然具有包含变量的优点的方法是在包装函数中使用捕获的变量：

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()

您显然可以将其与使用类作为命名空间结合起来，并在外部函数中解压内容。

当我用一个非常基本的例子尝试这个并使用 inspect_types 我得到

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

（其中A 只是我的常数）。我怀疑这就是你想要的。我确实尝试查看生成的汇编程序 (os.environ['NUMBA_DUMP_ASSEMBLY']='1')，但我在汇编程序方面做得不够好，无法挑选出相关行并确认它是否符合您的要求。但是 - 它确实使用nopython=True 编译，这至少表明它正在有效地工作。

有点小技巧，但希望能很好地工作！