-
Notifications
You must be signed in to change notification settings - Fork 250
[hail] tail-recursive loops #7614
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Awesome! I will take a closer look, but here's my quick reaction to the interface. Was I was hoping to write was:
This is basically modeled after letrec in lisp/scheme/ml which would look something like:
Another option is to get rid of
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!) |
@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? |
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. |
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:
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. |
FWIW, I'm also strongly in favor of: |
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.
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.
My proposal has two main differences. In
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
or, if we're giving names to loops, it might be simpler to pass the break and recur functions to the lambda:
The second difference is in the typing. In this PR, the In the type checker, 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 In the IR, I don't think this requires much change. If we're already adding a continuation context (as mentioned above), then 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.
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.
As a side note, @iitalics stream emitter can handle streams of multiple values fine. Effectively, you can make a |
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.
Ah, I thought @catoverdrive was referring to IR level streams. |
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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This was a very legible PR, nice work.
from hail.typecheck import anytype, typecheck | ||
from hail.utils.java import Env | ||
|
||
# FIXME, infer loop type? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we remove this FIXME?
# FIXME, infer loop type? | ||
@typecheck(f=anytype, typ=hail_type, exprs=expr_any) | ||
def loop(f: Callable, typ, *exprs): | ||
"""Expression for writing tail recursive expressions. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we make this a little less "expression"-y?
@@ -130,7 +130,9 @@ def wrapper(func, *args, **kwargs): | |||
|
|||
|
|||
def assert_evals_to(e, v): | |||
assert hl.eval(e) == v | |||
res = hl.eval(e) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
does pytest not generate a nice error here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
it didn't for me. Don't know if that was some sort of setting thing.
@@ -886,6 +886,21 @@ object Interpret { | |||
SafeRow(coerce[PTuple](t), ctx.r, resultOffset).get(0) | |||
case x: ReadPartition => | |||
fatal(s"cannot interpret ${ Pretty(x) }") | |||
case TailLoop(name, args, body) => |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't really care if we implement new nodes in Interpret. It's not long for the world.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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)
case Some(n) => n | ||
case None => | ||
if (!allowFreeVariables) | ||
throw new RuntimeException(s"found free variable in normalize: $name") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 TailLoop
s 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
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
@@ -37,7 +37,7 @@ This is the list of things you need to do to add a new IR node. | |||
|
|||
- Implement it in Emit (the compiler) or add it to Compilable as false | |||
|
|||
- If it binds a variable, add support to Binds and Subst | |||
- If it binds a variable, add support to Bindings |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
👍
TInt32()))) | ||
|
||
assertEvalsTo(triangleSum, FastIndexedSeq(5 -> TInt32()), 15 + 10 + 5) | ||
} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
can we add a test of the semantics when various pieces of the loop return missing?
Args, accumulators, condition, etc.
@tpoterba I also rewrote the docs, if you want to take a look. |
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:
but I'm not sure I really like any of them.
Recur
is difficult to check and enforce.Recur
unambiguously refers to a loop defined in the surrounding scope, mostly treatingRecur
as something of a function reference.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!)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.Recur
in a method. I don't believe there are any cases where we wrapIf
conditions orLet
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.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=1cc @cseed @patrick-schultz @chrisvittal