knn.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <raft/spatial/knn/detail/processing.hpp>  // MetricProcessor
 
#include <cuvs/distance/distance.hpp>
#include <cuvs/neighbors/ivf_flat.hpp>
#include <cuvs/neighbors/ivf_pq.hpp>
 
namespace raft {
class handle_t;
}
 
namespace ML {
 
void brute_force_knn(const raft::handle_t& handle,
                     std::vector<float*>& input,
                     std::vector<int>& sizes,
                     int D,
                     float* search_items,
                     int n,
                     int64_t* res_I,
                     float* res_D,
                     int k,
                     bool rowMajorIndex                  = false,
                     bool rowMajorQuery                  = false,
                     cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
                     float metric_arg                    = 2.0f,
                     std::vector<int64_t>* translations  = nullptr);
 
void rbc_build_index(const raft::handle_t& handle,
                     std::uintptr_t& rbc_index,
                     float* X,
                     int64_t n_rows,
                     int64_t n_cols,
                     cuvs::distance::DistanceType metric);
 
void rbc_knn_query(const raft::handle_t& handle,
                   const std::uintptr_t& rbc_index,
                   uint32_t k,
                   const float* search_items,
                   uint32_t n_search_items,
                   int64_t dim,
                   int64_t* out_inds,
                   float* out_dists);
 
void rbc_free_index(std::uintptr_t rbc_index);
 
struct knnIndex {
  cuvs::distance::DistanceType metric;
  float metricArg;
  int nprobe;
  std::unique_ptr<raft::spatial::knn::MetricProcessor<float>> metric_processor;
 
  std::unique_ptr<cuvs::neighbors::ivf_flat::index<float, int64_t>> ivf_flat;
  std::unique_ptr<cuvs::neighbors::ivf_pq::index<int64_t>> ivf_pq;
 
  int device;
};
 
struct knnIndexParam {
  virtual ~knnIndexParam() {}
};
 
struct IVFParam : knnIndexParam {
  int nlist;
  int nprobe;
};
 
struct IVFFlatParam : IVFParam {};
 
struct IVFPQParam : IVFParam {
  int M;
  int n_bits;
  bool usePrecomputedTables;
};
 
void approx_knn_build_index(raft::handle_t& handle,
                            knnIndex* index,
                            knnIndexParam* params,
                            cuvs::distance::DistanceType metric,
                            float metricArg,
                            float* index_array,
                            int n,
                            int D);
 
void approx_knn_search(raft::handle_t& handle,
                       float* distances,
                       int64_t* indices,
                       knnIndex* index,
                       int k,
                       float* query_array,
                       int n);
 
void knn_classify(raft::handle_t& handle,
                  int* out,
                  int64_t* knn_indices,
                  std::vector<int*>& y,
                  size_t n_index_rows,
                  size_t n_query_rows,
                  int k);
 
void knn_regress(raft::handle_t& handle,
                 float* out,
                 int64_t* knn_indices,
                 std::vector<float*>& y,
                 size_t n_index_rows,
                 size_t n_query_rows,
                 int k);
 
void knn_class_proba(raft::handle_t& handle,
                     std::vector<float*>& out,
                     int64_t* knn_indices,
                     std::vector<int*>& y,
                     size_t n_index_rows,
                     size_t n_query_rows,
                     int k);
};  // namespace ML