knn_mg.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <cuml/prims/opg/matrix/data.hpp>
#include <cuml/prims/opg/matrix/part_descriptor.hpp>
 
#include <raft/core/handle.hpp>
 
#include <vector>
 
namespace ML {
namespace KNN {
namespace opg {
 
void knn(raft::handle_t& handle,
         std::vector<MLCommon::Matrix::Data<int64_t>*>* out_I,
         std::vector<MLCommon::Matrix::floatData_t*>* out_D,
         std::vector<MLCommon::Matrix::floatData_t*>& idx_data,
         MLCommon::Matrix::PartDescriptor& idx_desc,
         std::vector<MLCommon::Matrix::floatData_t*>& query_data,
         MLCommon::Matrix::PartDescriptor& query_desc,
         bool rowMajorIndex,
         bool rowMajorQuery,
         int k,
         size_t batch_size,
         bool verbose);
 
void knn_classify(raft::handle_t& handle,
                  std::vector<MLCommon::Matrix::Data<int>*>* out,
                  std::vector<std::vector<float*>>* probas,
                  std::vector<MLCommon::Matrix::floatData_t*>& idx_data,
                  MLCommon::Matrix::PartDescriptor& idx_desc,
                  std::vector<MLCommon::Matrix::floatData_t*>& query_data,
                  MLCommon::Matrix::PartDescriptor& query_desc,
                  std::vector<std::vector<int*>>& y,
                  std::vector<int*>& uniq_labels,
                  std::vector<int>& n_unique,
                  bool rowMajorIndex = false,
                  bool rowMajorQuery = false,
                  bool probas_only   = false,
                  int k              = 10,
                  size_t batch_size  = 1 << 15,
                  bool verbose       = false);
 
void knn_regress(raft::handle_t& handle,
                 std::vector<MLCommon::Matrix::Data<float>*>* out,
                 std::vector<MLCommon::Matrix::floatData_t*>& idx_data,
                 MLCommon::Matrix::PartDescriptor& idx_desc,
                 std::vector<MLCommon::Matrix::floatData_t*>& query_data,
                 MLCommon::Matrix::PartDescriptor& query_desc,
                 std::vector<std::vector<float*>>& y,
                 bool rowMajorIndex,
                 bool rowMajorQuery,
                 int k,
                 int n_outputs,
                 size_t batch_size,
                 bool verbose);
 
};  // END namespace opg
};  // namespace KNN
};  // namespace ML