Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
This is the end result of the design process in #31253. # Overview This PR implements a new kind of closure, called an `opaque closure`. 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, but I'm expecting it will find other use cases as well. From the end user perspective, opaque closures basically behave like regular closures, expect that they are introduced by adding the `@opaque` macro (not part of this PR, but will be added after). In front of the existing closure. In particular, all scoping, capture, etc. rules are identical. For such user written closures, the primary difference is in the performance characteristics. In particular: 1) Passing an opaque closure to a high order function will specialize on the argument and return types of the closure, but not on the closure identity. (This also means that the opaque closure will not be eligible for inlining into the higher order function, unless the inliner can see both the definition and the call site). 2) The optimizer is allowed to modify the capture environment of the opaque closure (e.g. dropping unused captures, or reducing `Box`ed values back to value captures). The `opaque` part of the naming comes from the notion that semantically, nothing is supposed to inspect either the code or the capture environment of the opaque closure, since the optimizer is allowed to choose any value for these that preserves the behavior of calling the opaque closure itself. # 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 nothing : return a 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 nothing : return x.a 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. This provides a very straightforward implementation of this feature, as opaque closures may be inserted at any point during the compilation process (unlike types, which may only be inserted by the frontend). # Mechanism The primary concept introduced by this PR is the `OpaqueClosure{A<:Tuple, R}` type, constructed, by the new `Core._opaque_closure` builtin, with the following signature: ``` _opaque_closure(argt::Type{<:Tuple}, lb::Type, ub::Type, source::CodeInfo, captures...) Create a new OpaqueClosure taking arguments specified by the types `argt`. When called, this opaque closure will execute the source specified in `source`. The `lb` and `ub` arguments constrain the return type of the opaque closure. In particular, any return value of type `Core.OpaqueClosure{argt, R} where lb<:R<:ub` is semantically valid. If the optimizer runs, it may replace `R` by the narrowest possible type inference was able to determine. To guarantee a particular value of `R`, set lb===ub. ``` Captures are available to the CodeInfo as `getfield` from Slot 1 (referenced by position). # Examples I think the easiest way to understand opaque closures is look through a few examples. These make use of the `@opaque` macro which isn't implemented yet, but makes understanding the semantics easier. Some of these examples, in currently available syntax can be seen in test/opaque_closure.jl ``` oc_trivial() = @opaque ()::Any->1 @show oc_trivial() # ()::Any->◌ @show oc_trivial()() # 1 oc_inf() = @opaque ()->1 # Int return type is inferred @show oc_inf() # ()::Int->◌ @show oc_inf()() # 1 function local_call(b::Int) f = @opaque (a::Int)->a + b f(2) end oc_capture_opt(A) = @opaque (B::typeof(A))->ndims(A)*B @show oc_capture_opt([1; 2]) # (::Vector{Int},)::Vector{Int}->◌ @show sizeof(oc_capture_opt([1; 2]).env) # 0 @show oc_capture_opt([1 2])([3 4]) # [6 8] ```
- Loading branch information
