Improving speed of runtime dispatch detector

At least you probably want to avoid the current half-way behavior, and either give submodules the full treatment or don’t touch them at all. I don’t think I have an opinion about which alternative is better.

You’re right! Total red herring. But I suppose you might want to support single-line modules as well?

Of course — It’s just an edge case I didn’t know could affect the expression structure.

And yeah I think full support of submodules makes sense.

I guess begin would be nice as well. That way you could wrap an entire package by putting a @stable begin in the main file, wrapping everything.

Not sure how to fix all of this though…

2 Likes

How can one opt-out of errors thrown by @stable at callsite?
That is,

@stable f(x) = ...

g(x) = ... f(x) ...

result = @allow_unstable g(123)  # some argument that makes f(x) unstable

I don’t see any obvious way for now, but some optout is crucial to have if @stable is to gain adoption.

Very nice idea! However, the following does not work at present (on main):

julia> @eval @stable module M
       
       f(x) = x
       
       end
ERROR: LoadError: AssertionError: module_body.head == :block

Some other comments:

  • In my experience, calls to Core.Compiler.return_type (or Base.promote_op) cannot be nested, at least not always. If that also holds here, then applying @stable to a whole module has a good change of breaking something.

  • How does @stable interact with macros like @inline, @propagate_inbounds and @assume_effects? It might be a good idea to address this in the documentation.

  • It might help users to explain that the return type Union{} means that the function does not return normally (which most likely is the result of a programming mistake). Currently one gets:

julia> @stable f() = 1/'2'
f (generic function with 1 method)

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

Discussed above:

3 Likes

On a branch I have a @stable warnonly=true ... syntax so it will still run but just warn you once. Not sure that’s what you meant

What do you mean by nested? Also the promote_op should just compile away anyways if there’s no error.

Good idea. What does Test.@inferred do in this situation?

Maybe in such a case the default behavior would be to just run it, so that any errors just show up? Especially if you are doing this on entire modules, it seems like it would be annoying if all errors get captured by it.

Maybe one way to disable at the call site would be to replace Base.isconcretetype(T) with an overload-able function that a macro could overwrite to be always true…?

So it would basically be like

result = @allow_unstable g(123)  # some argument that makes f(x) unstable

gets transformed into

result = begin
    @eval DD.type_instability(::Type{T}) where {T} = false
    output = g(123)
    @eval DD.type_instability(::Type{T}) where {T} = !Base.isconcretetype(T)
    output
end

Obviously some major issues with this but it would “work”.

I’m not sure how you would do it otherwise.

Maybe there could be some global compilation flag?? You would want it to be constant, though, so that the check can still compile away.

You could also do

@stable g(x) = ...
g(x::Int) = ...

and the stable check will only apply to the general method, but not the one on Int where it’s known to be unstable.

Perhaps that’s the right solution here?

Okay all the earlier ideas are implemented. You can do @stable on begin...end blocks and also have include within submodules work:

using DispatchDoctor

@stable begin
    f(x) = x > 0 ? x : 0.0
    g(x) = x > 0 ? x : 0.0
    module A
        module B
            include("myfile.jl") # all functions will be @stable
        end

        h(::Number) = rand(Bool) ? 1.0 : 0


        using DispatchDoctor: @unstable
        @unstable begin
            h(::Int) = rand(Bool) ? 1.0 : 0
            # ^ Can turn off for specific input types!
        end
    end
end

which gives

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

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

julia> A.h(1)
1.0

Actually now include just gets mapped to _stabilizing_include which means you can do

@stable begin
    include("myfile1.jl")
    include("myfile2.jl")
end

and it will work. If those files have their own include, those will also get mapped (since we use the include(mapexpr::Function, m::Module, path) feature).

This is now fixed on main:

using DispatchDoctor

@stable module A
    f(x) = x > 0 ? x : 0.0
end

@test A.f(1.0) == 1.0
@test_throws TypeInstabilityError A.f(1)

This is now implemented for any Union{} return value, meaning we get:

julia> @stable my_bad_function(x) = x / "blah";

julia> my_bad_function(1)
ERROR: MethodError: no method matching /(::Int64, ::String)

whereas before it would throw an instability error.

As of the recent change to allow arbitrary code blocks, @stable can now be used in combination with any other macro, and in any order. All it does is recurse through expressions until it finds a function definition. So

