JET: are union incurred runtime dispatches really happening?

General Usage Performance
,
1

I tend to think that the compiler is pretty smart and staff like union-splitting are handled gracefully under the hood.

However, I was currently using JET.@code_opt and I saw some dynamic dispatch where I thought there was none. Now, since JET is so dependent on the julia internals and compiler, I guess it doesn’t lie ?

The minimum working example is:

using JET

struct MyStruct
    alist::Union{Nothing, Vector{Int}}
end

function dosomething(m)
    return m.alist[1] + m.alist[2]
end

m = MyStruct([1,2,3]);

JET.@report_opt dosomething(m)

giving

═════ 2 possible errors found ═════
β”Œ dosomething(m::MyStruct) @ Main /u/home/wima/fchrstou/code/julia/jet/myjet.jl:6
β”‚ runtime dispatch detected: (%1::Union{Nothing, Vector{Int64}})[1]::Int64
└────────────────────
β”Œ dosomething(m::MyStruct) @ Main /u/home/wima/fchrstou/code/julia/jet/myjet.jl:6
β”‚ runtime dispatch detected: (%10::Union{Nothing, Vector{Int64}})[2]::Int64
└────────────────────

P.S.
I know that the solution is to define

struct MyStruct2{T<:Union{Nothing, Vector{Int}}}
    alist::T
end

But I am wondering, on whether I can trust JET that julia indeed uses dynamic dispatch there ?

3 Likes
2

Try posting in Performance - Julia Programming Language instead, the field here is mathematical optimization, not optimization of code so you might have more answers there :wink:

1 Like
3

JET.jl does have bugs sometimes, but in this case it seems correct (checked with Cthulhu.jl). Not sure why union splitting isn’t helping here.

The weird thing is that even a ::Vector{Int} type assertion can’t help here, it seems. I’ll try to put up a Julia bug report later.

1 Like
4

The results of JET’s report_opt accurately reflect the behavior of the Julia compiler in most cases, so generally speaking you can trust them.
Looking at the results of @code_typed, you will see several :calls to Base.getindex(...):

