Runtime function

[BTV25]

I'm running into performance issues with the current inrange function when applying it to large numbers of target points—mainly due to memory allocated for returning all indices. To address this, I created [PR #209] that introduces an enhanced inrange variant allowing users to pass in a runtime callback function. This function is executed whenever in range indices are found, which avoids storing a vector of vectors and helps reduce memory usage.

I’ve included tests and documentation, and while it’s only slightly faster in practice, the memory efficiency improvement is significant. Hoping this addition can benefit others with similar constraints. Would appreciate any feedback!

[KristofferC]

mainly due to memory allocated for returning all indices.

Curious what you think about something like #167.

[BTV25]

While I think that would be a beneficial change, it still doesn't cut down much on the memory being allocated. What I am suggesting cuts down on that memory significantly. Unfortunately, a minimum working example is a bit long but here is one:

Using the standard inrange function and looping through the indicies:
1.335535 seconds (2.01 M allocations: 152.174 MiB, 3.75% gc time)

While using the in range_runtime! function:
1.278695 seconds (3 allocations: 224 bytes)

using NearestNeighbors

#arbitrary non-allocating calculation
function runtime_function(point_index, neighbor_index, point, neighbors, total_distance) 
    dx = point[1] - neighbors[1, neighbor_index]
    dy = point[2] - neighbors[2, neighbor_index]
    dz = point[3] - neighbors[3, neighbor_index]
    total_distance[point_index] += sqrt(dx*dx + dy*dy + dz*dz)
    return nothing
end

function test_inrange(targets, tree, d, neighbors)
    d .= 0.0
    f(point_index, neighbor_index, point) = runtime_function(point_index, neighbor_index, point, neighbors, d)
    idxs = inrange(tree, targets, 0.01)

    for i in eachindex(idxs)
        for j in idxs[i]
            f(i, j, view(targets,:,i))
        end
    end
    return d
end

function test_inrange_runtime(targets, tree, d, neighbors)
    d .= 0.0
    f(point_index, neighbor_index, point) = runtime_function(point_index, neighbor_index, point, neighbors, d)
    inrange_runtime!(tree, targets, 0.01, f)
    return d
end

function test(n_targets)
    println("Testing with $n_targets points and $n_targets neighbors...")
    d = zeros(n_targets)
    targets = rand(3, n_targets)
    neighbors = rand(3,n_targets)
    tree = KDTree(neighbors)

    println("inrange")
    test_inrange(targets, tree, d, neighbors)
    @time test_inrange(targets, tree, d, neighbors)

    println("inrange_runtime")
    test_inrange_runtime(targets, tree, d, neighbors)
    @time test_inrange_runtime(targets, tree, d, neighbors)
    println()
    return nothing
end

test(10^6)

[KristofferC]

it still doesn't cut down much on the memory being allocated.

Not in absolute terms but it drastically cuts down on the number of objects being allocated which influences GC times etc.

I'm running into performance issues with the current inrange function when applying it to large numbers of target points—mainly due to memory allocated for returning all indices.

Out of curiosity do you have a benchmark that shows how this gets improved with your PR? On its own, it doesn't look like the performance is very noticable.

[BTV25]

Sorry, I didn't put the benchmark in a good spot. Using the above MWE to compare inrange and inrange_runtime! I get the following timing results

Using the standard inrange function and looping through the indicies:
1.335535 seconds (2.01 M allocations: 152.174 MiB, 3.75% gc time)

While using the in range_runtime! function:
1.278695 seconds (3 allocations: 224 bytes)

[KristofferC]

I'm running into performance issues

So the performance issue you run into is because of those 0.07 extra seconds?

[BTV25]

No, the main problem is the 152 MiB that gets allocated. I anticipate that the number of points I am doing inrange searches on could be several orders of magnitude larger than what is shown here, which allocates significantly more memory. I also have to do the inrange search several thousand times so even small time savings could be worth it.

For example, if I increase the problem size by a factor of 10 the benchmark becomes:

Using the standard inrange function and looping through the indicies:
45.984809 seconds (34.95 M allocations: 11.400 GiB, 14.76% gc time)

While using the in range_runtime! function:
35.079411 seconds (3 allocations: 224 bytes)

[KristofferC]

Okay, if you are not going to use the returned indices at all, I understand why you don't want to allocate them. Let's continue in the PR then.