mapreduce.jl

# https://github.com/JuliaLang/julia/issues/38558

using BenchmarkTools, MappedArrays, Test

f1(x, y) = mapreduce(==, +, x, y)
f1(x, y, dims) = mapreduce(==, +, x, y; dims)

function f2(x, y)  # from @oscardssmith
    total = 0
    for i in eachindex(x, y)
        total += @inbounds x[i] == y[i]
    end
    return total
end

f3(x,y) = reduce(+, mappedarray(==, x, y))  # from @jishnub
f3(x,y,dims) = reduce(+, mappedarray(==, x, y); dims) 

f4(x,y) = mapreduce(((x,y),)->x==y,+,zip(x,y))  # from @stevengj 

function mapreduce_same(f, op, A::Vararg{AbstractArray,N}; kw...) where {N} # from @wheeheee
    tup_f(i) = f(ntuple(j -> @inbounds(A[j][i]), Val(N))...)
    mapreduce(tup_f, op, eachindex(A...); kw...)
end
function mapreduce_same_reshape(f, op, A::Vararg{AbstractArray,N}; kw...) where {N}
    tup_f(i) = f(ntuple(j -> @inbounds(A[j][i]), Val(N))...)
    # inds = reshape(eachindex(A...), axes(A[1]))  # this change allows for dims=2 etc
    ei = eachindex(A...)  # when this is CartesianIndices, reshaping it makes it slow
    inds = ei isa AbstractUnitRange ? reshape(ei, axes(A[1])) : ei
    mapreduce(tup_f, op, inds; kw...)
end

f5(x, y) = mapreduce_same(==, +, x, y; init=0) # from @wheeheee, but renamed not to clash

f6(x, y) = mapreduce_same_reshape(==, +, x, y)
f6(x, y, dims) = mapreduce_same_reshape(==, +, x, y; dims)

@testset for fun in [f1, f2, f3, f4, f5, f6]
    x = [1 2; 3 4]
    y = [1 0; 3 0]
    @test fun(x, y) == mapreduce(==, +, x, y)
    z = y |> permutedims |> transpose  # needs cartesian indexing
    @test fun(x, z) == mapreduce(==, +, x, z)
end;

let x = randn(10240); y = similar(x)  # same vectors as @oscardssmith
    print("Base:          "); @btime f1($x, $y)
    print("loop:          "); @btime f2($x, $y)
    print("MappedArrays:  "); @btime f3($x, $y)
    print("zip:           "); @btime f4($x, $y)
    print("index:         "); @btime f5($x, $y)
    print("index_reshape: "); @btime f6($x, $y)
end;

let x = randn(Float32, 40, 256); y = similar(x) |> permutedims |> transpose  # same length
    print("Base:          "); @btime f1($x, $y)
    print("loop:          "); @btime f2($x, $y)
    print("MappedArrays:  "); @btime f3($x, $y)
    print("zip:           "); @btime f4($x, $y)
    print("index:         "); @btime f5($x, $y)
    print("index_reshape: "); @btime f6($x, $y)
end;

@testset for fun in [f1, f3, f6]
    x = [1 2; 3 4]
    y = [1 0; 3 0]
    @test fun(x, y, 1) == mapreduce(==, +, x, y; dims=1)
    @test fun(x, y, 2) == mapreduce(==, +, x, y; dims=2)
    z = y |> permutedims |> transpose  # needs cartesian indexing
    @test fun(x, z, 1) == mapreduce(==, +, x, z; dims=1)
    @test fun(x, z, 2) == mapreduce(==, +, x, z; dims=2)
end;


let x = randn(Float32, 40, 256); y = similar(x) |> permutedims |> transpose
    for dims in (1, 2)
        @show dims
        print("Base:          "); @btime f1($x, $y, $dims)
        # print("loop:          "); @btime f2($x, $y)
        print("MappedArrays:  "); @btime f3($x, $y, $dims)
        # print("zip:           "); @btime f4($x, $y)
        # print("index:         "); @btime f5($x, $y)
        print("index_reshape: "); @btime f6($x, $y, $dims)
    end
end;


# Times on M1 mac:
#=

julia> let x = randn(10240); y = similar(x)  # same vectors as @oscardssmith
           print("Base:          "); @btime f1($x, $y)
           print("loop:          "); @btime f2($x, $y)
           print("MappedArrays:  "); @btime f3($x, $y)
           print("zip:           "); @btime f4($x, $y)
           print("index:         "); @btime f5($x, $y)
           print("index_reshape: "); @btime f6($x, $y)
       end;
