python - 如何将多个Python函数融合在一起

Question

我想将许多操作“融合”在一起。假设有 3 种可能的操作:

sq = lambda x: x**2
add = lambda x: x+3
mul = lambda x: x*5

我还有一系列操作:

ops = [add, sq, mul, sq]

然后我可以通过这些操作创建一个函数:

def generateF(ops):
    def inner(x):
        for op in ops:
            x = op(x)
        return x
    return inner
f = generateF(ops)
f(3) # returns 32400

fastF = lambda x: (5*(x+3)**2)**2

f 和 fastF 执行相同的操作，但 fastF 比 f 快 1.7-2 倍左右根据我的基准，这是有道理的。我的问题是，如何编写 generateF 函数来返回与 fastF 一样快的函数？这些操作仅限于基本操作，如 __add__、__mul__、__matmul__、__rrshift__ 等(基本上是 numeric operations ) 。 generateF 可以花您想要的时间，因为它会在到达热代码之前完成。

上下文是，这是我的库的一部分，因此我可以定义每个合法操作，从而确切地知道它们是什么。操作定义不是由最终用户随机提供给我们的(用户只能选择操作的顺序)，因此我们可以利用有关它们的所有外部知识。

这可能看起来像是过早的优化，但事实并非如此，因为 f 是热代码。放弃使用 C 不是一个选择，因为操作可能很复杂(想想 PyTorch 张量乘法)，并且 x 可以是任何类型。目前，我正在考虑修改 python 的字节码，但这非常不愉快，因为每个 Python 版本的字节码规范都会发生变化，所以我想在深入讨论该解决方案之前先在这里询问。

Answer 1

这是从给定函数的字节码合成新函数的非常hacky版本。基本技术是仅在第一个函数的开头保留 LOAD_FAST 操作码，并在最后一个函数的末尾之外去除 RETURN_VALUE 操作码。这使得在函数(最初的)之间的堆栈上操作的值。完成后，您不会进行任何函数调用。

import dis, inspect

sq = lambda x: x**2
add = lambda x: x+3
mul = lambda x: x*5

ops = [add, sq, mul, sq]

def synthF(ops):
    bytecode = bytearray()
    constants = []
    stacksize = 0
    for i, op in enumerate(ops):
        code = op.__code__
        # works only with functions having one argument and no other vars
        assert code.co_argcount == code.co_nlocals == 1
        assert not code.co_freevars
        stacksize = max(stacksize, code.co_stacksize)
        opcodes = bytearray(code.co_code)
        # starts with LOAD_FAST argument 0 (i.e. we're doing something with our arg)
        assert opcodes[0] == dis.opmap["LOAD_FAST"] and opcodes[1] == 0
        # ends with RETURN_VALUE
        assert opcodes[-2] == dis.opmap["RETURN_VALUE"] and opcodes[-1] == 0
        if bytecode:        # if this isn't our first function, our variable is already on the stock
            opcodes = opcodes[2:]
        # adjust LOAD_CONSTANT opcodes. each function can have constants,
        # but their indexes start at 0 in each function.  since we're
        # putting these into a single function we need to accumulate the
        # constants used in each function and adjust the indexes used in
        # the function's bytecode to access the value by its index in the
        # accumulated list.
        offset = 0
        if bytecode:
            while True:
                none = code.co_consts[0] is None
                offset = opcodes.find(dis.opmap["LOAD_CONST"], offset)
                if offset < 0:
                    break
                if not offset % 2 and (not none or opcodes[offset+1]):
                    opcodes[offset+1] += len(constants) - none
                offset += 2
            # first constant is always None. don't include multiple copies
            # (to be safe, we actually check that)
            constants.extend(code.co_consts[none:])
        else:
            assert code.co_consts[0] is None
            constants.extend(code.co_consts)
        # add our adjusted bytecode, cutting off the RETURN_VALUE opcode
        bytecode.extend(opcodes[:-2])
    bytecode.extend([dis.opmap["RETURN_VALUE"], 0])

    func = type(ops[0])(type(code)(1, 1, 0, 1, stacksize, inspect.CO_OPTIMIZED, bytes(bytecode),
                tuple(constants), (), ("x",), "<generated>", "<generated>", 0, b''),
                globals())

    return func

f = synthF(ops)
assert f(3) == 32400

恶心，还有很多警告(在注释中指出)，但它有效，并且应该与您的表达式一样快，因为它编译为几乎相同的字节码。它需要一些工作来支持连接更复杂的函数。