knn_mg.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2024, 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 <cumlprims/opg/matrix/data.hpp>
#include <cumlprims/opg/matrix/part_descriptor.hpp>
#include <raft/core/handle.hpp>
 
#include <vector>
 
namespace ML {
namespace KNN {
namespace opg {
 
using namespace MLCommon;
 
void knn(raft::handle_t& handle,
         std::vector<Matrix::Data<int64_t>*>* out_I,
         std::vector<Matrix::floatData_t*>* out_D,
         std::vector<Matrix::floatData_t*>& idx_data,
         Matrix::PartDescriptor& idx_desc,
         std::vector<Matrix::floatData_t*>& query_data,
         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<Matrix::Data<int>*>* out,
                  std::vector<std::vector<float*>>* probas,
                  std::vector<Matrix::floatData_t*>& idx_data,
                  Matrix::PartDescriptor& idx_desc,
                  std::vector<Matrix::floatData_t*>& query_data,
                  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<Matrix::Data<float>*>* out,
                 std::vector<Matrix::floatData_t*>& idx_data,
                 Matrix::PartDescriptor& idx_desc,
                 std::vector<Matrix::floatData_t*>& query_data,
                 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