Significant slowdown of nanmean over slow axis dimensions

[JamesWrigley]

I've run into a case on some large-ish arrays where nanmean is quite slow compared to mean:

using Statistics, NaNStatistics

data = rand(Float64, (2880, 2880, 700));
@time mean(data; dims=3)
@time nanmean(data; dims=3)

2.132474 seconds (9 allocations: 63.282 MiB)
85.518496 seconds (3 allocations: 63.281 MiB)

Whereas with numpy:

import numpy as np

# Note the flipped dimensions because of row-major vs column-major ordering
data = np.random.rand(700, 2880, 2880).astype(np.float64)
%time np.mean(data, axis=0)
%time np.nanmean(data, axis=0)

CPU times: user 2.12 s, sys: 9.91 ms, total: 2.13 s
Wall time: 2.14 s
CPU times: user 9.28 s, sys: 3.86 s, total: 13.1 s
Wall time: 13.2 s

I assume this is what #48 (comment) refers to. Do you reckon there's any way of recovering that performance?

[brenhinkeller]

Yeah, I think that's exactly it. Statistics.mean probably still does some basic loop reordering in comparison.

If LoopModels takes off, that would probably be a good option, but no timeline for that that I know of. In the meanwhile, LoopVectorization.@turbo does still seem to mostly work so I suppose we could try to add LV back in, but not sure if that's the best plan long term unless LoopVectorization becomes more actively maintained again..

[JamesWrigley]

Agreed that it'd be best to avoid LV 😅 But, I wonder if it's possible to rewrite some of the functions using the machinery from Base? I had a look at mean and the definition is actually really simple, it's basically just sum() with type promotion rules: https://github.com/JuliaStats/Statistics.jl/blob/22dee82f9824d6045e87aa4b97e1d64fe6f01d8d/src/Statistics.jl#L192

e.g. in some trivial tests so far this outperforms nansum():

function nansum2(A; dims=:)
    sum(x -> ifelse(isnan(x), zero(x), x), A; dims)
end

data = rand(1000, 1000)
@time nansum(data; dims=2)
@time nansum2(data; dims=2)

0.000719 seconds (3 allocations: 7.891 KiB)
0.000186 seconds (3 allocations: 7.891 KiB)

And on the original (2880, 2880, 700) array nansum2(data; dims=3) only takes 1.9s. The only downside I can see is that we'd be relying on compiler/Base magic to do the right thing. What do you think? What was the original motivation to do the metaprogramming stuff vs using high-level functions?

[brenhinkeller]

Not opposed to trying it, as long as we don’t lose too much performance on smaller arrays! The generated functions play nice with `LoopVectorization.@turbo` was the original motivation, and have had various pros and cons after that. As a pro, it turns out the generated functions are also nice for StaticCompiler and such (they don’t require the Julia standard library, so can be compiled into standalone binaries as long as they don’t allocate), but then that’s maybe less of a priority now with the `juliac` sort of approach. As a con, they used to have larger compile times which could be annoying, but turn out to work quite well with the code caching system so that’s basically a non-issue now. Overall, I’ve gotten the impression that the current approach has been doing reasonably on small arrays and at keeping allocations to a minimum, but I haven’t done a systematic comparison to other approaches in a long time!

[JamesWrigley]

So what I see so far is that there's a only a performance improvement for reductions over a slow scan axis. Here's some results for nansum on vectors:

And nansum on small matrices (the X labels are bit misleading, the 'array size' is the length of one side so x=50 means a 50x50 matrix):

TL;DR on this particular node it's a constant factor of 30% for vectors and 50% for matrices, unless you're reducing over a slow axis of a matrix in which case the Base implementation scales way better.

I would say that's a bit too much performance to give up, especially since those constant factors seem to hold for large arrays. But it does make me think that the existing code is doing something right and it might just need tweaking to fix the performance on slow axes. Or alternatively we could have a dead-simple Base implementation for reductions where dims != 1.

[brenhinkeller]

Nice benchmarks! Yeah, that makes sense. If we knew how base (or I suppose LV) implemented its loop reordering (or whatever other tricks might apply) it's possible that it wouldn't be too hard to tweak the generated functions here to replicate that..

[JamesWrigley]

After staring at static_nansum_quote for a while I realized that it's building up the loops from the first to the last dimension, which is the wrong order for column-major arrays 🙃 That already helped a bit, and refactoring the code to iterate over the dimensions in memory order instead of splitting it up in non-reduce/reduce dimensions matched the performance of nansum2. But that only works for reducing over the last dimensions of an array, I'm not sure how to generalize it. I suspect the easiest solution is keep the existing code for reducing over the fast dimensions and fall back to mapreduce/sum for other dimensions.

[brenhinkeller]

That sounds reasonable to me

[JamesWrigley]

cough took me a while to get back to this 🙈 This implementation is pretty solid:

nanmean2(A::AbstractVector; dims=:) = NaNStatistics.nanmean(A; dims)
nanmean2(A; dims=:) = _nanmean2(A, dims)

_nanmean2(A, dims::Colon) = NaNStatistics.nanmean(A; dims)

function _nanmean2(A, dims::Integer)
    if dims == 1
        NaNStatistics.nanmean(A; dims)
    else
        _nanmean2(A, (dims,))
    end
end

macro _mapreduce_impl(A, out, expr)
    # A copy of the fallback branch in _mapreducedim!() in base/reducedim.jl
    quote
        indsAt = Base.safe_tail(axes($A))
        indsRt = Base.safe_tail(axes($out))
        keep, Idefault = Broadcast.shapeindexer(indsRt)

        for IA in CartesianIndices(indsAt)
            IR = Broadcast.newindex(IA, keep, Idefault)
            @inbounds for ir in axes($A, 1)
                x = A[ir, IA]
                $expr
            end
        end
    end |> esc
