Attention

The vector search and clustering algorithms in RAFT are being migrated to a new library dedicated to vector search called cuVS. We will continue to support the vector search algorithms in RAFT during this move, but will no longer update them after the RAPIDS 24.06 (June) release. We plan to complete the migration by RAPIDS 24.10 (October) release and they will be removed from RAFT altogether in the 24.12 (December) release.

Matrix Ordering#

Argmax#

#include <raft/matrix/argmax.cuh>

namespace raft::matrix

template<typename math_t, typename idx_t, typename matrix_idx_t>
void argmax(raft::resources const &handle, raft::device_matrix_view<const math_t, matrix_idx_t, row_major> in, raft::device_vector_view<idx_t, matrix_idx_t> out)#

Argmax: find the col idx with maximum value for each row.

Parameters:

  • handle – [in] raft handle

  • in – [in] input matrix of size (n_rows, n_cols)

  • out – [out] output vector of size n_rows

Argmin#

#include <raft/matrix/argmin.cuh>

namespace raft::matrix

template<typename math_t, typename idx_t, typename matrix_idx_t>
void argmin(raft::resources const &handle, raft::device_matrix_view<const math_t, matrix_idx_t, row_major> in, raft::device_vector_view<idx_t, matrix_idx_t> out)#

Argmin: find the col idx with minimum value for each row.

Parameters:

  • handle – [in] raft handle

  • in – [in] input matrix of size (n_rows, n_cols)

  • out – [out] output vector of size n_rows

Select-K#

#include <raft/matrix/select_k.cuh>

namespace raft::matrix

enum class SelectAlgo : uint8_t#

Algorithm used to select the k largest neighbors.

Details about how the the select-k algorithms in RAFT work can be found in the paper “Parallel Top-K Algorithms on GPU: A Comprehensive Study and New Methods” https://doi.org/10.1145/3581784.3607062. The kRadix* variants below correspond to the ‘Air Top-k’ algorithm described in the paper, and the kWarp* variants correspond to the ‘GridSelect’ algorithm.

Values:

enumerator kAuto#

Automatically pick the select-k algorithm based off the input dimensions and k value

enumerator kRadix8bits#

Radix Select using 8 bits per pass

enumerator kRadix11bits#

Radix Select using 11 bits per pass, fusing the last filter step

enumerator kRadix11bitsExtraPass#

Radix Select using 11 bits per pass, without fusing the last filter step

enumerator kWarpAuto#

Automatically switches between the kWarpImmediate and kWarpFiltered algorithms based off of input size

enumerator kWarpImmediate#

This version of warp_sort adds every input element into the intermediate sorting buffer, and thus does the sorting step every Capacity input elements.

This implementation is preferred for very small len values.

enumerator kWarpFiltered#

This version of warp_sort compares each input element against the current estimate of k-th value before adding it to the intermediate sorting buffer. This makes the algorithm do less sorting steps for long input sequences at the cost of extra checks on each step.

This implementation is preferred for large len values.

enumerator kWarpDistributed#

This version of warp_sort compares each input element against the current estimate of k-th value before adding it to the intermediate sorting buffer. In contrast to warp_sort_filtered, it keeps one distributed buffer for all threads in a warp (independently of the subwarp size), which makes its flushing less often.

enumerator kWarpDistributedShm#

The same as warp_sort_distributed, but keeps the temporary value and index buffers in the given external pointers (normally, a shared memory pointer should be passed in).

template<typename T, typename IdxT>
void select_k(raft::resources const &handle, raft::device_matrix_view<const T, int64_t, row_major> in_val, std::optional<raft::device_matrix_view<const IdxT, int64_t, row_major>> in_idx, raft::device_matrix_view<T, int64_t, row_major> out_val, raft::device_matrix_view<IdxT, int64_t, row_major> out_idx, bool select_min, bool sorted = false, SelectAlgo algo = SelectAlgo::kAuto)#

Select k smallest or largest key/values from each row in the input data.

If you think of the input data in_val as a row-major matrix with len columns and batch_size rows, then this function selects k smallest/largest values in each row and fills in the row-major matrix out_val of size (batch_size, k).

Example usage 