Base:            min 3.047 μs, mean 4.354 μs (3 allocations, 10.06 KiB)
loop:            min 2.051 μs, mean 2.082 μs (0 allocations)
MappedArrays:    min 2.125 μs, mean 2.146 μs (0 allocations)
zip:             min 3.682 μs, mean 3.819 μs (0 allocations)
index:           min 2.176 μs, mean 2.228 μs (0 allocations)
index_reshape:   min 2.134 μs, mean 2.151 μs (0 allocations)

julia> let x = randn(Float32, 40, 256); y = similar(x) # same length
           print("Base:          "); @btime f1($x, $y)
           print("loop:          "); @btime f2($x, $y)
           print("MappedArrays:  "); @btime f3($x, $y)
           print("zip:           "); @btime f4($x, $y)
           print("index:         "); @btime f5($x, $y)
           print("index_reshape: "); @btime f6($x, $y)
       end;
Base:            min 2.259 μs, mean 3.605 μs (6 allocations, 10.17 KiB)
loop:            min 1.396 μs, mean 1.417 μs (0 allocations)
MappedArrays:    min 1.500 μs, mean 1.518 μs (0 allocations)
zip:             min 3.994 μs, mean 4.047 μs (0 allocations)
index:           min 1.408 μs, mean 1.440 μs (0 allocations)
index_reshape:   min 1.500 μs, mean 1.524 μs (0 allocations)

julia> let x = randn(Float32, 40, 256); y = similar(x) |> permutedims |> transpose
           print("Base:          "); @btime f1($x, $y)
           print("loop:          "); @btime f2($x, $y)
           print("MappedArrays:  "); @btime f3($x, $y)
           print("zip:           "); @btime f4($x, $y)
           print("index:         "); @btime f5($x, $y)
           print("index_reshape: "); @btime f6($x, $y)
       end;
Base:            min 9.958 μs, mean 11.193 μs (8 allocations, 10.20 KiB)
loop:            min 3.510 μs, mean 3.535 μs (0 allocations)
MappedArrays:    min 6.667 μs, mean 6.800 μs (0 allocations)
zip:             min 12.333 μs, mean 12.491 μs (0 allocations)
index:           min 6.658 μs, mean 6.731 μs (4 allocations, 128 bytes)
index_reshape:   min 6.575 μs, mean 6.630 μs (0 allocations)

julia> let x = randn(Float32, 40, 256); y = similar(x)
           for dims in (1, 2)
               @show dims
               print("Base:          "); @btime f1($x, $y, $dims)
               # print("loop:          "); @btime f2($x, $y)
               print("MappedArrays:  "); @btime f3($x, $y, $dims)
               # print("zip:           "); @btime f4($x, $y)
               # print("index:         "); @btime f5($x, $y)
               print("index_reshape: "); @btime f6($x, $y, $dims)
           end
       end;
dims = 1
Base:            min 2.634 μs, mean 4.058 μs (6 allocations, 12.16 KiB)
MappedArrays:    min 2.153 μs, mean 2.393 μs (3 allocations, 2.08 KiB)
index_reshape:   min 2.157 μs, mean 2.395 μs (3 allocations, 2.08 KiB)
dims = 2
Base:            min 2.204 μs, mean 3.564 μs (5 allocations, 10.48 KiB)
MappedArrays:    min 2.481 μs, mean 2.523 μs (2 allocations, 416 bytes)
index_reshape:   min 2.481 μs, mean 2.526 μs (2 allocations, 416 bytes)

julia> let x = randn(Float32, 40, 256); y = similar(x) |> permutedims |> transpose
           for dims in (1, 2)
               @show dims
               print("Base:          "); @btime f1($x, $y, $dims)
               # print("loop:          "); @btime f2($x, $y)
               print("MappedArrays:  "); @btime f3($x, $y, $dims)
               # print("zip:           "); @btime f4($x, $y)
               # print("index:         "); @btime f5($x, $y)
               print("index_reshape: "); @btime f6($x, $y, $dims)
           end
       end;
dims = 1
Base:            min 10.375 μs, mean 11.896 μs (6 allocations, 12.16 KiB)
MappedArrays:    min 3.552 μs, mean 3.823 μs (3 allocations, 2.08 KiB)
index_reshape:   min 3.615 μs, mean 4.069 μs (3 allocations, 2.08 KiB)
dims = 2
Base:            min 9.917 μs, mean 11.059 μs (5 allocations, 10.48 KiB)
MappedArrays:    min 4.030 μs, mean 4.090 μs (2 allocations, 416 bytes)
index_reshape:   min 4.028 μs, mean 4.087 μs (2 allocations, 416 bytes)

julia> versioninfo()
Julia Version 1.12.0-DEV.3
Commit ddd7afb951 (2024-02-16 16:36 UTC)
Platform Info:
  OS: macOS (arm64-apple-darwin21.6.0)
  CPU: 8 × Apple M1
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 4 virtual cores)

=#