CodeInfo(
1 ── %1  = Base.getfield(m, :alist)::Union{Nothing, Vector{Int64}}
β”‚    %2  = (isa)(%1, Vector{Int64})::Bool
└───       goto #8 if not %2
2 ── %4  = Ο€ (%1, Vector{Int64})
β”‚    %5  = $(Expr(:boundscheck, true))::Bool
└───       goto #6 if not %5
3 ── %7  = Base.sub_int(1, 1)::Int64
β”‚    %8  = Base.bitcast(Base.UInt, %7)::UInt64
β”‚    %9  = Base.getfield(%4, :size)::Tuple{Int64}
β”‚    %10 = $(Expr(:boundscheck, true))::Bool
β”‚    %11 = Base.getfield(%9, 1, %10)::Int64
β”‚    %12 = Base.bitcast(Base.UInt, %11)::UInt64
β”‚    %13 = Base.ult_int(%8, %12)::Bool
└───       goto #5 if not %13
4 ──       goto #6
5 ── %16 = Core.tuple(1)::Tuple{Int64}
β”‚          invoke Base.throw_boundserror(%4::Vector{Int64}, %16::Tuple{Int64})::Union{}
└───       unreachable
6 ┄─ %19 = Base.getfield(%4, :ref)::MemoryRef{Int64}
β”‚    %20 = Base.memoryrefnew(%19, 1, false)::MemoryRef{Int64}
β”‚    %21 = Base.memoryrefget(%20, :not_atomic, false)::Int64
└───       goto #7
7 ──       goto #9
8 ── %24 = Base.getindex(%1, 1)::Int64
└───       goto #9
9 ┄─ %26 = Ο† (#7 => %21, #8 => %24)::Int64
β”‚    %27 = Base.getfield(m, :alist)::Union{Nothing, Vector{Int64}}
β”‚    %28 = (isa)(%27, Vector{Int64})::Bool
└───       goto #16 if not %28
10 ─ %30 = Ο€ (%27, Vector{Int64})
β”‚    %31 = $(Expr(:boundscheck, true))::Bool
└───       goto #14 if not %31
11 ─ %33 = Base.sub_int(2, 1)::Int64
β”‚    %34 = Base.bitcast(Base.UInt, %33)::UInt64
β”‚    %35 = Base.getfield(%30, :size)::Tuple{Int64}
β”‚    %36 = $(Expr(:boundscheck, true))::Bool
β”‚    %37 = Base.getfield(%35, 1, %36)::Int64
β”‚    %38 = Base.bitcast(Base.UInt, %37)::UInt64
β”‚    %39 = Base.ult_int(%34, %38)::Bool
└───       goto #13 if not %39
12 ─       goto #14
13 ─ %42 = Core.tuple(2)::Tuple{Int64}
β”‚          invoke Base.throw_boundserror(%30::Vector{Int64}, %42::Tuple{Int64})::Union{}
└───       unreachable
14 β”„ %45 = Base.getfield(%30, :ref)::MemoryRef{Int64}
β”‚    %46 = Base.memoryrefnew(%45, 2, false)::MemoryRef{Int64}
β”‚    %47 = Base.memoryrefget(%46, :not_atomic, false)::Int64
└───       goto #15
15 ─       goto #17
16 ─ %50 = Base.getindex(%27, 2)::Int64
└───       goto #17
17 β”„ %52 = Ο† (#15 => %47, #16 => %50)::Int64
β”‚    %53 = Base.add_int(%26, %52)::Int64
└───       return %53
) => Int64

These correspond to the union split branches where m.alist is nothing. In such cases, a MethodError is raised, and since such β€œmust-throw” case is not optimized, it results in dynamic dispatch. However, Inline statically known method errors. by gbaraldi Β· Pull Request #54972 Β· JuliaLang/julia Β· GitHub might change this behavior.

3 Likes
5

What’s going on here, It seems like it says if nothing is there, indexing it is inferred as Int64. Why would it infer that when it would throw an error here?

6

Hey @nsajko . How exactly did you find out that with Cthulhu ?
I am only getting the following

image
image1375Γ—333 47.1 KB

, from which I cannot deduct dynamic dispatch is happening

7

That is pretty huge. So code like the following which I thought was dynamic-dispatch-free is actually not:


struct MyStruct
    el::Union{Nothing, Int}
end

function dosomething(m)
    # or   if !(m.el isa Nothing)
    # or   if !(m.el === nothing)
    if !isnothing(m.el) 
        return m.el + 1
    end
end

julia> m = MyStruct(1);

julia> JET.@report_opt dosomething(m)
═════ 1 possible error found ═════
β”Œ dosomething(m::MyStruct) @ Main /u/home/wima/fchrstou/code/julia/jet/myjet.jl:9
β”‚ runtime dispatch detected: (%13::Union{Nothing, Int64} + 1)::Int64
└────────────────────

my best hopes for Inline statically known method errors. by gbaraldi Β· Pull Request #54972 Β· JuliaLang/julia Β· GitHub :pray:

8

If I’m reading aviatesk’s comment correctly, the dynamic dispatch occurs when indexing nothing throws an error, not anywhere in the vector branch. That lines up with non-erroring calls having zero allocations and the overall method being inferrable. Something similar happens with FunctionWrappers.jl; it is often used to eliminate dynamic dispatch, but JET.jl will still report one in a precompile-time initialization branch that usually never runs.

9

Julia is (currently, IIRC there were plans to improve this) pretty bad at preserving these kinds of invariants on field access across different calls to getproperty. There won’t be dynamic dispatch if you write it like this:

julia> function dosomething(m)
           el = m.el
           if !isnothing(el) 
               return el + 1
           end
       end
dosomething (generic function with 1 method)

julia> JET.@report_opt dosomething(m)
No errors detected

This is because keeping that !isnothing information around is correctly tracked for the same variable there. The erroring path being the culprit is IMO a red herring.

3 Likes
10

Press T, for [T]yped code. EDIT: I’m not sure any more, though, does an entry marked as runtime mean run time dispatch, or does it mean something else…

11

Fresh REPL:

julia> using JET

julia> struct MyStruct
           alist::Union{Nothing,Vector{Int}}
       end

julia> function f(m)
           a = m.alist
           a::Vector{Int}
           a[1]
       end
f (generic function with 1 method)

julia> report_opt(f, Tuple{MyStruct})
═════ 1 possible error found ═════
β”Œ f(m::MyStruct) @ Main ./REPL[3]:4
β”‚ runtime dispatch detected: (%1::Union{Nothing, Vector{Int64}})[1]::Int64
└────────────────────


julia> function g(m)
           a = m.alist
           (a::Vector{Int})[1]
       end
g (generic function with 1 method)

julia> report_opt(g, Tuple{MyStruct})
No errors detected

So the type assertion prevents run time dispatch in one simple case, but not in the other? This was with Julia v1.11, because JET.jl doesn’t work for nightly Julia.

@aviatesk how can I check if this issue persists on nightly Julia.

EDIT: issue here: