Dynamic dispatch with Union{Nothing, ...}

When benchmarking a data structure we encountered an interesting phenomenon. The data structure is

mutable struct Node{K,D}
    parent::Union{Node{K,D},Nothing}
    left::Union{Node{K,D},Nothing}
    right::Union{Node{K,D},Nothing}
    key::K
    bf::Int8
    data::D
end 

and the code contains functions like

@inline function rotate_left(t::AVLTree{K,D}, x::Node{K,D}, x_right::Node{K,D}) where {K,D}
    y = x_right

    if y.left !== nothing
        x.right = y.left
        y.left.parent = x
    else
        x.right = nothing
    end
    y.left = x

    xp = x.parent
    if xp === nothing
        t.root = y
    else
        if xp.left == x
            xp.left = y
        else
            xp.right = y
        end
    end

    y.parent = xp
    x.parent = y

    x.bf -= y.bf * (y.bf >= zero(Int8)) + one(Int8)
    y.bf += x.bf * (x.bf < zero(Int8)) - one(Int8)

    return y
end

First observation: dynamic dispatch generated by the compiler is different for 1.6.3 and 1.7.0-rc1. Is this expected?
Second observation: instead of fixing individual type instabilities by hand we came up with the idea to specialize Base.getproperty and Base.setproperty like so:

_getproperty(x::Nothing, f) = @assert false
_getproperty(x::Node{K,D}, f) where {K,D} = getfield(x, f)
Base.getproperty(x::Union{Nothing, Node{K,D}}, f::Symbol) where {K,D} =
    _getproperty(x, f)

_setproperty!(x::Nothing, f, v) = @assert false
_setproperty!(x::Node{K,D}, f, v) where {K,D} =
    # setfield!(x, f, convert(fieldtype(typeof(x), f), v))
    setfield!(x, f, v)
_setproperty!(x::Node{K,D}, f, ::Nothing) where {K,D} =
    setfield!(x, f, nothing)
_setproperty!(x::Node{K,D}, f, v::Node{K,D}) where {K,D} =
    setfield!(x, f, v)
Base.setproperty!(x::Union{Nothing, Node{K,D}}, f::Symbol, v) where {K,D} =
    _setproperty!(x, f, v)
Base.setproperty!(x::Union{Nothing, Node{K,D}}, f::Symbol, v::Union{Nothing, Node{K,D}}) where {K,D} =
    _setproperty!(x, f, v)

Is this a known technique? Any drawbacks we should know about before we commit this atrocity?

Could you give us a bit more to go on here? Like, output from @code_warntype or something that shows the difference?

The issue is here: Performance regression from Julia 1.6.3 and 1.8.0-DEV Β· Issue #18 Β· krynju/AVLTrees.jl Β· GitHub. Benchmark code is here: goerch / Allocators.jl . I’m usually working with Juno Profiler and not that much used to @code_warntype. I could easily provide profiler logs. Would that help?

Edit: I will try to produce @code_warntype output tomorrow.

1 Like

@code_warntype output looks pretty similar. But running

t = AVLTrees.AVLTree{Int, Nothing}()
x = AVLTrees.Node{Int, Nothing}(1, nothing, nothing)
@code_typed AVLTrees.rotate_left(t, x, x)

generates

