Help me understand constant propagation

General Usage
1

Consider the following simple example. I have overloaded getindex and getproperty for MyType in the exact same way, but one allows constant propagation and the other doesn’t. I’m not even using x[:c] and x.c; I’m calling them both the same way.

using BenchmarkTools

struct Axis{IndexMap} end
Axis(;kwargs...) = Axis{(;kwargs...)}()

Base.getindex(ax::Axis, s::Symbol) = _getindex(ax, Val(s))
@generated _getindex(::Axis{IM}, ::Val{s}) where {IM,s} = :($(getfield(IM, s)))


struct MyType{T,N,A<:AbstractArray{T,N},Axes}
    data::A
    ax::Axes
end

Base.getindex(x::MyType, s::Symbol) = MyType(getfield(x, :data), getindex(getfield(x, :ax), s))
Base.getproperty(x::MyType, s::Symbol) = MyType(getfield(x, :data), getindex(getfield(x, :ax), s))


function test_index(x)
    return getindex(x, :c)
end

function test_prop(x)
    return getproperty(x, :c)
end


ax = Axis(a=1, b=2:4, c=(a=5:6, b=7))
mt = MyType(rand(7), ax)

@btime test_index($mt) # 4.600 μs (1 allocation: 16 bytes)
@btime test_prop($mt) # 5.701 ns (1 allocation: 16 bytes)

What’s strange is I can remove just the call to the MyType constructor and things are fast again:

Base.getindex(x::MyType, s::Symbol) = (getfield(x, :data), getindex(getfield(x, :ax), s))
@btime test_index($mt) # 5.701 ns (1 allocation: 16 bytes)

Then I thought it might help to wrap s as a Val immediately and call _getindex directly, but that didn’t help either:

Base.getindex(x::MyType, s::Symbol) = MyType(getfield(x, :data), _getindex(getfield(x, :ax), Val(s)))
@btime test_index($mt) # 4.628 μs (1 allocation: 16 bytes)

What’s going on here?

2

Try with

@inline Base.getindex(x::MyType, s::Symbol) = MyType(getfield(x, :data), getindex(getfield(x, :ax), s))

looks like this one doesn’t inline.

2 Likes
3

Thanks for the response. That works for this problem, but for some reason not my non-MWE. I’ll have to see what is different and possibly pull together a better example.

4

I have made a habit of always annotating my getindex definitions with @inline, are there any potential downsides to that, provided that the operation is simple?

5

Alright, here is a better example. This is a pared-down version of what I’m doing in ComponentArrays.jl. I’m using @inline for getindex here and am still seeing the issue.

idx_ax(x) = (x, NamedTuple())
idx_ax(x::Tuple) = x

struct Axis{IdxMap} end
Axis(IdxMap) = Axis{IdxMap}()

struct ComponentArray{Axes,T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
    data::A
    axes::Axes
    ComponentArray(data::A, ax::Ax) where {A<:AbstractArray{T,N},Ax<:Axis} where {T,N} = new{Ax,T,N,A}(data, ax)
    ComponentArray(data, ax) = data
end

@inline getdata(x::ComponentArray) = getfield(x, :data)
@inline getaxes(x::ComponentArray) = getfield(x, :axes)
@inline getaxes(::Type{ComponentArray{Ax,T,N,A}}) where {Ax<:Axis,T,N,A} = Ax

Base.size(x::ComponentArray) = size(getdata(x))

@inline Base.getindex(x::ComponentArray, s::Symbol) = _getindex(x, Val(s))
@inline Base.getindex(x::ComponentArray, idx) = getdata(x)[idx]
@inline Base.getindex(::Type{Axis{IdxMap}}, s::Symbol) where IdxMap = idx_ax(getfield(IdxMap, s))

@generated function _getindex(x::ComponentArray, ::Val{s}) where s
    ind_tup = getindex(getaxes(x), s)
    idx = ind_tup[1]
    new_ax = Axis(ind_tup[2])
    return :(Base.@_inline_meta; ComponentArray(Base.maybeview(getdata(x), $idx), $new_ax))
end

@inline Base.getproperty(x::ComponentArray, s::Symbol) = _getindex(x, Val(s))

Constant propagation works just fine on getproperty, but I can’t seem to get it to work with getindex:

ax = Axis((a=1, b=2:4, c=(5:8, (a=1:3, b=4))))
ca = ComponentArray(rand(8), ax)


using BenchmarkTools

@btime $ca.a       # 1.099 ns (0 allocations: 0 bytes)
@btime $ca[:a]     # 4.314 μs (1 allocation: 16 bytes)

@btime $ca.c.b     # 1.099 ns (0 allocations: 0 bytes)
@btime $ca[:c][:b] # 8.933 μs (3 allocations: 80 bytes)

@btime $ca.c.a     # 13.300 ns (2 allocations: 64 bytes)
@btime $ca[:c][:a] # 8.833 μs (4 allocations: 128 bytes)


function test_prop(x)
    return getproperty(x, :c)
end

function test_index(x)
    return getindex(x, :c)
end

@btime test_prop($ca)  # 13.400 ns (2 allocations: 64 bytes)
@btime test_index($ca) # 4.386 μs (2 allocations: 64 bytes)