extrema implementations

extrema.jl

using Base.Cartesian, Base.Test, BenchmarkTools

extrema_mapslices(v,region) = mapslices(extrema, v, region)

function extrema_functor(a::AbstractArray, dim)
    dim = tuple(dim...)
    mi = minimum(a,dim)
    ma = maximum(a,dim)
    reshape(collect(zip(mi,ma)), size(mi))
end

function extrema_cartesian(A::AbstractArray, dims)
    sz = [size(A)...]
    sz[[dims...]] = 1
    B = Array{Tuple{eltype(A),eltype(A)}}(sz...)
    extrema_cartesian!(B, A)
end

function extrema_cartesian_simd(A::AbstractArray, dims)
    sz = [size(A)...]
    sz[[dims...]] = 1
    B = Array{Tuple{eltype(A),eltype(A)}}(sz...)
    extrema_cartesian_simd!(B, A)
end

function extrema_cartesianrange(A::AbstractArray, dims)
    sz = [size(A)...]
    sz[[dims...]] = 1
    B = Array{Tuple{eltype(A),eltype(A)}}(sz...)
    extrema_cartesianrange!(B, A)
end

@generated function extrema_cartesian!(B,A::Array{T,N}) where {T,N}
    return quote
        sA = size(A)
        sB = size(B)
        @nloops $N i B begin
            AI = @nref $N A i
            (@nref $N B i) = (AI, AI)
        end
        Bmax = sB
        @inbounds @nloops $N i A begin
            AI = @nref $N A i
            @nexprs $N d->(j_d = min(Bmax[d], i_{d}))
            BJ = @nref $N B j
            #if AI < BJ[1]
            #    (@nref $N B j) = (AI, BJ[2])
            #elseif AI > BJ[2]
            #    (@nref $N B j) = (BJ[1], AI)
            #end
            BJ = ifelse(AI < BJ[1], (AI, BJ[2]), BJ)
            BJ = ifelse(AI > BJ[2], (BJ[1], AI), BJ)
            (@nref $N B j) = BJ
        end
        return B
    end
end

@generated function extrema_cartesian_simd!(B,A::Array{T,N}) where {T,N}
    return quote
        sA = size(A)
        sB = size(B)
        @nloops $N i B begin
            AI = @nref $N A i
            (@nref $N B i) = (AI, AI)
        end
        Bmax = sB
        @inbounds @nloops $(N-1) i d -> indices(A,d+1) begin
            @inbounds @simd for i_0 = indices(A,1)
                AI = @nref $N A d->i_{d-1}
                @nexprs $N d->(j_d = min(Bmax[d], i_{d-1}))
                BJ = @nref $N B j
                #if AI < BJ[1]
                #    (@nref $N B j) = (AI, BJ[2])
                #elseif AI > BJ[2]
                #    (@nref $N B j) = (BJ[1], AI)
                #end
                BJ = ifelse(AI < BJ[1], (AI, BJ[2]), BJ)
                BJ = ifelse(AI > BJ[2], (BJ[1], AI), BJ)
                (@nref $N B j) = BJ
            end
        end
        return B
    end
end

@noinline function extrema_cartesianrange!(B, A)
    sA = size(A)
    sB = size(B)
    for I in CartesianRange(sB)
        AI = A[I]
        B[I] = (AI, AI)
    end
    Bmax = CartesianIndex(sB)
    @inbounds @simd for I in CartesianRange(sA)
        J = min(Bmax,I)
        BJ = B[J]
        AI = A[I]
        #if AI < BJ[1]
        #    B[J] = (AI, BJ[2])
        #elseif AI > BJ[2]
        #    B[J] = (BJ[1], AI)
        #end
        BJ = ifelse(AI < BJ[1], (AI, BJ[2]), BJ)
        BJ = ifelse(AI > BJ[2], (BJ[1], AI), BJ)
        B[J] = BJ
    end
    return B
end

foo = Dict(
    10 => rand(10,11,12),
    100 => rand(100,101,102),
    1000 => rand(1000,1001,1002))

global dim, sz
for dim in [1,2,(1,2)], sz in sort(collect(keys(foo)))
    info("dim=",dim,", sz=",sz)

    @test extrema_mapslices(foo[sz],dim) == extrema_functor(foo[sz],dim)
    @test extrema_mapslices(foo[sz],dim) == extrema_cartesian(foo[sz],dim)
    @test extrema_mapslices(foo[sz],dim) == extrema_cartesian_simd(foo[sz],dim)
    @test extrema_mapslices(foo[sz],dim) == extrema_cartesianrange(foo[sz],dim)

    gc(); print("mapslices     "); @btime extrema_mapslices(foo[sz],dim);
    gc(); print("functor       "); @btime extrema_functor(foo[sz],dim);
    gc(); print("cartesian     "); @btime extrema_cartesian(foo[sz],dim);
    gc(); print("cartesian_simd"); @btime extrema_cartesian_simd(foo[sz],dim);
    gc(); print("cartesianrange"); @btime extrema_cartesianrange(foo[sz],dim);
end

strange that @simd did not improve the Cartesian variant much, as it had a huge effect on the CartesianRange variant.