CodeInfo(
1 ── %1  = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %2  = (%1 === AVLTrees.nothing)::Bool
β”‚    %3  = Core.Intrinsics.not_int(%2)::Bool
└───       goto #13 if not %3
2 ── %5  = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %6  = (isa)(%5, Nothing)::Bool
└───       goto #4 if not %6
3 ──       Base.setfield!(x, :right, nothing)::Nothing
└───       goto #7
4 ── %10 = (isa)(%5, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #6 if not %10
5 ── %12 = Ο€ (%5, AVLTrees.Node{Int64, Nothing})
β”‚          Base.setfield!(x, :right, %12)::AVLTrees.Node{Int64, Nothing}
└───       goto #7
6 ──       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
7 ┄─ %17 = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %18 = Base.setproperty!::Core.Const(setproperty!)
β”‚    %19 = (isa)(%17, Nothing)::Bool
└───       goto #9 if not %19
8 ── %21 = Ο€ (%17, Nothing)
β”‚          invoke %18(%21::Nothing, :parent::Symbol, _3::AVLTrees.Node{Int64, Nothing})::Any
└───       goto #12
9 ── %24 = (isa)(%17, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #11 if not %24
10 ─ %26 = Ο€ (%17, AVLTrees.Node{Int64, Nothing})
β”‚          Base.setfield!(%26, :parent, x)::AVLTrees.Node{Int64, Nothing}
└───       goto #12
11 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
12 β”„       goto #14
13 ─       Base.setfield!(x, :right, nothing)::Nothing
14 β”„       Base.setfield!(x_right, :left, x)::AVLTrees.Node{Int64, Nothing}
β”‚    %34 = Base.getfield(x, :parent)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %35 = (%34 === AVLTrees.nothing)::Bool
└───       goto #16 if not %35
15 ─       Base.setfield!(t, :root, x_right)::AVLTrees.Node{Int64, Nothing}
└───       goto #24
16 ─ %39 = Ο€ (%34, AVLTrees.Node{Int64, Nothing})
β”‚    %40 = Base.getfield(%39, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %41 = (isa)(%40, Nothing)::Bool
└───       goto #18 if not %41
17 ─       goto #21
18 ─ %44 = (isa)(%40, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #20 if not %44
19 ─ %46 = Ο€ (%40, AVLTrees.Node{Int64, Nothing})
β”‚    %47 = (%46 === x)::Bool
└───       goto #21
20 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
21 β”„ %51 = Ο† (#17 => false, #19 => %47)::Bool
└───       goto #23 if not %51
22 ─ %53 = Ο€ (%34, AVLTrees.Node{Int64, Nothing})
β”‚          Base.setfield!(%53, :left, x_right)::AVLTrees.Node{Int64, Nothing}
└───       goto #24
23 ─ %56 = Ο€ (%34, AVLTrees.Node{Int64, Nothing})
└───       Base.setfield!(%56, :right, x_right)::AVLTrees.Node{Int64, Nothing}
24 β”„ %58 = Base.setproperty!::Core.Const(setproperty!)
β”‚    %59 = (isa)(%34, Nothing)::Bool
└───       goto #26 if not %59
25 ─ %61 = Ο€ (%34, Nothing)
β”‚          invoke %58(_4::AVLTrees.Node{Int64, Nothing}, :parent::Symbol, %61::Nothing)::Any
└───       goto #29
26 ─ %64 = (isa)(%34, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #28 if not %64
27 ─ %66 = Ο€ (%34, AVLTrees.Node{Int64, Nothing})
β”‚          Base.setfield!(x_right, :parent, %66)::AVLTrees.Node{Int64, Nothing}
└───       goto #29
28 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
29 β”„       Base.setfield!(x, :parent, x_right)::AVLTrees.Node{Int64, Nothing}
β”‚    %72 = Base.getfield(x, :bf)::Int8
β”‚    %73 = Base.getfield(x_right, :bf)::Int8
β”‚    %74 = Base.getfield(x_right, :bf)::Int8
β”‚    %75 = Base.sle_int(0, %74)::Bool
β”‚    %76 = Core.bitcast(Core.Int8, %75)::Int8
β”‚    %77 = Core.and_int(%76, 1)::Int8
β”‚    %78 = Base.mul_int(%73, %77)::Int8
β”‚    %79 = Base.add_int(%78, 1)::Int8
β”‚    %80 = Base.sub_int(%72, %79)::Int8
β”‚          Base.setfield!(x, :bf, %80)::Int8
β”‚    %82 = Base.getfield(x_right, :bf)::Int8
β”‚    %83 = Base.getfield(x, :bf)::Int8
β”‚    %84 = Base.getfield(x, :bf)::Int8
β”‚    %85 = Base.slt_int(%84, 0)::Bool
β”‚    %86 = Core.bitcast(Core.Int8, %85)::Int8
β”‚    %87 = Core.and_int(%86, 1)::Int8
β”‚    %88 = Base.mul_int(%83, %87)::Int8
β”‚    %89 = Base.sub_int(%88, 1)::Int8
β”‚    %90 = Base.add_int(%82, %89)::Int8
β”‚          Base.setfield!(x_right, :bf, %90)::Int8
└───       return x_right
) => AVLTrees.Node{Int64, Nothing}

on Julia 1.6.3 and

CodeInfo(
1 ── %1  = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %2  = (%1 === AVLTrees.nothing)::Bool
β”‚    %3  = Core.Intrinsics.not_int(%2)::Bool
└───       goto #13 if not %3
2 ── %5  = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %6  = Base.setproperty!::typeof(setproperty!)
β”‚    %7  = (isa)(%5, Nothing)::Bool
└───       goto #4 if not %7       
3 ── %9  = Ο€ (%5, Nothing)
β”‚          invoke %6(_3::AVLTrees.Node{Int64, Nothing}, :right::Symbol, %9::Nothing)::Any
└───       goto #7
4 ── %12 = (isa)(%5, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #6 if not %12
5 ── %14 = Ο€ (%5, AVLTrees.Node{Int64, Nothing})
β”‚          invoke %6(_3::AVLTrees.Node{Int64, Nothing}, :right::Symbol, %14::AVLTrees.Node{Int64, Nothing})::Any
└───       goto #7
6 ──       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
7 ┄─ %19 = Base.getfield(x_right, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %20 = Base.setproperty!::typeof(setproperty!)
β”‚    %21 = (isa)(%19, Nothing)::Bool
└───       goto #9 if not %21
8 ── %23 = Ο€ (%19, Nothing)
β”‚          invoke %20(%23::Nothing, :parent::Symbol, _3::AVLTrees.Node{Int64, Nothing})::Any
└───       goto #12
9 ── %26 = (isa)(%19, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #11 if not %26
10 ─ %28 = Ο€ (%19, AVLTrees.Node{Int64, Nothing})
β”‚          invoke %20(%28::AVLTrees.Node{Int64, Nothing}, :parent::Symbol, _3::AVLTrees.Node{Int64, Nothing})::Any
└───       goto #12
11 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
12 β”„       goto #14
13 ─       Base.setfield!(x, :right, nothing)::Nothing
14 β”„       Base.setfield!(x_right, :left, x)::AVLTrees.Node{Int64, Nothing}
β”‚    %36 = Base.getfield(x, :parent)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %37 = (%36 === AVLTrees.nothing)::Bool
└───       goto #16 if not %37
15 ─       Base.setfield!(t, :root, x_right)::AVLTrees.Node{Int64, Nothing}
└───       goto #24
16 ─ %41 = Ο€ (%36, AVLTrees.Node{Int64, Nothing})
β”‚    %42 = Base.getfield(%41, :left)::Union{Nothing, AVLTrees.Node{Int64, Nothing}}
β”‚    %43 = (isa)(%42, Nothing)::Bool
└───       goto #18 if not %43
17 ─       goto #21
18 ─ %46 = (isa)(%42, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #20 if not %46
19 ─ %48 = Ο€ (%42, AVLTrees.Node{Int64, Nothing})
β”‚    %49 = (%48 === x)::Bool
└───       goto #21
20 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
21 β”„ %53 = Ο† (#17 => false, #19 => %49)::Bool
└───       goto #23 if not %53
22 ─ %55 = Ο€ (%36, AVLTrees.Node{Int64, Nothing})
β”‚          Base.setfield!(%55, :left, x_right)::AVLTrees.Node{Int64, Nothing}
└───       goto #24
23 ─ %58 = Ο€ (%36, AVLTrees.Node{Int64, Nothing})
└───       Base.setfield!(%58, :right, x_right)::AVLTrees.Node{Int64, Nothing}
24 β”„ %60 = Base.setproperty!::typeof(setproperty!)
β”‚    %61 = (isa)(%36, Nothing)::Bool
└───       goto #26 if not %61
25 ─ %63 = Ο€ (%36, Nothing)
β”‚          invoke %60(_4::AVLTrees.Node{Int64, Nothing}, :parent::Symbol, %63::Nothing)::Any
└───       goto #29
26 ─ %66 = (isa)(%36, AVLTrees.Node{Int64, Nothing})::Bool
└───       goto #28 if not %66
27 ─ %68 = Ο€ (%36, AVLTrees.Node{Int64, Nothing})
β”‚          invoke %60(_4::AVLTrees.Node{Int64, Nothing}, :parent::Symbol, %68::AVLTrees.Node{Int64, Nothing})::Any
└───       goto #29
28 ─       Core.throw(ErrorException("fatal error in type inference (type bound)"))::Union{}
└───       unreachable
29 β”„       Base.setfield!(x, :parent, x_right)::AVLTrees.Node{Int64, Nothing}
β”‚    %74 = Base.getfield(x, :bf)::Int8
β”‚    %75 = Base.getfield(x_right, :bf)::Int8
β”‚    %76 = Base.getfield(x_right, :bf)::Int8
β”‚    %77 = Base.sle_int(0, %76)::Bool
β”‚    %78 = Core.bitcast(Core.Int8, %77)::Int8
β”‚    %79 = Core.and_int(%78, 1)::Int8
β”‚    %80 = Base.mul_int(%75, %79)::Int8
β”‚    %81 = Base.add_int(%80, 1)::Int8
β”‚    %82 = Base.sub_int(%74, %81)::Int8
β”‚          Base.setfield!(x, :bf, %82)::Int8
β”‚    %84 = Base.getfield(x_right, :bf)::Int8
β”‚    %85 = Base.getfield(x, :bf)::Int8
β”‚    %86 = Base.getfield(x, :bf)::Int8
β”‚    %87 = Base.slt_int(%86, 0)::Bool
β”‚    %88 = Core.bitcast(Core.Int8, %87)::Int8
β”‚    %89 = Core.and_int(%88, 1)::Int8
β”‚    %90 = Base.mul_int(%85, %89)::Int8
β”‚    %91 = Base.sub_int(%90, 1)::Int8
β”‚    %92 = Base.add_int(%84, %91)::Int8
β”‚          Base.setfield!(x_right, :bf, %92)::Int8
└───       return x_right
) => AVLTrees.Node{Int64, Nothing}

on 1.7.0-rc1.

You are doing type piracy on e.g. Base.setproperty!(::Nothing) which can lead to invalidations.

Since this seems to be a quite significant performance regression v1.6 I would open an issue on Base Julia with as much information as possible so that it can be tracked properly and hopefully fixed.

2 Likes

Here we go: Performance regression for Union types in 1.7.0-rc1 Β· Issue #42754 Β· JuliaLang/julia Β· GitHub

1 Like

The performance regression part might seem more pressing, but the interesting aspect for me is that the proposed specializations for getproperty and setproperty heal existing dispatch problems in 1.6.3 and the performance regression in 1.7.0-rc1 likewise. There are two possible reasons for this.

  1. Using the original definition from Base
   setfield!(x, f, convert(fieldtype(typeof(x), f), v))

instead of

    setfield!(x, f, v)

introduces type instabilities.

  1. Using Union{Nothing, …) for these recursive structures introduces the need for specializations like
_setproperty!(x::Node{K,D}, f, ::Nothing) where {K,D} =
    setfield!(x, f, nothing)
_setproperty!(x::Node{K,D}, f, v::Node{K,D}) where {K,D} =
    setfield!(x, f, v)

I seem to need both techniques, but can’t Base do this for us?

For now, if you’re on v1.8, you can do callsite inlining to recover the performance:

❯ git diff --no-index noinlined.jl inlined.jl 
diff --git a/noinlined.jl b/inlined.jl
index 3d12124..ca96fc8 100644
--- a/noinlined.jl
+++ b/inlined.jl
@@ -12,6 +12,7 @@ mutable struct AVLTree{K,D}
 end
 
 @inline function rotate_left(t::AVLTree{K,D}, x::Node{K,D}, x_right::Node{K,D}) where {K,D}
+    @inline begin
     y = x_right
 
     if y.left !== nothing
@@ -40,6 +41,7 @@ end
     y.bf += x.bf * (x.bf < zero(Int8)) - one(Int8)
 
     return y
+    end # @inline begin
 end
 
 t = AVLTree{Int, Nothing}(nothing)
2 Likes

The defect is fixed. Thanks @aviatesk!