RFC: Add an alternative implementation of closures

[Keno]

Overview

This PR is a sketch implementation of an alternative mechanism of
implementing closures. It is designed to be complementary to the
existing closure mechanism and makes some different trade offs.
The motivation for this mechanism comes primarily from closure-based
AD tools like Zygote, but I'm expecting it will find other use cases
as well. It's a little hard to name this, because it's just another
way to implement closures. I discussed this with jeff and options
that were considered were "Arrow Closures" or "Non-Nominal Closures",
but for now I'm just calling them "Yet Another Kind of Closure" (YAKC,
pronounced yak-c, rhymes with yahtzee). This PR is in very early
stages, so in order to explain what this is and what this does, I
may describe features that are not yet implemented. See the end of the
commit message to see the current status.

Motivation

Optimization across closure boundaries

Consider the following situation (type annotations are inference
results, not type asserts)

function foo()
    a = expensive_but_effect_free()::Any
    b = something()::Float64
    ()->!isa(b, Float64) ? return a : return nothing
end

now, the traditional closure mechanism will lower this to:

struct ###{T, S}
   a::T
   b::S
end
(x::###{T,S}) = !isa(x.b, Float64) ? return x.a : return nothing
function foo()
    a = expensive_but_effect_free()::Any
    b = something()::Float64
    new(a, b)
end

the problem with this is apparent: Even though after inference,
we know that a is unused in the closure (and thus would be
able to delete the expensive call were it not for the capture),
we may not delete it, simply because we need to satisfy the full
capture list of the closure. Ideally, we would like to have a mechanism
where the optimizer may modify the capture list of a closure in
response to information it discovers.

Closures from Casette transforms

Compiler passes like Zygote would like to generate new closures
from untyped IR (i.e. after the frontend runs) (and in the future
potentially typed IR also). We currently do not have a great mechanism
to support this (it is somewhat possible by constructing an object
that redoes the primal analysis, but it's awkward at the very least).
This provides a very straightforward implementation of this feature.

Mechanism

The primary concept introduced by this PR is the YAKC type, defined
as follows:

struct YAKC{A <: Tuple, R}
     env::Any
     ci::CodeInfo
 end

 function (y::YAKC{A, R})(args...) where {A,R}
     typeassert(args, A)
     ccall(:jl_invoke_yakc, Any, (Any, Any), y, args)::R
 end

The dynamic semantics are that calling the yakc will run whatever code
is stored in the .ci object, using env as the self argument. This
is augmented by special support in inference and the optimizer to
co-optimize yakcs that appear in bodies of functions along with their
containing functions in order to enable things like cross-closure DCE.

Note that argument types and return types are explicitly specified,
rather than inferred. The reason for this is to prevent YAKCs from
participating in inference cycles and to allow return-type information
to be provided to inference without having to look at the contained
CodeInfo (because the contents of the CodeInfo are not interprocedurally
valid and may in general depend on what the optimizer was able to
figure out about the program). This is also done with a few to an
extension where the CodeInfo inside the yakc is not generated until
later in the optimization pipeline. It would be possible to optionally
allow return-type inference of the yatc based on the unoptimized
CodeInfo (if available), but that is not currently within the scope
of my planned work in this PR.

Status

The PR has the bare bones support for yakcs, including some initial
support for inling and inference, though that support is known to
be incorrect and incomplete. There are also currently no nice
front-end forms. For explanatory and testing purposes, I would like to
provide a macro of the form:

function foo()
    a = expensive_but_effect_free()::Any
    b = something()::Float64
    @yakc ()->!isa(b, Float64) ? return a : return nothing
end

but that is not implemented yet. At the moment, the codeinfo needs
to be spliced in manually (here we're abusing code_lowered to construct
us a CodeInfo of the appropriate form), e.g.:

julia> bar() = 1
bar (generic function with 1 method)

julia> ci = @code_lowered bar();

julia> @eval function foo1()
          f = $(Expr(:new, :(Core.YAKC{Tuple{}, Int64}), nothing, ci))
       end
foo1 (generic function with 1 method)

julia> @eval function foo2()
          f = $(Expr(:new, :(Core.YAKC{Tuple{}, Int64}), nothing, ci))
          f()
       end
foo2 (generic function with 1 method)

julia> foo1()
Core.YAKC{Tuple{},Int64}(nothing, CodeInfo(
1 ─     return 1
))

julia> foo1()()
1

julia> @code_typed foo2()
CodeInfo(
1 ─     return 1
) => Int64

julia> struct Test
           a::Int
           b::Int
       end

julia> (a::Test)() = getfield(a, 1) + getfield(a, 2)

julia> ci2 = @code_lowered Test(1, 2)();

julia> @eval function foo2()
          f = $(Expr(:new, :(Core.YAKC{Tuple{}, Int64}), (1, 2), ci2))
          f()
       end
foo2 (generic function with 1 method)

julia> @code_typed foo2()
CodeInfo(
1 ─     return 3
) => Int64

TODO:

• Show that this actually helps the Zygote use case I care about
• Frontend support
• Better optimizations (detection of need for recursive inlining)
• Codegen for yakcs (right now they're interpreted)

[Keno]

I've tried this out with Zygote and there is a lot of things I like about it. The biggest problem I see at this point and (as I found out) in its current form doesn't work too well to solve the nested AD problem (for nested AD, you'd generally want to co-optimize and in this proposal there isn't a great place to stash that CodeInfo). I think it is quite feasible to modify this to split out an object that encapsulates the set of functions that need to be co-optimized (jl_yakc_herd_t?) which would then be the natural place to cache the transformed versions (and e.g. for the second order derivative, you'd co-optimize all the various codeinfos up to second order). For our own reference attaching the state of the white board at the end of the discussion about this:

I'll need to let this stew for a little while longer.

[Keno]

in its current form doesn't work too well to solve the nested AD problem

The nested AD problem is resolved by switching the representation.

I'm planning to revive this branch, so if there are comments, let me know.

[Keno]

Capturing some discussion with @vtjnash and @JeffBezanson:

The question we were discussing is what to do with inferred return types for YAKC.
Right now, the way I have it set up is that YAKC has type parameters YAKC{A,R}
for the YAKC's argument and return types and the optimizer is allowed to adjust the
R parameter based on inference results. Inference dependent program behavior
(like return_type) is a general cause for concern, so we discussed various options.
In the end we arrived at the following:

1. Rather than punning on Expr(:new), there will be a dedicated yakc constructor yakc(env, lb, ub, ci) where the implementation (interpreter, optimizer, etc.) is allowed to chose the return type between lb and ub according to whaterver rules it likes.
2. YAKC will have a custom lattice element that has full provenance information and can use the CodeInfo to do inference at the proper call site. It will not use the type bounds for this.
3. During the same inference session, we may not lift the inference result into the type domain. This is because the optimizer is what chooses the ultimate representation and optimization is logically (and depending on when the choice happens, actually), sequenced after inference, so inference may not try to use such information.
4. It is nevertheless ok to specialize on the type parameter afterwards if the yakc escapes from its optimization context (but inference may no longer look at the CodeInfo, since it lost provenance information at that point).

[Keno]

One remaining question is what to call this thing. Some options:

• delayed closure
• opaque closure
• thunk
• optic closure
• dynamic closure
• arrow closure
• nnclosure
• Just keep yakc

[oxinabox]

opaque closure

I think this one.

I guess users would creat these via @opaque ()->x+y, or similar?

---

A two birds with one stone option might be to make them also disallow rebinding variables so as to make boxing impossible.
Then could take FastClosure.jl's @closure macro, since it would probably take over all uses of FastClosure.jl

[StefanKarpinski]

Calling it a closure when it captures variables by value would be rather perverse: the defining characteristic of a closure is that it closes over bindings. It should be called @nonclosure in that case.

[Keno]

Capture vs non-capture is a frontend decision. I was planning to use the same lowering, meaning that bindings are captured by references and Boxes are introduced as they are now. The difference is that the optimizer would be allowed to optimize it back to a value capture if it can prove the dynamic semantics are unnecessary (since the opaqueness of the environment guarantees that nobody else is allowed to look at the box).

[Keno]

Alright, please take a look at test/yakc.jl and see if everyone likes the semantics. Frontend support is still missing, but @JeffBezanson said he'd do that for me and I since I don't need it for my use cases, I'll leave that out from this PR for the moment.

Obviously the implementation is primitive right now, lacks some of the corner cases, and most optimizations. However, it does pass the test/yakc.jl test, which I think covers (basic versions of) all the major different features I want it to have, so I'm reasonably confident that this is reasonable, so I'd like to get this on the path to get merged, and then iterate on it on master, as I try to use it for my actual application.

So do we like opaque closure then? If so, I'll do that rename tomorrow.

[JeffreySarnoff]

Naming this YAKC would help the reader avoid misunderstanding by clarifying their unknowing. Another approach to helping the reader would name this opaque closure type OpaqueClosure.

[StefanKarpinski]

Opaque closure seems good to me.

[tpapp]

DarkClosure would be shorter, and kind of cool.

[StefanKarpinski]

Or delayed closure; less cool but more accurate, imo.

[JeffreySarnoff]

delayed closure is better, good suggestion
for the cool kids: now delay is performant

[oxinabox]

Delayed Closure feels like it is exposing an implementation detail.

The opaqueness is the user facing difference.
Isn't it?

[JeffreySarnoff]

everyone is correct -- what to do, what to do :)

[JeffBezanson]

I like OpaqueClosure as well; the type hides what is called (and is intended to hide the captured variables to the extent possible), which is a key difference with our current closures.

[JeffreySarnoff]

great!

[Keno]

Alright, I feel like "opaque closures" have it then.

[JeffBezanson]

Move to tfuncs.jl?

[JeffBezanson]

Do you want to encourage inlining a yakc into its caller, or the function containing the call into its caller? (perhaps both?)

[Keno]

Both, e.g. consider:

call(f::YAKC) = f()
bar() = call(_yakc(...))

We would like to elide the allocation of the yakc and just inline the whatever the code the yakc would have had right into the caller, but in order to be able to do that, we need to see both the call and the definition to the yakc, so a high inline penalty for calling a yakc would be counterproductive.

[JeffBezanson]

Is m.unspecialized guaranteed to be defined here? (Maybe I just need to keep reading)

[JeffBezanson]

?

[Keno]

rebase artifact

[JeffBezanson]

Can we deal with this elsewhere? It would be better if new only has to worry about allocating the type passed to it, instead of also determining what type to allocate.

[Keno]

Yes, sorry. This is outdated and can be removed.

[Keno]

I've done the rename. I'll push it up to a new PR, so that we can have a clean place to do any review and get it merged, while keeping this history here.