end

function _nanmean2_impl!(A, out, counts)
    @_mapreduce_impl(A, out,
            begin
                if !isnan(x)
                    out[ir, IR] += x
                    counts[ir, IR] += 1
                end
            end)

    nan_val = convert(eltype(out), NaN)
    @inbounds for i in eachindex(out)
        if iszero(counts[i])
            out[i] = nan_val
        else
            out[i] /= counts[i]
        end
    end
end

function _nanmean2(A, dims::Tuple)
    # nanmean() is faster if reducing over the first dimension
    if 1 in dims
        return NaNStatistics.nanmean(A; dims)
    end

    out = similar(A, eltype(A), Base.reduced_indices(A, dims))
    fill!(out, zero(eltype(A)))

    # Compute the length of the dimensions we're reducing over to pick
    # the smallest valid eltype for the counts array. On large arrays this
    # can improve performance by ~25%.
    to_reduce_dims = ntuple(i -> size(out)[i] == 1 ? size(A)[i] : 1, ndims(out))
    max_reduction_len = prod(to_reduce_dims)
    counts_T = if max_reduction_len < typemax(UInt8)
        UInt8
    elseif max_reduction_len < typemax(UInt16)
        UInt16
    elseif max_reduction_len < typemax(UInt32)
        UInt32
    else
        UInt64
    end
    counts = similar(out, counts_T)
    fill!(counts, 0)

    _nanmean2_impl!(A, out, counts)

    out
end

The current implementation is faster in basically all cases for small arrays and for reducing over the first dimensions, but when reducing over the slow dimensions of large arrays nanmean2() is significantly faster:

x = rand(Float64, 1000, 1000, 1000);
for i in rand(eachindex(x), 1000)
    x[i] = NaN
end

@time nanmean(x; dims=3);
#  7.136803 seconds (3 allocations: 7.629 MiB)

@time nanmean2(x; dims=3);
#  2.922653 seconds (6 allocations: 9.537 MiB)

[brenhinkeller]

Ah, cool!

[JamesWrigley]

Ok I see the pattern now, the Base implementation is faster once we're hitting memory. This is the benchmark for reducing over the third dimension of (n, n, n) Float64 arrays (the vertical black lines are L1/L2/L3 cache sizes):

And the second dimension of Float64 matrices:

Unfortunately I don't think there's a portable way to get the L3 cache size in Julia 😅 Maybe we could make it an argument and default it to 32MB or 64MB? On the other hand I do wonder if there's a cache-oblivious algorithm for this.

[brenhinkeller]

Ah, nice analysis! Yeah, optimizing for specific cache sizes could be a lot of work.. I wonder how much variation there will be here between different mixes of cache sizes, memory bandwidth, etc..

I suppose there are a few possible philosophies we could go, either
(A) try to take on the problem, add some branches to try to generically optimize the memory-limited case in some way to the extent possible without compromising performance for smaller arrays or
(B) note the potential for memory-limitation to impact performance but leave it all up to the user to chunk/organize their problem in some memory-friendly way

with I suppose the possibility of some third option between these two where we pick some sort of default "chunk size" -- though I suppose I'm not entirely sure how to do that in practice

[JamesWrigley]

I think for now it's ok to go with that third option where we pick a default size limit (16MB seems reasonable?) and let the user change it if they want to optimize it more. Would you be ok with a PR to do that? We could also add an extension for Hwloc to allow the user to pass e.g. small_size_threshold=:L1 to allow passing the cache level they want to use as a threshold. That might be useful because it seems this can be very hardware-dependent, e.g. on another machine the threshold is actually the L1 cache:

In any case, on this machine both implementations handily beat numpy :)

x = np.random.rand(20, 1000, 1000)
%timeit np.nanmean(x, axis=0)
125 ms ± 154 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

julia> x = rand(Float64, (1000, 1000, 20))

julia> @benchmark nanmean($x; dims=3)
BenchmarkTools.Trial: 53 samples with 1 evaluation per sample.
 Range (min … max):  93.703 ms … 99.396 ms  ┊ GC (min … max): 0.00% … 4.07%
 Time  (median):     95.930 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   95.954 ms ±  1.278 ms  ┊ GC (mean ± σ):  0.30% ± 0.80%

     ▃  █     ▃ ▃▃ ▃     █▃  ▃          █    █  ▃ ▃▃           
  ▇▁▁█▁▇█▁▁▇▁▁█▁██▁█▇▇▁▁▁██▇▁█▁▇▇▇▁▇▇▇▇▇█▇▇▁▁█▁▁█▁██▇▇▁▁▇▁▁▁▇ ▁
  93.7 ms         Histogram: frequency by time        98.2 ms <

 Memory estimate: 7.63 MiB, allocs estimate: 3.

julia> @benchmark nanmean2($x; dims=3)
BenchmarkTools.Trial: 163 samples with 1 evaluation per sample.
 Range (min … max):  29.499 ms …  33.343 ms  ┊ GC (min … max): 0.00% … 9.76%
 Time  (median):     30.287 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   30.741 ms ± 982.729 μs  ┊ GC (mean ± σ):  2.11% ± 2.76%

      ▂█▆▃▁      ▁                          ▅▃▁                 
  ▄▆▃▆█████▄▃▅▄▆▅█▃▁▄▄▃▃▅▁▃▃▁▄▁▁▁▁▁▁▃▃▁▁▃▁▇▆███▇▃▃▁▁▄▅▅▃▄▁▃▁▆▄ ▃
  29.5 ms         Histogram: frequency by time         32.5 ms <

 Memory estimate: 8.58 MiB, allocs estimate: 11.

[brenhinkeller]

Go for it for a PR as long as it's speeding things up!