Questions about the ClosureLifetimeFixup pass

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.

  1. 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?

  2. 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?

  3. Is there a doc/PR/commit with more details about the pass, that I can read?

  4. What happens if we don't have this pass? Could someone provide a SIL example that might fail to compile without this pass?

@Arnold I saw that you are the original author of the pass. Would you mind answering my questions?

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.

2 Likes

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.

1 Like

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.

1 Like

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.

2 Likes

Thanks for those pointers! I'm in the middle of another task at the moment, but I'll start digging into this once I'm done with that.