加速用户定义的函数

Question

我有一个模拟，其中最终用户可以提供任意多个函数，然后在最内层循环中调用这些函数。类似于：

class Simulation:

    def __init__(self):
        self.rates []
        self.amount = 1

    def add(self, rate):
        self.rates.append(rate)

    def run(self, maxtime):
        for t in range(0, maxtime):
            for rate in self.rates:
                self.amount *= rate(t)

def rate(t):
    return t**2

simulation = Simulation()

simulation.add(rate)
simulation.run(100000)

作为一个 python 循环，这非常慢，但我无法使用我的正常方法来加速循环。

因为函数是用户定义的，所以我不能"numpyfy"最内层的调用（重写使得最内层的工作由优化的 numpy 代码完成）。

我首先尝试了numba，但是numba不允许将函数传递给其他函数，即使这些函数也是numba编译的。它可以用闭包，但是因为一开始不知道有多少函数，所以觉得用不上。关闭函数列表失败：

@numba.jit(nopython=True)
def a()
    return 1

@numba.jit(nopython=True)
def b()
    return 2

fs = [a, b]

@numba.jit(nopython=True)
def c()
    total = 0
    for f in fs:
        total += f()
    return total

c()

失败并出现错误：

[...]
  File "/home/syrn/.local/lib/python3.6/site-packages/numba/types/containers.py", line 348, in is_precise
    return self.dtype.is_precise()
numba.errors.InternalError: 'NoneType' object has no attribute 'is_precise' 
[1] During: typing of intrinsic-call at <stdin> (4)

我找不到来源，但我认为 numba 的文档在某处声明这不是错误，但预计不会起作用。

像下面这样的东西可能会解决从列表中调用函数的问题，但似乎是个坏主意：

def run(self, maxtime):
    len_rates = len(rates)
    f1 = rates[0]
    if len_rates >= 1:
        f2 = rates[1]
    if len_rates >= 2:
        f3 = rates[2]
    #[... repeat until some arbitrary limit]
    @numba.jit(nopython=True)
    def inner(amount):
        for t in range(0, maxtime)
            amount *= f1(t)
            if len_rates >= 1:
                amount *= f2(t)
            if len_rates >= 2:
                amount *= f3(t)
            #[... repeat until the same arbitrary limit]
        return amount

    self.amount = inner(self.amount)

我想也可以进行一些字节码破解：使用 numba 编译函数，将包含函数名称的字符串列表传递给 inner，执行类似 call(func_name) 的操作，然后然后重写字节码，使其变成func_name(t)。

对于 cython 只编译循环和乘法可能会加速一点，但如果用户定义的函数仍然 python 只是调用 python 函数可能仍然很慢（虽然我没有'配置文件）。我并没有真正找到关于 cython 的 "dynamically compiling" 函数的太多信息，但我想我需要以某种方式向用户提供的函数添加一些类型信息，这似乎..很难。

有没有什么好的方法可以加速用户定义函数的循环，而无需从中解析和生成代码？

Answer 1

我不认为你可以加速用户的功能 - 最终编写高效代码是用户的责任。您可以做的，是提供一种以高效方式与您的程序交互的可能性，而无需支付开销。

您可以使用 Cython，如果用户也是使用 cython 的游戏，与纯 python-solution.

相比，你们俩都可以实现大约 100 的加速。

作为基准，我稍微更改了您的示例：函数 rate 做了更多工作。

class Simulation:

    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, maxtime):
        for t in range(0, maxtime):
            for rate in self.rates:
                self.amount += rate(t)

def rate(t):
    return t*t*t+2*t

产量：

>>> simulation=Simulation([rate])
>>> %timeit simulation.run(10**5)
43.3 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

我们可以使用 cython 来加快速度，首先是你的 run 函数：

%%cython
cdef class Simulation:
    cdef int amount
    cdef list rates
    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        for t in range(maxtime):
            for rate in self.rates:
                self.amount *= rate(t)

这几乎给了我们因素 2：

>>> %timeit simulation.run(10**5)
23.2 ms ± 158 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

用户还可以使用 Cython 来 speed-up 他的计算：

%%cython
def rate(int t):
  return t*t*t+2*t

>>> %timeit simulation.run(10**5)
7.08 ms ± 145 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

使用Cython已经给了我们speed-up 6，现在的bottle-neck是什么？我们仍然对 polymorphism/dispatch 使用 python，这非常昂贵，因为为了使用它，必须创建 Python-objects（即这里的 Python-integers）。我们可以用 Cython 做得更好吗？是的，如果我们在编译时为传递给 run 的函数定义一个接口：

%%cython   
cdef class FunInterface:
   cpdef int calc(self, int t):
      pass

cdef class Simulation:
    cdef int amount
    cdef list rates

    def __init__(self, rates):
        self.rates=list(rates)
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        cdef FunInterface f
        for t in range(maxtime):
            for f in self.rates:
                self.amount *= f.calc(t)

cdef class  Rate(FunInterface):
    cpdef int calc(self, int t):
        return t*t*t+2*t

这产生额外的 speed-up 7:

 simulation=Simulation([Rate()])
 >>>%timeit simulation.run(10**5)
 1.03 ms ± 20.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

上面代码中最重要的部分是行：

self.amount *= f.calc(t)

它不再需要 python 进行分派，而是使用与 c++ 中的虚函数非常相似的机制。这种 c++ 方法的开销只有非常小的一个 indirection/look-up。这也意味着，函数的结果和参数都必须转换为 Python-objects。为此，Rate 必须是 cpdef-function，您可以查看了解更多详细信息，继承如何为 cpdef-function 工作。

bottle-neck现在是for f in self.rates行，因为每一步我们还要做很多python-interaction。这是一个示例，如果我们可以对此进行改进：

%%cython
.....
cdef class Simulation:
    cdef int amount
    cdef FunInterface f  #just one function, no list

    def __init__(self, fun):
        self.f=fun
        self.amount = 1

    def run(self, int maxtime):
        cdef int t
        for t in range(maxtime):
                self.amount *= self.f.calc(t)

...

 >>>  simulation=Simulation(Rate())
 >>> %timeit simulation.run(10**5)
 408 µs ± 1.41 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

另一个因素 2，但您可以决定是否需要更复杂的代码来存储 FunInterface 对象列表而不需要 python-interaction 是否真的值得。

加速用户定义的函数

Speeding up user-defined functions

python

performance

cython

numba