Brute-Force#

#include <raft/neighbors/brute_force.cuh>

namespace raft::neighbors::brute_force

template<typename value_t, typename idx_t>
inline void knn_merge_parts(raft::resources const &handle, raft::device_matrix_view<const value_t, idx_t, row_major> in_keys, raft::device_matrix_view<const idx_t, idx_t, row_major> in_values, raft::device_matrix_view<value_t, idx_t, row_major> out_keys, raft::device_matrix_view<idx_t, idx_t, row_major> out_values, size_t n_samples, std::optional<raft::device_vector_view<idx_t, idx_t>> translations = std::nullopt)#

Performs a k-select across several (contiguous) row-partitioned index/distance matrices formatted like the following:

part1row1: k0, k1, k2, k3 part1row2: k0, k1, k2, k3 part1row3: k0, k1, k2, k3 part2row1: k0, k1, k2, k3 part2row2: k0, k1, k2, k3 part2row3: k0, k1, k2, k3 etc…

The example above shows what an aggregated index/distance matrix would look like with two partitions when n_samples=3 and k=4.

When working with extremely large data sets that have been broken over multiple indexes, such as when computing over multiple GPUs, the ids will often start at 0 for each local knn index but the global ids need to be used when merging them together. An optional translations vector can be supplied to map the starting id of each partition to its global id so that the final merged knn is based on the global ids.

Usage example: 

#include <raft/core/resources.hpp>
#include <raft/neighbors/brute_force.cuh>
using namespace raft::neighbors;

raft::raft::resources handle;
...
compute multiple knn graphs and aggregate row-wise
(see detailed description above)
...
brute_force::knn_merge_parts(handle, in_keys, in_values, out_keys, out_values, n_samples);

Template Parameters:

  • idx_t –

  • value_t –

Parameters:

  • handle – [in]

  • in_keys – [in] matrix of input keys (size n_samples * n_parts * k)

  • in_values – [in] matrix of input values (size n_samples * n_parts * k)

  • out_keys – [out] matrix of output keys (size n_samples * k)

  • out_values – [out] matrix of output values (size n_samples * k)

  • n_samples – [in] number of rows in each partition

  • translations – [in] optional vector of starting global id mappings for each local partition

template<typename idx_t, typename value_t, typename matrix_idx, typename index_layout, typename search_layout, typename epilogue_op = raft::identity_op>
void knn(raft::resources const &handle, std::vector<raft::device_matrix_view<const value_t, matrix_idx, index_layout>> index, raft::device_matrix_view<const value_t, matrix_idx, search_layout> search, raft::device_matrix_view<idx_t, matrix_idx, row_major> indices, raft::device_matrix_view<value_t, matrix_idx, row_major> distances, distance::DistanceType metric = distance::DistanceType::L2Unexpanded, std::optional<float> metric_arg = std::make_optional<float>(2.0f), std::optional<idx_t> global_id_offset = std::nullopt, epilogue_op distance_epilogue = raft::identity_op())#

Flat C++ API function to perform a brute force knn on a series of input arrays and combine the results into a single output array for indexes and distances. Inputs can be either row- or column-major but the output matrices will always be in row-major format.

Usage example: 

#include <raft/core/resources.hpp>
#include <raft/neighbors/brute_force.cuh>
#include <raft/distance/distance_types.hpp>
using namespace raft::neighbors;

raft::raft::resources handle;
...
auto metric = raft::distance::DistanceType::L2SqrtExpanded;
brute_force::knn(handle, index, search, indices, distances, metric);

Parameters:

  • handle – [in] the cuml handle to use

  • index – [in] vector of device matrices (each size m_i*d) to be used as the knn index

  • search – [in] matrix (size n*d) to be used for searching the index

  • indices – [out] matrix (size n*k) to store output knn indices

  • distances – [out] matrix (size n*k) to store the output knn distance

  • metric – [in] distance metric to use. Euclidean (L2) is used by default

  • metric_arg – [in] the value of p for Minkowski (l-p) distances. This is ignored if the metric_type is not Minkowski.

  • global_id_offset – [in] optional starting global id mapping for the local partition (assumes the index contains contiguous ids in the global id space)

  • distance_epilogue – [in] optional epilogue function to run after computing distances. This function takes a triple of the (value, rowid, colid) for each element in the pairwise distances and returns a transformed value back.

template<typename value_t, typename idx_t, typename idx_layout, typename query_layout>
void fused_l2_knn(raft::resources const &handle, raft::device_matrix_view<const value_t, idx_t, idx_layout> index, raft::device_matrix_view<const value_t, idx_t, query_layout> query, raft::device_matrix_view<idx_t, idx_t, row_major> out_inds, raft::device_matrix_view<value_t, idx_t, row_major> out_dists, raft::distance::DistanceType metric)#

Compute the k-nearest neighbors using L2 expanded/unexpanded distance.

This is a specialized function for fusing the k-selection with the distance computation when k < 64. The value of k will be inferred from the number of columns in the output matrices.

Usage example: 

#include <raft/core/resources.hpp>
#include <raft/neighbors/brute_force.cuh>
#include <raft/distance/distance_types.hpp>
using namespace raft::neighbors;

raft::raft::resources handle;
...
auto metric = raft::distance::DistanceType::L2SqrtExpanded;
brute_force::fused_l2_knn(handle, index, search, indices, distances, metric);

Template Parameters:

  • value_t – type of values

  • idx_t – type of indices

  • idx_layout – layout type of index matrix

  • query_layout – layout type of query matrix

Parameters:

  • handle – [in] raft handle for sharing expensive resources

  • index – [in] input index array on device (size m * d)

  • query – [in] input query array on device (size n * d)

  • out_inds – [out] output indices array on device (size n * k)

  • out_dists – [out] output dists array on device (size n * k)

  • metric – [in] type of distance computation to perform (must be a variant of L2)