Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
WIP: RFC: Add an alternative implementation of closures
# Overview
This PR is a sketch implementation of an alternative mechanism of
implementing closures. It is designed to be complimenatry 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(b, Float64) ? return 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)- Loading branch information
