[hail] tail-recursive loops

[catoverdrive]

I don't think I'm quite satisfied with this implementation, although I think that it will do the things we need it to do, for the most part. It mostly adheres to the structure that @chrisvittal was implementing in #5228, although I had some questions/notes about implementation/interface and would love some input:

• I'd like to implement (in python) some sort of while loop construct, but I'm not sure what it looks like. I think I'd like to think of that as our primary loop construct since the general recursive loop function can be confusing to start with in terms of what's allowed (what is tail-recursive? what is non-tail-recursive?) if you're just trying to implement some convergence criteria. Some initial thoughts on interface:

1:
loop = hl.WhileBuilder(i=0, x=0)
loop.cond(loop.i < 10)
    .update(x = loop.x + i, 
            i = loop.i + 1)
    .result(loop.x)

2: 
loop = hl.while_loop(
    lambda i, x: i < 10, 
    lambda i, x: (x + i, i + 1), 
    lambda i, x: x, 
    0, 0)

3:
loop = hl.while_loop(
    lambda i, x: 
    hl.loop.cond(i < 10)
           .update(i + 1, x + i)
           .result(x).
    0, 0)

but I'm not sure I really like any of them.

• the scoping for Recur is difficult to check and enforce.
  • It's difficult to check if a TailLoop is invalidly attempting to recur a function from a different loop (in a nested environment), except by the type signature. I think I could give each loop a name so that Recur unambiguously refers to a loop defined in the surrounding scope, mostly treating Recur as something of a function reference.
• The code generation is rather inconsistent, since the Recur node technically has the same type as the return type of the function, but the code generated needs to be a jump node with no actual value (which makes the generated EmitTriplet look a lot like it has type TVoid!)
  • @patrick-schultz proposed a similar design, but with two additional types to make the difference between the type of the Recur concept and the actual return type more explicit. I have not yet implemented it because I think it might make this python interface more difficult to support, and I rather like its simplicity.
• @patrick-schultz and I were talking about rewriting loops in the stream interface.
  • We can only emit calls to recur in the same method that the original loop is defined in, because we are using jumps to implement them. This means we need to know that we are not going to wrap any calls that contain Recur in a method. I don't believe there are any cases where we wrap If conditions or Let bodies in methods, so this is fine for now, but we're not enforcing it in any way. I believe that the iteration for the stream codegen stuff will always be in the same method, so we wouldn't have to deal with this specially there.
  • I've implemented a pretty simple version of here, as part of the Streamify pass. I believe it should handle all the valid tail-recursive cases, but I don't think we want to use it right now since it'll always allocate (as opposed to not allocating if all state is primitive). We could potentially revisit this once we can allow stream elements like this to basically store their elements in primitive fields, if that prevents allocation. You can take a look here: https://github.com/catoverdrive/hail/compare/loops...catoverdrive:loops-as-stream?expand=1

cc @cseed @patrick-schultz @chrisvittal

[cseed]

Awesome! I will take a closer look, but here's my quick reaction to the interface. Was I was hoping to write was:

hl.loop(
  lambda i, x:
    hl.cond(i < 10, hl.recur(i + 1, x + i), x),
  0, 0)

This is basically modeled after letrec in lisp/scheme/ml which would look something like:

  (letrec (f i x)
    (if (< i 10)
      (f (+ i 1) (+ x i))
      x)
    (f 0 0))

Another option is to get rid of recur and make the binding of the loop more explicit with something like:

hl.loop(
  lambda f, i, x:
    hl.cond(i < 10, f(i + 1, x + i), x),
  0, 0)

where the first argument to the lambda is the loop itself.

I think main problem with your proposals (except maybe 3) is that it assumes too much about the structure of the loop: namely, it embeds the exit condition into the structure of the loop. The loop may have several backwards calls or exit points (e.g. in a case or if tree) and there maybe shared and conditional work that happens inside that tree (and even nested loops!)

[catoverdrive]

@cseed I think my point is that I want an interface that assumes the structure of the loop. I'm happy to implement other interfaces as well--the one that I've currently got in this PR is basically your second suggestion--but I think it would also be nice to have a while loop construct that looks syntactically more similar to what you'd expect a while loop to look like in python, if that makes sense?

[catoverdrive]

I do agree that I think all of my proposed interfaces are missing the ability to do things like nest loops, which is part of the reason I'm not sold on any of them.

[cseed]

So I'm going to insist on the classical loop interface I described above, since it is strictly more powerful than the interfaces you've proposed. I don't have a strong feeling if you want to also add a Python-inspired while loop (although I personally would find the similarities misleading given the required differences, I understand others might feel differently). Your while loop should be naturally implementable in terms of mine, so I also suggest we focus on that first.

Giving each loop a name seems natural. Apart from the wrapping issue (the greatest existential threat our generation faces) I don't see any problem calling an outer loop from an inner loop.

Is Patrick's proposal for extra types written up anywhere? I don't like the idea of complicating the type hierarchy for internal bookkeeping like this.

