numba jitted函数中的集合的正确签名是什么?

时间:2019-03-22 13:35:47

标签: python set jit signature numba

如果我理解正确,则可以通过添加签名来提高numba函数的性能。示例:

@njit(int32(int32, int32))
def f(x, y):
    # A somewhat trivial example
    return x + y

现在我有需要两个集合的函数。什么是正确的签名?

@njit(int32(set(int32), set(int32)))
def f(set_1, set_2):
    # A somewhat trivial example
    return x

我认为签名(int32(set(int32), set(int32)))可能是正确的,但是没有任何反应。 print(numba.typeof(set_1))返回reflected set(int32)

1 个答案:

答案 0 :(得分:2)

  

如果我理解正确,则可以通过添加签名来提高numba函数的性能。

那是错误的-或仅部分正确。使用签名numba可以提前编译函数,而不是使用这些参数首次调用时。第一次通话后,两者应同样快。在某些情况下,如果没有签名,该功能可能会稍微快一些(尤其是在numba可以使用输入的数组对齐方式的数组中)。

  

现在我有需要两个集合的函数。什么是正确的签名?

对于包含整数的Python集,正确的签名是:

numba.types.Set(numba.int64, reflected=True)

因此,具有两个集合(并返回一个集合)的函数的签名为:

import numba as nb

reflected_int_set = nb.types.Set(nb.int64, reflected=True)

@nb.njit(reflected_int_set(reflected_int_set, reflected_int_set))
def f(set_1, set_2):
    return set_1

>>> f({1,2,3}, {3,4,5})
{1, 2, 3}

但是,因为它(很有可能)并不能提高性能,所以我根本不会理会签名。

另外请注意:numba会将Python集转换为内部的numba集,因此将Python set传递给numba函数,或将set从numba函数返回到Python上下文复制整套。在大多数情况下,开销比numba提供的潜在加速要重要得多。

以我的经验,仅将numba函数严格限制在set s和list s中使用numba才有意义。因此,如果将它们用作参数或将其返回(返回给非numba函数/上下文),则必须测量性能并检查是否确实获得了加速。