Internals

generate_ir(optimise::Bool, f, args...; kwargs...)

Returns (original_ir, transformed_ir) for the call f(args...; kwargs...).

The first element is the original IRCode that Julia generates for the call. The second is the transformed IRCode that Libtask would use to implement the produce/consume interface in a TapedTask.

optimise controls whether the transformed IR (i.e. for the TapedTask) is optimised or not. Apart from inspecting the effects of optimisation, setting optimise to false can also be useful for debugging, because the optimisation pass will also perform verification, and will error if the IR is malformed (which can happen if Libtask's transformation pass has bugs).

This is intended purely as a debugging tool, and is not exported.

produce_value(x::Expr)

Returns the value that a produce statement returns. For example, for the statment produce(%x), this function will return %x.

is_produce_stmt(x)::Bool

true if x is an expression of the form Expr(:call, produce, %x) or a similar :invoke expression, otherwise false.

stmt_might_produce(x, ret_type::Type)::Bool

true if x might contain a call to produce, and false otherwise.

inc_args(stmt::T)::T where {T}

Returns a new T which is equal to stmt, except any Arguments present in stmt are incremented by 1. For example

julia> Libtask.inc_args(Core.ReturnNode(Core.Argument(1)))
:(return _2)

get_type(info::ADInfo, x)

Returns the static / inferred type associated to x.

_typeof(x)

Central definition of typeof, which is specific to the use-required in this package. Largely the same as Base._stable_typeof, differing only in a handful of situations, for example:

julia> Base._stable_typeof((Float64,))
Tuple{DataType}

julia> Libtask._typeof((Float64,))
Tuple{Type{Float64}}

replace_captures(oc::Toc, new_captures) where {Toc<:OpaqueClosure}

Given an OpaqueClosure oc, create a new OpaqueClosure of the same type, but with new captured variables. This is needed for efficiency reasons – if build_rrule is called repeatedly with the same signature and intepreter, it is important to avoid recompiling the OpaqueClosures that it produces multiple times, because it can be quite expensive to do so.

replace_captures(mc::Tmc, new_captures) where {Tmc<:MistyClosure}

Same as replace_captures for Core.OpaqueClosures, but returns a new MistyClosure.

module BasicBlockCode

Copied over from Mooncake.jl in order to avoid making this package depend on Mooncake. Refer to Mooncake's developer docs for context on this file.

opaque_closure(
    ret_type::Type,
    ir::IRCode,
    @nospecialize env...;
    isva::Bool=false,
    do_compile::Bool=true,
)::Core.OpaqueClosure{<:Tuple, ret_type}

Construct a Core.OpaqueClosure. Almost equivalent to Core.OpaqueClosure(ir, env...; isva, do_compile), but instead of letting Core.compute_oc_rettype figure out the return type from ir, impose ret_type as the return type.

Warning

User beware: if the Core.OpaqueClosure produced by this function ever returns anything which is not an instance of a subtype of ret_type, you should expect all kinds of awful things to happen, such as segfaults. You have been warned!

Extended Help

This is needed because we make extensive use of our ability to know the return type of a couple of specific OpaqueClosures without actually having constructed them. Without the capability to specify the return type, we have to guess what type compute_ir_rettype will return for a given IRCode before we have constructed the IRCode and run type inference on it. This exposes us to details of type inference, which are not part of the public interface of the language, and can therefore vary from Julia version to Julia version (including patch versions). Moreover, even for a fixed Julia version it can be extremely hard to predict exactly what type inference will infer to be the return type of a function.

Failing to correctly guess the return type can happen for a number of reasons, and the kinds of errors that tend to be generated when this fails tell you very little about the underlying cause of the problem.

By specifying the return type ourselves, we remove this dependence. The price we pay for this is the potential for segfaults etc if we fail to specify ret_type correctly.

misty_closure(
    ret_type::Type,
    ir::IRCode,
    @nospecialize env...;
    isva::Bool=false,
    do_compile::Bool=true,
)

Identical to opaque_closure, but returns a MistyClosure closure rather than a Core.OpaqueClosure.

optimise_ir!(ir::IRCode, show_ir=false)

Run a fairly standard optimisation pass on ir. If show_ir is true, displays the IR to stdout at various points in the pipeline – this is sometimes useful for debugging.

build_callable(sig::Type{<:Tuple})

Returns a MistyClosure which is used by TapedTask to implement the produce-consume interface. If this method has been called using sig in the current world age, will make a copy of an existing MistyClosure. If not, will derive it from scratch (derive the IR + compile it etc).

LazyCallable

Used to implement static dispatch, while avoiding the need to construct the callable immediately. When constructed, just stores the signature of the callable and its return type. Constructs the callable when first called.

All type information is known, so it is possible to make this callable type stable provided that the return type is concrete.

DynamicCallable

Like LazyCallable, but without any type information. Used to implement dynamic dispatch.

callable_ret_type(sig, produce_types)

Computes the types which might possibly be returned from a TapedTask, where sig is the signature (something of the form Tuple{...}) of the function from which the TapedTask is constructed, and produce_types are the possible types which such a call might produce.

In general, computing produce_types requires analysing the produce type of any statment in the IR associated to sig which might produce. See locations where this function is called to see where this happens.

fresh_copy(mc::T) where {T<:MistyClosure}

Creates an independent copy of mc by (carefully) replacing the Refs it contains in its captures. The resulting MistyClosure is safe to run.

This is achieved by replacing most Refs with new Refs of the same (el)type, but with nothing stored in them – values will be stored in them when the new MistyClosure is called. The only exception are DynamicCallables and LazyCallables – these are constructed when the MistyClosures IR is derived, so fresh instances of them are placed in the associated Ref.

The position counter is reset to -1 (indicating that execution should proceed from the start of the IR, rather than eg. jumping to a line following a produce statement.