Multiple loop unrolling passes?

[jack-melchert]

I was wondering if its possible to support multiple passes of loop unrolling.
This is the type of nested loop I want to unroll:

def test_nested_unroll():
    z = 3
    @end_rewrite()
    @loop_unroll()
    @begin_rewrite()
    def foo():
        for i in ast_tools.macros.unroll(range(z)):
            for x in ast_tools.macros.unroll(range(i)):
                print(i, x)

    print(inspect.getsource(foo))

Currently, this code produces:

def foo():
    for x in ast_tools.macros.unroll(range(0)):
        print(0, x)
    for x in ast_tools.macros.unroll(range(1)):
        print(1, x)
    for x in ast_tools.macros.unroll(range(2)):
        print(2, x)

While it would be helpful if it could unroll it again to:

def foo():
    print(1, 0)
    print(2, 0)
    print(2, 1)

Its possible to do what I want using the if_inline but, as you can imagine, the code gets very long if you have large loops:

def test_nested_unroll():
    z = 3
    @end_rewrite()
    @loop_unroll()
    @if_inline()
    @begin_rewrite()
    def foo():
        if inline(z == 1):
            for x in ast_tools.macros.unroll(range(0)):
                print(0, x)
        if inline(z == 2):
            for x in ast_tools.macros.unroll(range(0)):
                print(0, x)
            for x in ast_tools.macros.unroll(range(1)):
                print(1, x)
        if inline(z == 3):
            for x in ast_tools.macros.unroll(range(0)):
                print(0, x)
            for x in ast_tools.macros.unroll(range(1)):
                print(1, x)
            for x in ast_tools.macros.unroll(range(2)):
                print(2, x)

    print(inspect.getsource(foo))

[rdaly525]

Does it work if you apply the loop_unroll decorator twice? That could probably unblock you.

[jack-melchert]

Yep, that worked. Thanks for the help!

[rdaly525]

Lets keep this open because I think It would be useful to think about how we want marcros to evaluate. @cdonovick @leonardt, thoughts?

[cdonovick]

Yeah macro evaluation order is tricky. One possible solution is that the loop_unroll keeps calling itself until it the AST converges. But what if I need to unroll, then inline, then unroll? Maybe we should have a do_macros pass or something that just keeps calling loop_unroll / if_inline / <new macros> until convergences?

[cdonovick]

@leonardt @phanrahan @willcrichton Thoughts?

[rdaly525]

Something like this, where macro1 and macro2 would keep getting applied until it converges?

@end_rewrite()
@while_not_converge()
@macro2()
@macro1()
@do_macros()
@begin_rewrte()
def foo(...)

[willcrichton]

Is there a reason in the first place why the user has to write @loop_unroll? If I explicitly annotate the body of the function with ast_tools.macros.unroll calls, then it's clear I want those to be unrolled.

[willcrichton]

If you still need precise control over the order of passes, then you could also do something like

@loop_inline('a')
def foo():
  for x in name('a', range(0)):
    ...

This is in the same vein of how Pyro deals with random variables in its Python-embedded probabilistic programming language.

[leonardt]

That's an interesting thought @willcrichton , I think the use of decorators for the pass was mainly as an entry point to the function as well as ordering as you mentioned. The use of the ast_tools.macros.unroll function was a guard for the pass to determine whether the loop was to be unrolled or not.

But, it seems like we could have a generic @macro function that has a standard collection of passes that will be applied to handle these various markers. Then, we could do something that inspects the tree to collect any macro markers, and run the corresponding pass, until no markers are left in the tree (handling the case where one macro produces invocations of other macros).

(p.s. I think this is the do_macros pass that caleb suggested, it also seems reasonable to provide the passes you'd like to run, however, it could be that one pass depends on another pass, which in this case the user would need to know these dependencies, where if we could somehow automatically infer them via tree inspection, the user could simply include the top level macro/passes that they'd like to use and the dependent passes would be managed implicitly)

[cdonovick]

@willcrichton @loop_unroll invokes the the loop unrolling pass. The loop unrolling pass finds for loops where (with some abuse of notation) isinstance(for.iter, ast_tools.macros.unroll) and unrolls the contained iterator*.

* currently only supports iterators which yield int but I have plans to generalize it to any elaboration time constant iterator.

[cdonovick]

@leonardt I think we should think about pass ordering and infrastructure for managing it. We already have some dependencies inline_if and loop_unroll must be run before ssa.

Maybe we could describe the types of AST nodes that a pass operates on and the type nodes that pass eliminates, and type of nodes that a pass can't except, and from these constraints generate an ordering? For example loop_unroll eliminates For nodes and ssa can't accept For nodes. So loop_unroll must come before ssa. I have no idea how this usually done in compilers though.

[rdaly525]

I can comment that in LLVM, you manually specify other passes as dependencies for the current pass. This is not quite the same as what you want though. Not sure about the "if you want to run this other pass, it needs to be run before current pass"

[willcrichton]

I think in most compilers, ssa wouldn't be a pass in the same sense as loop_unroll. ssa translates between two languages, which would be represented by distinct IRs, while loop_unroll would be intra-IR. Here, you're using the general Python AST to hold each IR.