Numba - 编译时间常数优化

Question

我有使用Numba进行JIT编辑的人工示例代码：

import numba
@numba.njit
def f(x, c):
    for i in range(3):
        if c == 0:
            x += 1
        elif c == 1:
            x *= 2
        else:
            assert False
    return x
print(f(3, 0), f(3, 1))

c如果是来自小的最终整数值集的某个常数，例如，如何优化它0 <= c <= 3？

不允许将if有关条件c移出循环，因为这将是一个明显的优化。因为它是一个人为的示例，所以在实际示例中，循环可能会很大，并且仅更改主体内部的 1-2 行即可重复循环两次是负担不起的。

基本上我想让c编译时保持不变。与C++一样，可以将编译时常量作为模板参数，并执行与此相关的所有代码优化。

所以我想告诉 Numba 参数c0、c1等是编译时常量，应该以这样的方式处理：每当使用不同的cNumba 值调用函数时，都应该创建和编译一个单独的函数实例并进行必要的优化，就像C++生成一样当传递新的模板参数集时，模板化函数的每次实例化。

换句话说，这些常量的可能值c应该包含在函数的签名中，据我了解，Numba 会为每个不同的签名生成单独的函数代码，就像 C++ 模板一样。

另一种需要优化的情况是，whenc不是函数的参数，而是循环开始之前某些复杂表达式计算的结果。在这种情况下，结果c也是来自一小组可能值。也许可以通过创建一个接受参数的子函数将这种情况转换为上面的情况c。

一般来说，我有一些嵌套循环，每个循环都有一些在循环运行之前计算的常量。然后在相应的循环内部，这些常量在条件内部使用if。如何强制 Numba 进行优化，将这些常量 if 条件移出这些循环体。

您有什么建议？