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

Question

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

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

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

不幸的是，当我查看注释时，numba 似乎不理解这些值是常量，并且它总是回退到 python 对象上的慢速对象访问。这就是注释所说的：

#   .2 = global(SomeConstantsContainer: <class 'constants.SomeConstantContainer'>)  :: pyobject
#   .3 = getattr(attr=SOME_VARIABLE, value=.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

这给出：

    #   .3 = global(SomeConstantsContainer: <enum 'SomeConstantsContainer'>)  :: pyobject
    #   .4 = getattr(attr=SOME_OTHER_CONSTANT, value=.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()

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

当我尝试使用一个非常基本的示例并使用 inspect_types 时，我得到

[=11=].1 = freevar(A: 10.0)  :: float64

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

有点乱七八糟，但希望效果很好！

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

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

python

enums

constants

numba