Hi,
I am looking at an issue in Swift Autodiff where we're generating unoptimized code when the derivative of a function is called from a conditional control-flow context.
Upon digging further I found out that the ClosureLifetimeFixup pass is modifying the code in a way such that Swift AD's usual peephole and inlining optimizations aren't able to fully optimize the call to the derivative.
I had some high level questions, mostly around the handling of the convert_escape_to_noescape instruction.
I want to understand the purpose of this pass. At a high-level I believe that the pass is required, when an escaping closure is converted to a noescape closure, to ensure that the original escaping closure stays alive for as long as the noescape closure is used. Is my understanding correct in general?
Does the pass need to ensure the "liveness" of the escaping closure so that the compiler doesn't move a release of the escaping closure to a point before the last use of the noescape closure?
Is there a doc/PR/commit with more details about the pass, that I can read?
What happens if we don't have this pass? Could someone provide a SIL example that might fail to compile without this pass?
I didn't write the pass, but I can try to answer some of your questions. The ideal purpose of the pass is to eliminate escaping closures completely and replace them with nonescaping representations. It's done as a separate pass because SILGen historically has not been set up to deal with parameters as anything other than independent rvalues or inout parameters, and since escaping closures fall more in the "independent rvalue" bucket, we would emit them as such during SILGen and use the convert_escape_to_noescape instruction as a marker for the ClosureLifetimeFixup pass to determine the proper nonescaping lifetime and rewrite the closure as nonescaping. If the pass can't directly rewrite the closure as nonescaping, then it will do as you said and try to extend the original escaping closure's lifetime long enough for the nonescaping closure conversion to be valid, but we generally want to avoid that fallback whenever possible, since the semantics of noncopyable and nonescaping types rely on the closure representation being nonescaping for borrow checking to work.
Because we have more fleshed-out support for noncopyable and nonescaping types now, it might be worth investigating fixing SILGen now to emit nonescaping closures with the correct representation directly and seeing whether we can retire the ClosureLifetimeFixup pass. We should be able to treat a nonescaping closure similar to a borrowed argument, where the lifetime of the argument begins in a formal access before the receiving function call, and we end the formal access after the call completes.
Thanks for the thorough explanation @Joe_Groff! I have some follow-up questions.
The ideal purpose of the pass is to eliminate escaping closures completely and replace them with nonescaping representations.
Is this basically done so we could allocate the closure context more efficiently -- from heap to stack?
If the pass can't directly rewrite the closure as nonescaping, then it will do as you said...
Perhaps I missed this in the code, but I did not see a case where the escaping closure is directly rewritten as a non-escaping closure. Are you referring to the case when the escaping closure appears in the entry block and so it's easier to figure out where to insert the retains and releases for it?
A more general question about escaping closures. My understanding of them was that semantically they're telling the compiler that they will be alive for "static" (or the duration of the program), so their context must be created accordingly (often necessitating heap allocations).
Is this understanding valid? If yes, why do we need to "fix up" their lifetimes when we're using these closures as noescape?
Or perhaps my understanding of "fix up" is not quite right and what the pass does is choose between assigning a "static" or a shortened lifetime to the closure (with a preference for a shortened lifetime).
Nonescaping closures can be allocated more efficiently, but they're also semantically different, because their limited lifetime means that they're able to reference inout parameters and borrows from the outer scope.
tryRewriteToPartialApplyStack is the part that rewrites closures as nonescaping.
You are correct that nonescaping closures have statically determinable lifetimes, so they never need to be heap allocated. Like I said, the fixup pass is primarily a workaround for SILGen historically being unable to establish the proper static lifetime up front naturally.
If I were to take the work (for emitting noescape closures directly in SILGen) on, would you be able to provide some guidance and pointers to get started?
The implementation in SILGen is hopefully greatly simplified by nonescaping function values still being confined to literals in call expressions. So you might start by looking at ArgEmitter::emit in SILGenApply.cpp. You can see there are some checks there already to tryEmitBorrowed when handling borrowing arguments to calls; when an argument can be emitted as a borrow, we set it up as a DelayedArgument so that it can be emitted in a formal access surrounding the actual call. It won't be exactly the same as the existing borrow/inout cases, but you might be able to introduce a new kind of delayed argument for nonescaping closures; similarly, we would form the closure as an on-stack partial_apply right before the call begins, then dealloc_stack it right after. @Andrew_Trick and @meg-gupta have been working on generalized nonescaping types and might have some ideas for how to do nonescaping closure emission in a way that makes it easier to generalize in the future too.
The closure fixup pass is longstanding gripe of mine because the SIL produced by SILGen has invalid ownership, is unverifiable, and horribly confusing to anyone trying to make sense of it.
For better or worse, access to ~Escapable types are handled similarly to @noescape functions.
The difference being that lifetime dependencies are resolved after their access scopes are "fixed". So at least we don't have invalid SIL in the meantime.
I was originally hoping someone would formalize "borrowed arguments" in SILGen, but neither of us touched SILGen to bringup ~Escapable.
In the current world, how come SILGen is unable to establish proper static lifetime for closures up front?
Since we already have @escaping closures, shouldn't we be able to classify all closure direct/transitive uses as escaping/non-escaping? The tryRewriteToPartialApplyStack function in the closure lifetime fix up pass seems to be doing something similar, but it's still pretty conservative. Is there a reason we can't do more and do it early on - in SILGen itself?
Or is this what Joe was hinting at when he said that we could have a special DelayedArgument type for closures that can be definitively treated as non-escaping.
Ideally, SILGen would emit nonescaping captures as borrows the way @Joe_Groff explained above.
The main purpose of ClosureLifetimeFixup is to extend the lifetime of captured values to "cover" applications of the nonescaping closure. That might not be possible without lifetime analysis, which is best done immediately after SILGen as part of a general "lifetime completion" pass. That's effectively the same as doing it during SILGen though.
Ideally, SILGen would emit nonescaping captures as borrows the way @Joe_Groff explained above.
Do you see any blockers in SILGen actually doing this now?
And even if we can do it, based on your second comment it sounds like we'd be more or less porting some of the logic from the ClosureLifetimeFixup pass into SILGen, but the pass itself will still be necessary since some closures will just need to be emitted as escaping. In that case, is it worth making the changes that Joe suggested?
I'm essentially trying to understand if I can simply handle Swift Autodiff specific cases in tryRewriteToPartialApplyStack. But if that just creates more debt for something we'd like to fix in the future then I'd like to fix things the "right" way now.
I also wanted to make sure I have a precise understanding of what you meant by "SILGen would emit nonescaping captures as borrows". It's the arguments captured by a nonescaping closures that should be emitted as borrows, right?
If we can reliably emit nonescaping closures with their correct lifetimes during SILGen, then ClosureLifetimeFixup shouldn't be necessary as a mandatory pass anymore, because escaping closures don't semantically require a definite lifetime. There may be parts of it that are still useful as an optimization.
@Joe_Groff gave a nice explanation of how to change SILGen to (1) emit arguments to (captures of) nonescaping closure as borrows and (2) directly emit partial_apply with the [on_stack] flag.
ClosureLifetimeFixup has a special case for handling optional closures. It allocates an optional value that holds the captured value on control flow paths that call the closure, and .none on other paths. That may need to be factored out.
I do think this is a fair amount of work because it changes the representation that many mandatory SIL passes will see. This means potentialy updating those passes and rewriting lots of test cases. So you can weight that against just hacking on tryRewriteToPartialApplyStack.