Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement opaque closures #37849

Closed
wants to merge 12 commits into from
Closed

Implement opaque closures #37849

wants to merge 12 commits into from

Commits on Dec 24, 2020

  1. Refactor JL_GC_PUSHFRAME to internalize layout assumptions

    Right now callers of JL_GC_PUSHFRAME need to be aware that the
    roots start three pointers after the end of the struct. That
    seems like an implementation detail of the macro, so refactor the
    macro to also take the assignment location for the locals instead.
    In the process fix an error where the gc-analyzer version of this
    code assumed a different layout.
    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    a74e698 View commit details
    Browse the repository at this point in the history
  2. Implement opaque closures

    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]
    ```
    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    f860f8b View commit details
    Browse the repository at this point in the history
  3. add front-end support for OpaqueClosure

    JeffBezanson authored and Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    be57eaf View commit details
    Browse the repository at this point in the history
  4. Address some review

    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    2ccaa2b View commit details
    Browse the repository at this point in the history
  5. Configuration menu
    Copy the full SHA
    61d9ba9 View commit details
    Browse the repository at this point in the history
  6. Configuration menu
    Copy the full SHA
    c129e34 View commit details
    Browse the repository at this point in the history
  7. Backport some fixes

    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    190f4e9 View commit details
    Browse the repository at this point in the history
  8. Some extra backports

    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    147502b View commit details
    Browse the repository at this point in the history
  9. Fix review nits

    Keno committed Dec 24, 2020
    Configuration menu
    Copy the full SHA
    4b884a5 View commit details
    Browse the repository at this point in the history

Commits on Dec 25, 2020

  1. Switch opaque closure creation to IR form only

    As requested in review.
    Keno committed Dec 25, 2020
    Configuration menu
    Copy the full SHA
    6f97ef0 View commit details
    Browse the repository at this point in the history
  2. Fix opaque closure show method

    Keno committed Dec 25, 2020
    Configuration menu
    Copy the full SHA
    b40a7bd View commit details
    Browse the repository at this point in the history

Commits on Dec 27, 2020

  1. Fix tag number

    Keno committed Dec 27, 2020
    Configuration menu
    Copy the full SHA
    5998195 View commit details
    Browse the repository at this point in the history