knn.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2023, 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/distance/distance_types.hpp>
#include <raft/spatial/knn/ann_common.h>
#include <raft/spatial/knn/ball_cover_types.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,
                     raft::distance::DistanceType metric = raft::distance::DistanceType::L2Expanded,
                     float metric_arg                    = 2.0f);
 
void rbc_build_index(const raft::handle_t& handle,
                     raft::spatial::knn::BallCoverIndex<int64_t, float, uint32_t>& index);
 
void rbc_knn_query(const raft::handle_t& handle,
                   raft::spatial::knn::BallCoverIndex<int64_t, float, uint32_t>& index,
                   uint32_t k,
                   const float* search_items,
                   uint32_t n_search_items,
                   int64_t* out_inds,
                   float* out_dists);
void approx_knn_build_index(raft::handle_t& handle,
                            raft::spatial::knn::knnIndex* index,
                            raft::spatial::knn::knnIndexParam* params,
                            raft::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,
                       raft::spatial::knn::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