No type inference on struct fields

[Edmundee]

When functions are given as an input a struct whose field is an OffsetArray type inference on that function fails if the code tries to access an element of the array and defaults to ::Any. If my understanding of the profiler is correct, this causes run-time dispatching for subsequent code that depends on this and significantly deteriorates the performance of the code.

using OffsetArrays

struct struct_normal
    x::Array{Int64,2}
end

struct struct_offset
    x::OffsetArray{Int64,2}
end

x = zeros(Int64, 1:10,1:10)
tst1 = struct_normal(x)
tst2 = struct_offset(x)

function fn(x)
    return x.x[1,1]
end

@code_typed fn(tst1)
@code_typed fn(tst2)

Produces the following output:

CodeInfo(
1 ─ %1 = Base.getfield(x, :x)::Matrix{Int64}
│   %2 = Base.arrayref(true, %1, 1, 1)::Int64
└──      return %2
) => Int64

CodeInfo(
1 ─ %1 = Base.getfield(x, :x)::OffsetMatrix{Int64, AA} where AA<:AbstractMatrix{Int64}
│   %2 = Base.getindex(%1, 1, 1)::Any
└──      return %2
) => Any

Now, the usual approach to fixing this would be to let julia optimize on the function boudaries, by separating the function so that all types are known at compile time.

function fn_x(x::OffsetArray{Int64,2})
    return x[1,1]
end

function fn_alternative(x::struct_offset)
    return fn_x(x.x)
end

@code_typed fn_x(tst2.x)
@code_typed fn_alternative(tst2)

However, this still results in untyped code, despite the fact that the intermediate function is correctly typed.

CodeInfo(
1 ─       goto #6 if not true
2 ─ %2  = Core.tuple(1, 1)::Tuple{Int64, Int64}
│   %3  = Base.getfield(x, :parent)::Matrix{Int64}
│   %4  = Base.arraysize(%3, 1)::Int64
│   %5  = Base.arraysize(%3, 2)::Int64
│   %6  = Base.slt_int(%4, 0)::Bool
│   %7  = Core.ifelse::typeof(Core.ifelse)
│   %8  = (%7)(%6, 0, %4)::Int64
│   %9  = Base.slt_int(%5, 0)::Bool
│   %10 = Core.ifelse::typeof(Core.ifelse)
│   %11 = (%10)(%9, 0, %5)::Int64
│   %12 = Base.getfield(x, :offsets)::Tuple{Int64, Int64}
│   %13 = Base.getfield(%12, 1, true)::Int64
│   %14 = Base.getfield(%12, 2, true)::Int64
│   %15 = Base.sub_int(1, %13)::Int64
│   %16 = Base.sle_int(1, %15)::Bool
│   %17 = Base.sle_int(%15, %8)::Bool
│   %18 = Base.and_int(%16, %17)::Bool
│   %19 = Base.sub_int(1, %14)::Int64
│   %20 = Base.sle_int(1, %19)::Bool
│   %21 = Base.sle_int(%19, %11)::Bool
│   %22 = Base.and_int(%20, %21)::Bool
│   %23 = Base.and_int(%22, true)::Bool
│   %24 = Base.and_int(%18, %23)::Bool
└──       goto #4 if not %24
3 ─       Base.nothing::Nothing
└──       goto #5
4 ─       invoke Base.throw_boundserror(x::OffsetMatrix{Int64, Matrix{Int64}}, %2::Tuple{Int64, Int64})::Union{}
└──       unreachable
5 ─       nothing::Nothing
6 ┄ %31 = Base.getfield(x, :offsets)::Tuple{Int64, Int64}
│   %32 = Base.getfield(%31, 1, true)::Int64
│   %33 = Base.getfield(%31, 2, true)::Int64
│   %34 = Base.sub_int(1, %32)::Int64
│   %35 = Base.sub_int(1, %33)::Int64
│   %36 = Base.getfield(x, :parent)::Matrix{Int64}
│   %37 = Base.arrayref(false, %36, %34, %35)::Int64
└──       goto #7
7 ─       return %37
) => Int64

CodeInfo(
1 ─ %1 = Base.getfield(x, :x)::OffsetMatrix{Int64, AA} where AA<:AbstractMatrix{Int64}        
│   %2 = Main.fn_x(%1)::Any
└──      return %2
) => Any

Is this expected behaviour?

[jishnub]

Note that OffsetArray{Int64,2} isn't a concrete type, and a concrete type would be something like OffsetArray{Int64,2,Matrix{Int64}}. I would recommend parameterizing the struct as

julia> struct struct_offset2{A<:OffsetMatrix{Int64}}
               x::A
       end

which should always lead to a concretely typed instance, and would help with subsequent type inference.

julia> tst3 = struct_offset2(x);

julia> @code_typed fn(tst3)
CodeInfo(
1 ─ %1  = Base.getfield(x, :x)::OffsetMatrix{Int64, Matrix{Int64}}
└──       goto #6 if not true
2 ─ %3  = Core.tuple(1, 1)::Tuple{Int64, Int64}
│   %4  = Base.getfield(%1, :parent)::Matrix{Int64}
│   %5  = Base.arraysize(%4, 1)::Int64
│   %6  = Base.arraysize(%4, 2)::Int64
│   %7  = Base.slt_int(%5, 0)::Bool
│   %8  = Core.ifelse::typeof(Core.ifelse)
│   %9  = (%8)(%7, 0, %5)::Int64
│   %10 = Base.slt_int(%6, 0)::Bool
│   %11 = Core.ifelse::typeof(Core.ifelse)
│   %12 = (%11)(%10, 0, %6)::Int64
│   %13 = Base.getfield(%1, :offsets)::Tuple{Int64, Int64}
│   %14 = Base.getfield(%13, 1, true)::Int64
│   %15 = Base.getfield(%13, 2, true)::Int64
│   %16 = Base.sub_int(1, %14)::Int64
│   %17 = Base.sle_int(1, %16)::Bool
│   %18 = Base.sle_int(%16, %9)::Bool
│   %19 = Base.and_int(%17, %18)::Bool
│   %20 = Base.sub_int(1, %15)::Int64
│   %21 = Base.sle_int(1, %20)::Bool
│   %22 = Base.sle_int(%20, %12)::Bool
│   %23 = Base.and_int(%21, %22)::Bool
│   %24 = Base.and_int(%23, true)::Bool
│   %25 = Base.and_int(%19, %24)::Bool
└──       goto #4 if not %25
3 ─       Base.nothing::Nothing
└──       goto #5
4 ─       invoke Base.throw_boundserror(%1::OffsetMatrix{Int64, Matrix{Int64}}, %3::Tuple{Int64, Int64})::Union{}
└──       unreachable
5 ─       nothing::Nothing
6 ┄ %32 = Base.getfield(%1, :offsets)::Tuple{Int64, Int64}
│   %33 = Base.getfield(%32, 1, true)::Int64
│   %34 = Base.getfield(%32, 2, true)::Int64
│   %35 = Base.sub_int(1, %33)::Int64
│   %36 = Base.sub_int(1, %34)::Int64
│   %37 = Base.getfield(%1, :parent)::Matrix{Int64}
│   %38 = Base.arrayref(false, %37, %35, %36)::Int64
└──       goto #7
7 ─       return %38
) => Int64