slow `sum` function for custom arrays with `IndexCartesian`

[gasagna]

I have a custom array MyArray with Base.IndexStyle(a::MyArray) = IndexCartesian(). I noted that the builtin sum function is slower than a naive mysum implementation.

See this MWE

using BenchmarkTools

struct MyArray{T, A<:AbstractArray{T, 3}} <: AbstractArray{T, 3}
    data::A
end

Base.size(a::MyArray) = size(a.data)
Base.@propagate_inbounds @inline function Base.getindex(a::MyArray, i, j, k)
    @boundscheck checkbounds(a.data, i, j, k)
    @inbounds v = a.data[i, j, k]
    return v
end

# change this to IndexLinear to get ~5x speed up
Base.IndexStyle(a::MyArray) = IndexCartesian()

function mysum(a)
    S = zero(eltype(a))
    for k = 1:size(a, 3)
        for j = 1:size(a, 2)
            @simd for i = 1:size(a, 1)
                @inbounds S += a[i, j, k]
            end
        end
    end
    return S
end

N = 64
pa = rand(N, N, N)
a  = MyArray(pa)

@btime mysum($a)
@btime   sum($a)
@btime   sum($pa)

On my machine I get

  47.294 μs (0 allocations: 0 bytes)
  453.761 μs (0 allocations: 0 bytes)
  48.895 μs (0 allocations: 0 bytes)

which makes it 10x slower. I find surprising that disabling bound checking with the flag --check-bounds=no makes sum(::MyArray) twice as fast, i.e.

  46.498 μs (0 allocations: 0 bytes)
  216.276 μs (0 allocations: 0 bytes)
  48.483 μs (0 allocations: 0 bytes)

[KristofferC]

With IndexLinear() the following specialized method is hit:

julia/base/reduce.jl

Lines 147 to 168 in 1ee8965

	@noinline function mapreduce_impl(f, op, A::AbstractArray, ifirst::Integer, ilast::Integer, blksize::Int)
	if ifirst == ilast
	@inbounds a1 = A[ifirst]
	return mapreduce_first(f, op, a1)
	elseif ifirst + blksize > ilast
	# sequential portion
	@inbounds a1 = A[ifirst]
	@inbounds a2 = A[ifirst+1]
	v = op(f(a1), f(a2))
	@simd for i = ifirst + 2 : ilast
	@inbounds ai = A[i]
	v = op(v, f(ai))
	end
	return v
	else
	# pairwise portion
	imid = (ifirst + ilast) >> 1
	v1 = mapreduce_impl(f, op, A, ifirst, imid, blksize)
	v2 = mapreduce_impl(f, op, A, imid+1, ilast, blksize)
	return op(v1, v2)
	end
	end

which uses @simd and @inbounds, otherwise, we end up in

julia/base/reduce.jl

Lines 39 to 52 in 1ee8965

	function mapfoldl_impl(f::F, op::OP, nt::NamedTuple{(:init,)}, itr, i...) where {F,OP}
	init = nt.init
	# Unroll the while loop once; if init is known, the call to op may
	# be evaluated at compile time
	y = iterate(itr, i...)
	y === nothing && return init
	v = op(init, f(y[1]))
	while true
	y = iterate(itr, y[2])
	y === nothing && break
	v = op(v, f(y[1]))
	end
	return v
	end

which does not use these.

[brenhinkeller]

Some current timings on 1.8.3 on an M1 mac

julia> @btime mysum($a)
  75.458 μs (0 allocations: 0 bytes)
130963.16543331955

julia> @btime   sum($a)
  287.375 μs (0 allocations: 0 bytes)
130963.16543331827

julia> @btime   sum($pa)
  49.416 μs (0 allocations: 0 bytes)
130963.16543331878

[adienes]

still slow

[mbauman]

Yep, this is on my punchlist. It's an easy fix, just need to get #58241 in first. It's this fallback:

julia/base/reduce.jl

Line 450 in a909276

_mapreduce(f, op, ::IndexCartesian, A::AbstractArrayOrBroadcasted) = mapfoldl(f, op, A)