hilbert-ordering for pairwise! in Distances.jl

hilb_ordering.jl

@enum hilb_state _down _up _left _right

function pairwise_recurse_multithreaded!(res,d,A,B, thresh, nthreads)
   # thresh = threshcompute(A);
   jobs = Vector{UnitRange{Int}}(nthreads)
   resize!(jobs, 0)
   si = round(Int,size(A,2)/nthreads)
   el = 0
   for i in 1 : nthreads-1
       push!(jobs, 1 + el : el+si)
       el+= si
   end
   push!(jobs, el+1 : size(A,2))
   #dump(jobs)
   Threads.@threads for irange in jobs
       nohilbert_map_block(irange, 1:size(B,2), thresh, blockpairwise!, d, res, A, B);
   end;
end

function pairwise_hilbert_multithreaded!(res,d,A,B, thresh, nthreads)
   # thresh = threshcompute(A);
   jobs = Vector{UnitRange{Int}}(nthreads)
   resize!(jobs, 0)
   si = round(Int,size(A,2)/nthreads)
   el = 0
   for i in 1 : nthreads-1
       push!(jobs, 1 + el : el+si)
       el+= si
   end
   push!(jobs, nth, el+1 : size(A,2))
   dump(jobs)
   Threads.@threads for irange in jobs
       hilbert_map_block(irange, 1:size(B,2), thresh, blockpairwise!, d, res, A, B);
   end;
end

function pairwise_recurse_singlethreaded!(res,d,A,B, thresh)
   # thresh = threshcompute(A);
    nohilbert_map_block(1:size(A,2), 1:size(B,2), thresh, blockpairwise!, d, res, A, B);
end
function pairwise_hilbert_singlethreaded!(res,d,A,B, thresh)
    #thresh = threshcompute(A);
    hilbert_map_block(1:size(A,2), 1:size(B,2), thresh, Val{_right}() ,blockpairwise!, d, res, A, B);
end

function threshcompute(A, target_usage = 24_000)
    memsize = size(A,1)*sizeof(eltype(A))*2;
    max(8, convert(Int,round(Int, target_usage/memsize)))
end


@noinline function blockpairwise!(irange,jrange,d, res, A, B)
    @inbounds for j in jrange
        for i in irange
            res[i,j] = evaluate(d, fastview(A, :, i), fastview(B, :, j))
        end
    end
    ;
end

@noinline function nohilbert_map_block(irange::UnitRange, jrange::UnitRange, thresh,  f, arg1, arg2, arg3, arg4)
    ni = irange.stop - irange.start
    nj = jrange.stop-jrange.start
    L = irange.start : irange.start + ni >>1
    R = irange.start + ni >>1 +1 : irange.stop
    T = jrange.start: jrange.start + nj>>1
    B = jrange.start + nj>>1 +1 : jrange.stop
    if (ni > thresh) & (nj > thresh)
            nohilbert_map_block(L,T,thresh, f, arg1, arg2, arg3, arg4 )
            nohilbert_map_block(L,B,thresh, f, arg1, arg2, arg3, arg4 )
            nohilbert_map_block(R,B,thresh, f, arg1, arg2, arg3, arg4 )
            nohilbert_map_block(R,T,thresh, f, arg1, arg2, arg3, arg4 )
    elseif ni > thresh
            nohilbert_map_block(L,jrange,thresh, f, arg1, arg2, arg3, arg4 )
            nohilbert_map_block(R,jrange,thresh, f, arg1, arg2, arg3, arg4 )
    elseif nj > thresh
            nohilbert_map_block(irange,T,thresh, f, arg1, arg2, arg3, arg4 )
            nohilbert_map_block(irange,B,thresh, f, arg1, arg2, arg3, arg4 )
    else
        f(irange,jrange,arg1, arg2, arg3, arg4)
    end
    ;
end


@noinline function hilbert_map_block(irange::UnitRange, jrange::UnitRange, thresh, ::Val{hstate},  f, arg1, arg2, arg3, arg4) where {hstate}
    ni = irange.stop - irange.start
    nj = jrange.stop-jrange.start
    L = irange.start : irange.start + ni >>1
    R = irange.start + ni >>1 +1 : irange.stop
    T = jrange.start: jrange.start + nj>>1
    B = jrange.start + nj>>1 +1 : jrange.stop
    if (ni > thresh) & (nj > thresh)
        if hstate == _right
            hilbert_map_block(L,T,thresh, Val{_down}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,B,thresh, Val{_right}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,B,thresh, Val{_up}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,T,thresh, Val{_right}(), f, arg1, arg2, arg3, arg4 )
        elseif hstate == _down #TLTR
            hilbert_map_block(L,T,thresh, Val{_right}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,T,thresh, Val{_down}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,B,thresh, Val{_left}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,B,thresh, Val{_down}(), f, arg1, arg2, arg3, arg4 )
        elseif hstate == _up #TLTR
            hilbert_map_block(R,B,thresh, Val{_left}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,B,thresh, Val{_up}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,T,thresh, Val{_right}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,T,thresh, Val{_up}(), f, arg1, arg2, arg3, arg4 )
        elseif hstate == _left #TLTR
            hilbert_map_block(R,B,thresh, Val{_up}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,T,thresh, Val{_left}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,T,thresh, Val{_down}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,B,thresh, Val{_left}(), f, arg1, arg2, arg3, arg4 )
        else
            dump(hstate)
            assert(false)
        end
    elseif ni > thresh
        if hstate== _left
            hilbert_map_block(R,jrange,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(L,jrange,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
        else
            hilbert_map_block(L,jrange,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(R,jrange,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
        end
    elseif nj > thresh
        if hstate== _up
            hilbert_map_block(irange,B,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(irange,T,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
        else
            hilbert_map_block(irange,T,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
            hilbert_map_block(irange,B,thresh, Val{hstate}(), f, arg1, arg2, arg3, arg4 )
        end
    else
        f(irange, jrange, arg1, arg2, arg3, arg4)
    end
    ;
end