DispatchDoctor 🩺

The doctor's orders: no type instability allowed!

Usage

This package provides the @stable macro to enforce that functions have type stable return values.

using DispatchDoctor: @stable

@stable function relu(x)
    if x > 0
        return x
    else
        return 0.0
    end
end

Calling this function will throw an error for any type instability:

julia> relu(1.0)
1.0

julia> relu(0)
ERROR: TypeInstabilityError: Instability detected in function `relu`
with arguments `(Int64,)`. Inferred to be `Union{Float64, Int64}`,
which is not a concrete type.

Code which is type stable should safely compile away the check:

julia> @stable f(x) = x;

with @code_llvm f(1):

define i64 @julia_f_12055(i64 signext %"x::Int64") #0 {
top:
  ret i64 %"x::Int64"
}

Meaning there is zero overhead on this type stability check.

You can use @stable on blocks of code, including begin-end blocks, module, and anonymous functions. The inverse of @stable is @unstable which turns it off:

@stable begin

    f() = rand(Bool) ? 0 : 1.0
    f(x) = x

    module A
        # Will apply to code inside modules:
        g(; a, b) = a + b

        # Will recursively apply to included files:
        include("myfile.jl")

        module B
            # as well as nested submodules!

            # `@unstable` inverts `@stable`:
            using DispatchDoctor: @unstable
            @unstable h() = rand(Bool) ? 0 : 1.0

            # This can also apply to code blocks:
            @unstable begin
                h(x::Int) = rand(Bool) ? 0 : 1.0
                # ^ And target specific methods
            end
        end
    end
end

All methods in the block will be wrapped with the type stability check:

julia> f()
ERROR: TypeInstabilityError: Instability detected in function `f`.
Inferred to be `Union{Float64, Int64}`, which is not a concrete type.

(Tip: you cannot import or define macros within a begin...end block, unless it is at the "top level" of a submodule. So, if you are wrapping the contents of a package, you should either import any macros outside of @stable begin...end, or put them into a submodule.)

(Tip 2: in the REPL, you must wrap modules with @eval, because the REPL has special handling of the module keyword.)

You can disable stability errors for a single scope with the allow_unstable context:

julia> @stable f(x) = x > 0 ? x : 0.0

julia> allow_unstable() do
           f(1)
       end
1

although this will error if you try to use it simultaneously from two separate threads.

Options

You can provide the following options to @stable:

default_mode::String="error":
- Change the default mode from "error" to "warn" to only emit a warning, or "disable" to disable type instability checks by default.
- To locally or globally override the mode for a package that uses DispatchDoctor, you can use the "instability_check" key in your LocalPreferences.toml (typically configured with Preferences.jl).
default_codegen_level::String="debug":
- Set the code generation level to "min" to only generate a single function body for each stabilized function. The default, "debug", generates an entire duplicate function so that @code_warntype can be used.
- To locally or globally override the code generation level for a package that uses DispatchDoctor, you can use the "instability_check_codegen_level" key in your LocalPreferences.toml.
default_union_limit::Int=1:
- Sets the maximum elements in a union to be considered stable. The default is 1, meaning that all unions are considered unstable. A value of 2 would indicate that Union{Float32,Float64} is considered stable, but Union{Float16,Float32,Float64} is not.
- To locally or globally override the union limit for a package that uses DispatchDoctor, you can use the "instability_check_union_limit" key in your LocalPreferences.toml.

Each of these is denoted a default_ because you may set them globally or at a per-package level with Preferences.jl (see below).

Usage in packages

You might find it useful to only enable @stable during unit-testing, to have it check every function in a library, but not throw errors for downstream users. You may also want to have warnings instead of errors.

For this, use the default_mode keyword to set the default behavior:

module MyPackage
using DispatchDoctor
@stable default_mode="disable" begin

# Entire package code

end
end

"disable" as the mode will turn @stable into a no-op, so that DispatchDoctor has no effect on your code by default.

The mode is configurable via Preferences.jl, meaning that, within your test/runtests.jl, you could add a line before importing your package:

using Preferences: set_preferences!

set_preferences!("MyPackage", "instability_check" => "error")

You can also set to be "warn" if you would just like warnings.

You might also find it useful to set the default_codegen_level parameter to "min" instead of the default "debug". This will result in no code duplication, improving precompilation time (although @code_warntype and error messages will be less useful). As with the default_mode, you can configure the codegen level with Preferences.jl by using the "instability_check_codegen_level" key.

Note that for code coverage to work as expected over stabilized code, you will also need to use default_codegen_level="min".

Special Cases

[!NOTE] There are several scenarios and special cases for which type instabilities will be ignored. These are discussed below.

During precompilation.
In unsupported Julia versions.
When loading code changes with Revise.jl*.
- *Basically, @stable will attempt to travel through any include's. However, if you edit the included file and load the changes with Revise.jl, instability checks will get stripped (see Revise#634). The result will be that the @stable will be ignored.
Within certain code blocks and function types:
- Within an @unstable block
- Within a @generated block

- Within a `quote ... end` block
- Within a `macro ... end` block
- Within an incompatible macro, such as
	- `@eval`
	- `@generated`
	- `@assume_effects`
	- `@pure`
	- Or anything else registered as incompatible with `register_macro!`
- Parameterized functions like `MyType{T}(args...) = ...`
- Functions with an expression-based name like `(::MyType)(args...) = ...`
- A function inside another function (a closure).
	- But note the outer function will still be stabilized. So, e.g., `@stable f(x) = map(xi -> xi^2, x)` would stabilize `f`, but not `xi -> xi^2`. Though if `xi -> xi^2` were unstable, `f` would likely be as well, and it would get caught!

Note that you can safely use @stable over all of these cases, and special cases will automatically be skipped. Although, if you use @stable internally in some of these cases, like calling @stable within a function on a closure, such as directly on the xi -> xi^2, then it can still apply.

Eliminating Type Instabilities

Say that you start using @stable and you run into a type instability error. What then? How should you fix it?

The first thing you can try is using @code_warntype on the function in question, which will highlight each individual variable's type with a special color for any instabilities.

Note that some of the lines you will see are from DispatchDoctor's inserted code. If those are bothersome, you can disable the checking with Preferences.set_preferences!("MyPackage", "instability_check" => "disable") followed by restarting Julia.

Other, much more powerful options to try include Cthulhu.jl and JET.jl, which can provide more detailed type instability reports in an easier-to-read format than @code_warntype. Both packages can also descend into your function calls to help you locate the source of the instability.

Caveats

Using @stable is likely to increase precompilation time. (To reduce this effect, try the default_codegen_level above)
Using @stable over an entire package may result in flagging type instabilities on small functions that act as aliases and may otherwise be inlined by the Julia compiler. Try putting @unstable on any suspected such functions if needed.

Credits

Many thanks to @chriselrod and @thofma for tips on this discourse thread.