Question: Are you interested to support automatic loop unrolling into batches of iterations?

[joshring]

Given a trivial code example:

def some_fn():
    for i in range(100):
        result += x * i
    return result

unrolled in batches of say 10 iterations:

def some_fn():
    for i in range(0, 100, 10):
        result += x * i
        result += x * (i+1)
        result += x * (i+2)
        result += x * (i+3)
        result += x * (i+4)
        result += x * (i+5)
        result += x * (i+6)
        result += x * (i+7)
        result += x * (i+8)
        result += x * (i+9)
    return result

• I know it's possible to do such a transform by hand with an unrolled inner loop and a normal outer loop and manually adjusting the loop iterations.
• I figured it would make it a less manual process (and less error-prone) if there was an option to automate this transform taking a parameter for the number of iterations to unroll at a time, and a corresponding decrease in the loop iterations.

[leonardt]

Yes this seems like a useful feature. Would you be interested in adding it? We would be happy to help guide you on how to do it. If not, we will look into adding it when we get a chance.

[joshring]

With the caveat that I am quite new to looking at ASTs, I am interested to have a look at it sure.

[leonardt]

No problem :) I think this would be the perfect task for learning more about ASTs as well as this library. It's not too complicated, but requires some knowledge of the nodes and can leverage some of the reusable tools we've built here.

First, I'll mention this doc as a useful guide to the python ast: https://greentreesnakes.readthedocs.io/en/latest/nodes.html
You'll probably want it handy as you're learning the node names and their attributes.

For ast_tools, we use a specific pass pattern to make it easy to write and compose passes. To write a new pass, you'll subclass the Pass class and define a rewrite method. Here's an example for the current loop unrolled:
https://github.com/leonardt/ast_tools/blob/master/ast_tools/passes/loop_unroll.py#L9-L14

For this pass, you won't need to use themetadata parameters (at least for now) and so you can simply just return that unchaged/unused. You'll want to take tree and env parameter, modify it to do your loop unrolling, and return the modified tree.

Here's the core logic for the current loop unroller: https://github.com/leonardt/ast_tools/blob/master/ast_tools/transformers/loop_unroller.py

It defines a visit_For method which will be invoked on For loop nodes inside the AST. It checks if the loop iter can be evaluated to a constant value and that value is an instance of the unroll object. If true, it unrolls the loop by building a list of statements to return. For each iteration to be unrolled, it appends the contents of the For loop body to the return list after replacing any instances of the loop variable with the current iteration value. This is done using the symbol table https://github.com/leonardt/ast_tools/blob/master/ast_tools/transformers/loop_unroller.py#L33 passed to the replace_symbol function.

Here, instead of replacing the loop variable with a constant (ast.Num node), you can instead construct an ast node for the desired expression (e.g. i + 1).

I think this could possibly be folded into the same logic as the current unroll pass. Perhaps we can call this macro unroll_by (referring to the constant factor to unroll_by, right now unroll tries to unroll the entire loop). So, the syntax could look something like
for i in ast_tools.macro.unroll_by(range(16), 4).

In this case, we'll probably want to restrict the arguments to be a simple range with a step of 1 for now, and you could rewrite the loop by using the step argument to range, e.g. for i in range(16, step=4). You could alternatively divide the end of the range by the unroll factor and multiply the iteration variable with the factor.

The unroll macro is defined here: https://github.com/leonardt/ast_tools/blob/master/ast_tools/macros.py#L1-L6 you can define a similar macro but add an argument for the unroll factor. This would be the object you're looking for when inspecting/evaluting the for loop node iterator.

let me know if you have any questions or run into trouble, more than happy to help!

[joshring]

Thanks for that write up it was very helpful. 👍

A prototype is coming along, fairly early days but some experiments with basic ast manipulation look promising so far :)
Have replicated the previous approach's results using a batch style calculation, i'll now put that for loop into the ast as opposed to using it to fill the node's body, which was nice for debugging