@stable @inline f(x, i) = x[i]

will get parsed like

@inline $(_stable_fnc(:(f(x, i) = x[i])))

which would allow for the @inline to tag the resultant function.

Basically it should allow other macros to work as normal.

4 Likes

When promote_op analyzes a function that itself calls promote_op, then it may give up and return Any. I think this is the issue mentioned in this JuliaCon talk at 2m25s. Unfortunately, the screen is blank in the video for some reason. (@schlichtanders could you remind us what the example was?)

For this reason I would expect something like

@stable f(x) = # something involving g(x)
@stable g(x) = # something

to produce problems when the function bodies become sufficiently complex. I’ve tried to produce an example, but so far it has always worked. However, I’ve run into such a problem myself when inferring the return type of functions in LinearCombinations.jl. I’ve used Core.Compiler.return_type, but I think promote_op is just wrapper around it.

1 Like

Ok I got one approach working for this, no macros needed:

julia> using DispatchDoctor

julia> @stable f(x) = x > 0 ? x : 0.0
f (generic function with 1 method)

julia> g(x) = f(x);

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

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

the allow_unstable context manipulates a global flag (guarded by a ReentrantLock): DispatchDoctor.jl/src/DispatchDoctor.jl at f00ce7b711dc2443d1a24a4b83fdc092692514f7 · MilesCranmer/DispatchDoctor.jl · GitHub which is checked by the function body before throwing an error:

func[:body] = quote
    let $closure() = $(func[:body]), $T = $(Base).promote_op($closure)
        if $(type_instability)($T) && !$(is_precompiling)() && $(checking_enabled)()
            $err
        end

        return $closure()::$T
    end
end

You will also notice the is_precompiling function there. That’s a quality of life thing; I just found it was annoying to get the type instability errors during precompilation. Much easier to use if I can check them by hand later.

And since the type_instability($T) gets hardcoded, the LLVM IR is kept nice and safe:

julia> @stable f(x) = x
f (generic function with 1 method)

julia> @code_llvm f(1)
define i64 @julia_f_560(i64 signext %0) #0 {
top:
  ret i64 %0
}

Yeah I’m not sure either. I just wrapped some of my other Julia packages with @stable warnonly=true begin ... end and stuff seemed to work just fine. Maybe the most recent Julia is better at this?

The only difference I see is that promote_op has special behavior for Union{}: julia/base/promotion.jl at b93cd432a3fd1e88d4d6a5e36f32e6fbfa0a4160 · JuliaLang/julia · GitHub which maybe (?) is necessary for special handling of MethodErrors.

Also it looks like promote_op is used more throughout the stdlib so I feel a bit safer using it.

But yes if there’s a better way please let me know!

I’ve tried DispatchDoctor with my package SmallCollections.jl. It doesn’t work as expected. Without DispatchDoctor I get

julia> @code_typed SmallBitSet([1, 2, 3])
CodeInfo(
[...]
) => SmallBitSet{UInt64}

and with DispatchDoctor

julia> @code_typed SmallBitSet([1, 2, 3])
CodeInfo(
[...]
) => SmallBitSet

which is not a concrete type. This gives an error when running the test for the package.

Here is what I did it (probably not the most efficient way): I’ve copied the source (v0.1.2) to a new directory, changed the first lines file src/SmallCollections.jl to

using DispatchDoctor
@stable module SmallCollections

Then I change into the src directory, include SmallCollections.jl and say using .SmallCollections.

1 Like

You said there was an error when running it, what does the error message say?

Also sometimes @code_typed lies (see Can we support a way to make `@code_typed` stop lying to us all? :) · Issue #32834 · JuliaLang/julia · GitHub), so maybe try it with Test.@inferred or Cthulhu.jl?

Basically the before result might not actually be true, because @code_typed is using incorrect specialisation rules when compared to the actual Julia compiler. @code_typed tends to be overly optimistic. The @stable macro annotates the return value by the compiler-inferred type, which might be tricking @code_typed to be more realistic in the second case you observed.

For full-package use I’ve been trying it out in DynamicExpressions.jl by doing

module DynamicExpressions
using DispatchDoctor
@stable begin
# package contents
end
end

So far the type instabilities found look to be real.

1 Like

It looks like the culprit in your library is _push. If I annotate the return value of that function with ::SmallBitSet{U}, the type instability error goes away. Im not sure really sure why _push is type unstable though, maybe something about the for-loop makes the inference struggle.

