Inlining Python Function

Question

In a C program, inlining a function is a fairly intuitive optimization. If the inlined function's body is sufficiently small, you end up saving the jump to the function and creation of the stack frame, and you store the return value wherever the function's result would have been stored, jumping to the end of the inlined function's "body" rather than long-jumping to the return pointer.

I'm interested in doing the same thing in Python, converting two python functions into another valid python function where the first got "inlined" into the second. An ideal solution to this might look something like the following:

def g(x):
    return x ** 2

def f(y):
    return g(y + 3)

# ... Becomes ...

def inlined_f(y):
    return (y + 3) ** 2

Clearly, in a language as dynamic as Python, this isn't trivial to do automatically. The best generic solution I have come up with is to use dict to capture the arguments passed to the function, wrap the function body in a one-iteration for loop, use break to jump to the end of the function, and replace uses of arguments with indexes into the argument dictionary. The result looks something like the following:

def inlined_f(y):
    _g = dict(x=y + 3)
    for ____ in [None]:
        _g['return'] = _g['x'] ** 2
        break
    _g_return = _g.get('return', None)
    del _g
    return _g_return

I don't care that it's ugly, but I do care that it doesn't support returns from within loops. E.g.:

def g(x):
    for i in range(x + 1):
        if i == x:
            return i ** 2
    print("Woops, you shouldn't get here")

def inlined_f(y):
    _g = dict(x=y + 3)
    for ____ in [None]:
        for _g['i'] in range(_g['x'] + 1):
            if _g['i'] == _g['x']:
                _g['return'] _g['i'] ** 2
                break  # <-- Doesn't exit function, just innermost loop
        print("Woops, you shouldn't get here")
    _g_return = _g.get('return', None)
    del _g
    return _g_return

What approach could I take to this problem that avoids needing to use break to "jump" out of the inlined function's body? I'd also be open to an overall better, generic approach could I take to inline one Python function into another.

For reference, I'm working at the AST (abstract syntax tree) level, so using parsed Python code; clearly, outside of literal values, I don't know what value or type anything will have while performing this transformation. The resulting inlined function must behave identically to the original functions, and must support all features typically available when calling a function. Is this even possible in Python?

EDIT: I should clarify since I used the tag "optimization", that I'm not actually interested in a performance boost. The resulting code does not need to be faster, it just must not call the inlined function while still behaving identically. You can assume that both functions' source code is available as valid Python.

Thanks for the interesting question, but I am curious about your intended application for this transformation. — NPE
– NPE, Commented May 2, 2018 at 19:12
It may make more sense to do this transformation at the bytecode level. (If you're at the point where you're putting this level of effort into microoptimizations, though, using something like Cython will probably give you more bang for your buck.) — user2357112
– user2357112, Commented May 2, 2018 at 19:13
Hmm.. this doesn't seem like the normal definition of inlining (i.e. alpha-conversion + replacing the call with the code). You seem to be doing some sort of eager evaluation..? — thebjorn
– thebjorn, Commented May 2, 2018 at 19:19
@NPE Specifically I'm writing a library whose goal is to perform dynamic syntactic alteration to support other code generation libraries, such as numba.jit and tangent. Hence why I need the result to be a valid Python function. If I can dynamically produce equivalent Python functions that don't use certain code features, then I don't need to rely on those libraries to support those code features (such as function calls). — scnerd
– scnerd, Commented May 2, 2018 at 19:33
@scnerd: Numba currently doesn't support exception handling, but it does support function calls, so transforming function calls by using exception handling sounds unhelpful. — user2357112
– user2357112, Commented May 2, 2018 at 19:40

9000 · Accepted Answer · 2018-05-02 19:53:48Z

2

The only reasonable way on source level I see, simplified:

Parse the source into some AST (or just use the built-in AST).
Copy a subtree representing the function's body.
Rename the variables in the subtree, e.g. by adding an unique prefix.
At the call site, replace all passed arguments with assignments using the function's new variable names.
Remove the call and replace it with the function body you've prepared.
Serialize the AST back to source.

What poses real problems:

Generator functions; just don't inline them.
Returns from under try/finally that need to run the finally part. Might be pretty hard to rewrite correctly; imho, best left in-unlined.
Returns from under context managers that need to run the __exit__ parts. While not impossible, it's also tricky to rewrite preserving the semantics; likely also best left un-inlined.
Mid-function returns, especially from within multiple loop constructs. You might need to replace them with an extra variable and thread it into every condition of every while statement, and likely to add a conditional break to for statements. Again, not impossible but likely best left un-inlined.

answered May 2, 2018 at 19:53

9000

41k9 gold badges69 silver badges108 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

9000 Over a year ago

@thebjorn: Indeed! Though tail recursion can be replaced with a loop :)

thebjorn Over a year ago

Sure, and you could translate it all to CPS (continuation-passing-style) ;-) -- but whatever you try you'll very quickly run into the fact that Python is not friendly to simple forms of inlining. The PyPy inliner doesn't even try to do inlining at the source level, and those guys have worked on this for a while.

9000 Over a year ago

@thebjorn: Afaict Python here is an intermediate form: «Specifically I'm writing a library whose goal is to perform dynamic syntactic alteration to support other code generation libraries». I think simple forms of inlining (e.g. inlining a a function with a single return as the last statement) is doable. With something more advanced, complexity (and the chance of bugs) snowballs.

scnerd Over a year ago

For recursive functions, so far I've just been tracking the recursion depth, so you actually have _g_0, _g_1, etc., and the recursion limit can be explicitly capped (and should be a really low cap). Pretty sure generator functions are impossible to inline syntactically, so I've taken a "you better know what you're doing" approach and treated a call to a generator function f(x) to be equivalent to list(f(x)), which makes it inlinable.

scnerd Over a year ago

In cases of finally and context manager __exit__'s, they should get run properly as the return-exception cascades up past them, shouldn't they? That's one of the big points of those two paradigms is to be resilient to being exception'ed past, right?

|

Brendan Abel · Accepted Answer · 2018-05-02 19:19:02Z

1

Probably the closest analog to a return would be raising an Exception, which would work to pop out of nested loops to the top of the "inlined function".

class ReturnException(Exception):
    pass


g = dict(x=y + 3)
try:
    for j in some_loop:
        for _g['i'] in range(_g['x'] + 1):
            if _g['i'] == _g['x']:
                raise ReturnException(_g['i'] ** 2)
except ReturnException as e:
    _g['return'] = e.message
else:
    _g['return'] = None

I don't know how much overhead is associated with exceptions though or if that would be faster than simply calling the function.

answered May 2, 2018 at 19:19

Brendan Abel

38k15 gold badges101 silver badges126 bronze badges

3 Comments

scnerd Over a year ago

Oh, I like that approach. It'd require making sure that the exception type is available within the function body, but I can do that pretty easily. I'm going to give this a bit longer in case other people have suggestions, but I'd be happy to accept this answer.

user2357112 Over a year ago

@scnerd: You're going to have problems with except blocks inside the inlined function.

scnerd Over a year ago

@user2357112 True, if they are bare except: or except BaseException: blocks. I could inherit from BaseException to minimize the risk, but it's certainly not foolproof. Under proper coding practices, though, this exception approach is safer and more general than my original for/break approach

Collectives™ on Stack Overflow

Inlining Python Function

2 Answers 2

6 Comments

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

6 Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related