Dynamic dispatch with Union{Nothing, ...}

goerch · October 21, 2021, 7:29pm

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?

kristoffer.carlsson · October 21, 2021, 7:54pm

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

goerch · October 21, 2021, 7:58pm

The issue is here: https://github.com/krynju/AVLTrees.jl/issues/18. 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.

goerch · October 22, 2021, 5:47am

@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.

kristoffer.carlsson · October 22, 2021, 6:54am

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.

goerch · October 22, 2021, 8:19am

Here we go: https://github.com/JuliaLang/julia/issues/42754

goerch · October 22, 2021, 8:32am

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.

Using the original definition from Base

   setfield!(x, f, convert(fieldtype(typeof(x), f), v))

instead of

    setfield!(x, f, v)

introduces type instabilities.

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?

aviatesk · October 22, 2021, 9:43am

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)

goerch · October 28, 2021, 4:43am

The defect is fixed. Thanks @aviatesk!