slow sum function for custom arrays with IndexCartesian

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)

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.

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

still slow

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)

Navigation Menu

Repository navigation

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

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	@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

	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

slow sum function for custom arrays with IndexCartesian #33928

Description

Activity

KristofferC commented on Nov 23, 2019

brenhinkeller commented on Nov 20, 2022

adienes commented on May 5, 2025

mbauman commented on May 5, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

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