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

Question

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

---

我在python中做一些移位，想用numba加快速度。为此，我有很多常量整数值，必须以一种易于理解的方式进行处理。我想将它们组合在一起，成为枚举式对象，使所有常量在一个名称空间中，可以通过attribute-get运算符进行访问。当然，我也想让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，这至少表明它可以有效地工作。

有点破解，但希望效果很好！