So I'm going to remark that in the code generator it is often natural to build data structures to aid the organization of the code generator, and those data structures need not need to be types/IRs. Given that Recur has to be in tail position, and you know exactly when you're existing the loop (branches that don't contain recur nodes). So the compilation looks like:

  set initial loop variables
  # fall through into loop
Lloop:
  ...
  # recur
  loop variables = new values
  goto Lloop
Lan_exit_branch:
  result = compile(branch)
  goto Lafter
Lanother_exit_branch:
  result = compile(other_branch)
  goto Lafter
Lafter:
  use result ...

What I would do is "peel" off the ifs and lets (anything else?) that can sit in tail position and build a separate data structure for those nodes which I then traverse to emit the above code.

Using the stream interface seems wrong to me also. What's the type of the stream the loop turns into? Since loops carry multiple values (by design), memory allocating these to create a tuple stream is going to be a performance non-starter.

I'll comment more once I've looked over the code.

[danking]

FWIW, I'm also strongly in favor of: hl.loop(lambda loop, *parameters: ..., *arguments).

[patrick-schultz]

So I'm going to insist on the classical loop interface I described above, since it is strictly more powerful than the interfaces you've proposed.

I agree that the tail-recursion interface seems like the right primitive to expose in python, on top of which we could implement convenience methods for building while/for loops if we decide it's worth it.

Giving each loop a name seems natural. Apart from the wrapping issue (the greatest existential threat our generation faces) I don't see any problem calling an outer loop from an inner loop.

Also agree. This will require either adding another context of loops/continuations in the environment (valid places to jump to, and their argument types), or keeping them in the normal value context by adding a new continuation type.

Is Patrick's proposal for extra types written up anywhere?

My proposal has two main differences. In

hl.loop(
  lambda i, x:
    hl.cond(i < 10, hl.recur(i + 1, x + i), x),
  0, 0)

the point that jumps back to the top of the loop is explicit, but the point that jumps out of the loop is not. I suggested making this something like

hl.loop(
  lambda i, x:
    hl.cond(i < 10, hl.recur(i + 1, x + i), hl.break(x)),
  0, 0)

or, if we're giving names to loops, it might be simpler to pass the break and recur functions to the lambda:

hl.loop(
  lambda sum, ret, i, x:
    hl.cond(i < 10, sum(i + 1, x + i), ret(x)),
  0, 0)

The second difference is in the typing. In this PR, the hl.recur expression is given the type of the entire loop. I would add a single new type Bottom, and give all expressions which jump (both the recur and the break expressions) the type Bottom. Bottom is the empty type, so there can be no closed expressions of type Bottom.

In the type checker, Bottom is only allowed to appear in tail positions, and for If, we keep the rule that both branches must have the same type, so either both branches are Bottom or neither are. This keeps the semantics simple: an if statement either makes a value or it jumps away, there's no confusing mix.

One nice property of this setup is that if an expression has a non-bottom type, then it is guaranteed not to jump away from itself (it may jump internally), so it is safe to method-wrap.

This also make codegen very simple, and @iitalics has already implemented it! See JoinPoint and JoinPoint.CallCC.

In the IR, I don't think this requires much change. If we're already adding a continuation context (as mentioned above), then TailLoop just needs to bind both a recur and a break continuation, where recur takes the loop variable types, and break takes the loop result type. Then Recur can be replaced by a Jump node which calls (jumps to) a continuation in context.

There's also a middle ground where we make break continuations explicit in the IR, but we want to keep the scheme-like interface in python. Then the pass @cseed described for inferring where the loop exits are would just go in python instead of the emitter.

Using the stream interface seems wrong to me also.

When I mentioned this, I was thinking we could reuse the region management logic from the stream emitter for loops, but I've since changed my mind. I think loops will have hard region management no matter what.

What's the type of the stream the loop turns into? Since loops carry multiple values (by design), memory allocating these to create a tuple stream is going to be a performance non-starter.

As a side note, @iitalics stream emitter can handle streams of multiple values fine. Effectively, you can make a Stream[(Code[A], Code[B], Code[C])].

[cseed]

This proposal mounts to programming with explicit continuations. It doesn't increase the expressiveness of the loop construct that I can see. Our users are reluctant enough to learn functional programming, I think continuations is one step too far for the user interface. Internally, I don't care as much, although personally I would prefer to code up my solution. @catoverdrive's doing the work, so I'll let them decide.

As a side note, @iitalics stream emitter

Ah, I thought @catoverdrive was referring to IR level streams.

[catoverdrive]

ok, great. @cseed I've essentially already implemented the thing that you've described (although I need to actually duplicate the IR nodes in python for the loop function to work, and add some tests), although I haven't peeled off the ifs and lets to emit separately (and I don't know if we need to? seems to work fine with the current code generator); we check that all recurs are in tail position when we typecheck the IR. I will clean up the python and give the loops names and then assign this to someone.

[tpoterba]

This was a very legible PR, nice work.

[tpoterba]

can we remove this FIXME?

[tpoterba]

can we make this a little less "expression"-y?

[tpoterba]

does pytest not generate a nice error here?

[catoverdrive]

it didn't for me. Don't know if that was some sort of setting thing.

[tpoterba]

I don't really care if we implement new nodes in Interpret. It's not long for the world.

[catoverdrive]

it makes the tests run better for now, so I just threw together a quick thing (we still test against the interpret path for IRSuite by default)

[tpoterba]

this isn't really about free variables, right? This is more "you found a recur with no buddy TailLoop"?

We should probably always error here.

[patrick-schultz]

I think I disagree. Names of loops are basically variables, and recur nodes reference a loop variable. The allowFreeVariables flag basically means "let me run this on a subexpression sitting inside some unknown context". The set of containing TailLoops is part of that context. I don't see why we should forbid running NormalizeNames (or any other local transformation) on the body of a TailLoop.

[tpoterba]

Ah, OK, I'm convinced. I would really like to remove the allowFreeVars stuff (it was implemented as a hack to get stuff working in the state of our compiler) but there are still places where we run it on incomplete subtrees without information about context.

[tpoterba]

👍

[tpoterba]

can we add a test of the semantics when various pieces of the loop return missing?

Args, accumulators, condition, etc.

fixed!

[catoverdrive]

@tpoterba I also rewrote the docs, if you want to take a look.