using namespace raft;
// get a 2D row-major array of values to search through
auto in_values = {... input device_matrix_view<const float, int64_t, row_major> ...}
// prepare output arrays
auto out_extents = make_extents<int64_t>(in_values.extent(0), k);
auto out_values  = make_device_mdarray<float>(handle, out_extents);
auto out_indices = make_device_mdarray<int64_t>(handle, out_extents);
// search `k` smallest values in each row
matrix::select_k<float, int64_t>(
  handle, in_values, std::nullopt, out_values.view(), out_indices.view(), true);

Template Parameters:

  • T – the type of the keys (what is being compared).

  • IdxT – the index type (what is being selected together with the keys).

Parameters:

  • handle – [in] container of reusable resources

  • in_val – [in] inputs values [batch_size, len]; these are compared and selected.

  • in_idx – [in] optional input payload [batch_size, len]; typically, these are indices of the corresponding in_val. If in_idx is std::nullopt, a contiguous array 0...len-1 is implied.

  • out_val – [out] output values [batch_size, k]; the k smallest/largest values from each row of the in_val.

  • out_idx – [out] output payload (e.g. indices) [batch_size, k]; the payload selected together with out_val.

  • select_min – [in] whether to select k smallest (true) or largest (false) keys.

  • sorted – [in] whether to make sure selected pairs are sorted by value

  • algo – [in] the selection algorithm to use

inline std::ostream &operator<<(std::ostream &os, const SelectAlgo &algo)#
template<typename T, typename IdxT>
void select_k(raft::resources const &handle, raft::device_csr_matrix_view<const T, IdxT, IdxT, IdxT> in_val, std::optional<raft::device_vector_view<const IdxT, IdxT>> in_idx, raft::device_matrix_view<T, IdxT, raft::row_major> out_val, raft::device_matrix_view<IdxT, IdxT, raft::row_major> out_idx, bool select_min, bool sorted = false, SelectAlgo algo = SelectAlgo::kAuto)#

Selects the k smallest or largest keys/values from each row of the input matrix.

This function operates on a row-major matrix in_val with dimensions batch_size x len, selecting the k smallest or largest elements from each row. The selected elements are then stored in a row-major output matrix out_val with dimensions batch_size x k. If the total number of values in a row is less than K, then the extra position in the corresponding row of out_val will maintain the original value. This applies to out_idx

Template Parameters:

  • T – Type of the elements being compared (keys).

  • IdxT – Type of the indices associated with the keys.

Parameters:

  • handle – [in] Container for managing reusable resources.

  • in_val – [in] Input matrix in CSR format with a logical dense shape of [batch_size, len], containing the elements to be compared and selected.

  • in_idx – [in] Optional input indices [in_val.nnz] associated with in_val.values. If in_idx is std::nullopt, it defaults to a contiguous array from 0 to len-1.

  • out_val – [out] Output matrix [in_val.get_n_row(), k] storing the selected k smallest/largest elements from each row of in_val.

  • out_idx – [out] Output indices [in_val.get_n_row(), k] corresponding to the selected elements in out_val.

  • select_min – [in] Flag indicating whether to select the k smallest (true) or largest (false) elements.

  • sorted – [in] whether to make sure selected pairs are sorted by value

  • algo – [in] the selection algorithm to use

Column-wise Sort#

#include <raft/matrix/col_wise_sort.cuh>

namespace raft::matrix

template<typename in_t, typename out_t, typename matrix_idx_t, typename sorted_keys_t>
void sort_cols_per_row(raft::resources const &handle, raft::device_matrix_view<const in_t, matrix_idx_t, raft::row_major> in, raft::device_matrix_view<out_t, matrix_idx_t, raft::row_major> out, sorted_keys_t &&sorted_keys_opt)#

sort columns within each row of row-major input matrix and return sorted indexes modelled as key-value sort with key being input matrix and value being index of values

Template Parameters:

  • in_t – element type of input matrix

  • out_t – element type of output matrix

  • matrix_idx_t – integer type for matrix indexing

  • sorted_keys_t – std::optional<raft::device_matrix_view<in_t, matrix_idx_t, raft::row_major>> sorted_keys_opt

Parameters:

  • handle – [in] raft handle

  • in – [in] input matrix

  • out – [out] output value(index) matrix

  • sorted_keys_opt – [out] std::optional, output matrix for sorted keys (input)

template<typename ...Args, typename = std::enable_if_t<sizeof...(Args) == 3>>
void sort_cols_per_row(Args... args)#

Overload of sort_keys_per_row to help the compiler find the above overload, in case users pass in std::nullopt for one or both of the optional arguments.

Please see above for documentation of sort_keys_per_row.