I played around with some optimizations when answering this SO question. Sought is a function which given matrices x of size (m, l) and y of size (n, l) returns a matrix of size (n, k) where each row consist of the k row numbers in x closest to the corresponding row of y. The important requirement was that this should be done without allocation (except for the result array, obviously).

My best solution uses the following "tricks":

  • regular views instead of JuliennedArrays.Slices I used in my SO answer
  • a manually written non-allocating Euclidean distance, with @fastmath and no square root (which is unnecessary for sorting)
  • linear time replacement function instead of permutedims, using the slice of the result matrix as buffer
  • more @inbounds
function replacetop!(result, x; by=identity)
    f_max, i_max = findmax(by, result)
    if by(x) < f_max
        result[i_max] = x
    return result

function squared_distance(x, y)
    peeled = Iterators.peel(eachindex(x, y))
    isnothing(peeled) && error("vectors must not be empty")
    i_first, i_rest = peeled
    s = @fastmath (x[i_first] - y[i_first])^2
    @inbounds for i in i_rest
        @fastmath s += (x[i] - y[i])^2
    return @fastmath s

function knnslice!(result, x, y)
    @inbounds for (i_r, i_y) in zip(axes(result, 1), axes(y, 1))
        result_row = @view(result[i_r, :])
        fill!(result_row, first(axes(y, 1)))
        f(j) = squared_distance(@view(x[j, :]), @view(y[i_y, :]))
        @inbounds for j_x in axes(x, 1)
            replacetop!(result_row, j_x; by=f)
    return result

knnslice(x, y, k) = knnslice!(similar(x, Int, size(y, 1), k), x, y)

If we loosen the original interface to better respect column-major order, the following transposed implementation is able to shave off a bit more time:

function knnslice_t!(result, x, y)
    @inbounds for (i_r, i_y) in zip(axes(result, 2), axes(y, 2))
        result_row = @view(result[:, i_r])
        fill!(result_row, first(axes(y, 2)))
        f(j) = squared_distance(@view(x[:, j]), @view(y[:, i_y]))
        @inbounds for j_x in axes(x, 2)
            replacetop!(result_row, j_x; by=f)
    return result

knnslice_t(x, y, k) = knnslice_t!(similar(x, Int, k, size(y, 2)), x, y)

Test cases and benchmarks:

x, y = rand(10, 4), rand(5, 4);
xt, yt = permutedims.((x, y));
knnslice(x, y, 5) |> permutedims == knnslice_t(xt, yt, 5)
knnslice([1:10;;], [3;8;;], 3) == [2 3 4; 8 9 7]

@benchmark knnslice!(r, $x, $y) setup=(r = similar(x, Int, size(y, 1), 5))
@benchmark knnslice_t!(r, $xt, $yt) setup=(r = similar(x, Int, 5, size(y, 2)))

I'm open for feedback on everything, but especially on the numerics and loop parts; for example, I don't really know what to be careful about with @fastmath. Also, probably there's existing better implementations of squared_distance? Something with even SIMD, maybe? Does it even pay off with small l?

  • \$\begingroup\$ It seems to me knnslice! doesn't always give correct output. For example, knnslice([1:10;;], [1;;], 3) == [1 1 1] instead of [1 2 3]. If I understand your intent correctly, I think the fill! line should be something like result_row .= axes(x, 1)[1:length(result_row)] instead and the @inbounds for j_x line should be @inbounds for j_x in axes(x, 1)[length(result_row)+1:end]. Similarly for knnslice_t!. Could you check this? \$\endgroup\$
    – Vincent Yu
    May 7 at 6:38


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.