1 Like

You have figured it out already, but just for completeness: I do

julia> include("SmallCollections.jl")   # loading modified source code as described above
julia> using .SmallCollections
julia> SmallBitSet([1,2,3])
ERROR: TypeInstabilityError: Instability detected in `SmallBitSet` defined at SmallCollections.jl:3 with arguments `(Vector{Int64},)`. Inferred to be `SmallBitSet`, which is not a concrete type.

I still don’t think that there is a type instablity. Two more reasons:

  1. The package tests include something analogous to
julia> @inferred SmallBitSet([1,2,3])

(via my own macro @test_inferred), and this test passes.

  1. Look at the following benchmark:
function f(n)
    if n <= 8
        SmallBitSet{UInt8}(1:n)
    elseif n <= 16
        SmallBitSet{UInt16}(1:n)
    elseif n <= 32
        SmallBitSet{UInt32}(1:n)
    else
        SmallBitSet{UInt64}(1:n)
    end
end

g(n) = SmallBitSet(1:n)   # calls SmallBitSet{UInt}(1:n)
julia> n = 40; @b length(f($n))
25.530 ns (1 allocs: 16 bytes)

julia> n = 40; @b length(g($n))
3.167 ns

If there were an instability, then it wouldn’t be so fast and without allocations, I would say.

Weird I’m not sure. It could also be because you have some Vararg which are not specializing, but when you write the number of args explicitly in the benchmark, Julia constant propagation saves you (I see the first SmallBitSet is on args...). See Performance Tips · The Julia Language

I’m really not sure other than that. Would be great if you could make a MWE of this as otherwise this is a tough point to start from.

(Also maybe you could move it to a GitHub issue so we stop sending notifications to people in this thread :sweat_smile: )

1 Like

This looks like an MWE:

struct SmallBitSet{U}
    mask::U
    global _SmallBitSet(mask::U) where U = new{U}(mask)
end

@stable function _push(mask::U, iter) where U
    for n in iter
        mask |= one(U) << (Int(n)-1)
    end
    _SmallBitSet(mask)
end

With DispatchDoctor:

julia> _push(UInt(0), ())
ERROR: TypeInstabilityError: Instability detected in `_push`

Without DispatchDoctor:

julia> @code_llvm _push(UInt(0), ())
;  @ REPL[2]:1 within `_push`
define [1 x i64] @julia__push_145(i64 zeroext %0) #0 {
top:
;  @ REPL[2]:5 within `_push`
; ┌ @ REPL[1]:3 within `_SmallBitSet`
   %unbox.fca.0.insert = insertvalue [1 x i64] zeroinitializer, i64 %0, 0
   ret [1 x i64] %unbox.fca.0.insert
; └
}

(Please post it to the GitHub, not this thread. Also make sure it’s minimal — for example I’m sure the -1 and Int() likely aren’t needed…)

1 Like

One more design question… should @stable skip unknown macros? For example I realised that MacroTools.@capture permits function syntax to specify a similarity search. @stable would break that code.

Maybe @stable should only propagate through macros if they are a member of an internal list defining known compatible ones (like @inline, @inbounds, etc)?

Oh wait it looks like this isn’t just your package but something stranger about Julia type inference… Sorry about that, @matthias314!

It’s weird stuff. Check this out:

function or1(x::T, iter) where {T}
    for n in iter
        x |= one(T)
    end
    x
end

# Inferred concrete:
Base.promote_op(or1, UInt64, Tuple{}) # UInt64

But if we put the exact same thing in a closure, this doesn’t work:


function or2(x::T, iter) where {T}
    function closure()
        for n in iter
            x |= one(T)
        end
        x
    end
    return closure()
end

# Unstable inference:
Base.promote_op(or2, UInt64, Tuple{})  # Any

However, if we make the contents of the closure “more stable”, we can save it:

function or3(x::T, iter) where {T}
    function closure()
        reduce(|, map(_ -> one(T), iter); init=x)
    end
    return closure()
end

Base.promote_op(or3, UInt64, Tuple{})  # UInt64

Type inference is weird… (Is there any Base.@constprop-like magic to fix this?)

Anyways, at least this issue results in a “loud” error. You will know if there is failed inference somewhere. And a failure of inference within a closure might just point to a spot where a refactor is needed (like that for loop).