Implementation:Rapidsai Cuml KNN C API

Knowledge Sources	Rapidsai_Cuml
Domains	Machine_Learning, Nearest_Neighbors
Last Updated	2026-02-08 12:00 GMT

Overview

Provides a C-compatible API wrapper for GPU-accelerated brute-force k-nearest neighbors search, enabling interoperability with C code and foreign function interfaces.

Description

The knn_api.h header declares a single C-linkage function, knn_search, that wraps the cuML C++ brute-force KNN implementation. This function is designed for use by C callers or through foreign function interfaces (e.g., Python ctypes, Cython) where C++ name mangling would be problematic.

The function performs a brute-force KNN search across an array of input partitions and writes the resulting neighbor indices and distances into pre-allocated output buffers. It accepts the same set of parameters as the C++ brute_force_knn function but uses C-compatible types: raw pointer arrays instead of std::vector, an integer metric_type enum value instead of the C++ enum class, and a cumlHandle_t opaque handle instead of raft::handle_t.

The function returns a cumlError_t error code indicating success or the nature of any failure.

Usage

Use this C API when integrating cuML's KNN functionality into non-C++ environments such as Python bindings, C libraries, or other language FFIs. For pure C++ usage, prefer the C++ brute_force_knn function from knn.hpp.

Code Reference

Source Location

• Repository: Rapidsai_Cuml
• File: cpp/include/cuml/neighbors/knn_api.h

Signature

cumlError_t knn_search(const cumlHandle_t handle,
                       float** input,
                       int* size,
                       int n_params,
                       int D,
                       float* search_items,
                       int n,
                       int64_t* res_I,
                       float* res_D,
                       int k,
                       bool rowMajorIndex,
                       bool rowMajorQuery,
                       int metric_type,
                       float metric_arg,
                       bool expanded);

Import

#include <cuml/neighbors/knn_api.h>

I/O Contract

Inputs

Name	Type	Required	Description
handle	cumlHandle_t	Yes	The cuML handle for GPU resource management
input	float**	Yes	Array of pointers to input (index) arrays on device
size	int*	Yes	Array of row counts for each input partition
n_params	int	Yes	Number of input partitions (length of input and size arrays)
D	int	Yes	Dimensionality of the feature vectors
search_items	float*	Yes	Device pointer to query vectors of shape [n x D]
n	int	Yes	Number of query rows
k	int	Yes	Number of nearest neighbors to return
rowMajorIndex	bool	Yes	Whether index arrays are in row-major layout
rowMajorQuery	bool	Yes	Whether query array is in row-major layout
metric_type	int	Yes	Integer value from cuml::ML::distance::DistanceType enum specifying the distance metric
metric_arg	float	Yes	The value of p for Minkowski distances (ignored for non-Minkowski metrics)
expanded	bool	Yes	Whether lp-based distances should be returned in expanded form

Outputs

Name	Type	Description
res_I	int64_t*	Device array of resulting neighbor indices, size [n x k]
res_D	float*	Device array of resulting neighbor distances, size [n x k]
return value	cumlError_t	CUML_SUCCESS on success, or an error flag on failure

Usage Examples

#include <cuml/neighbors/knn_api.h>

// Assume cuML handle is already created
cumlHandle_t handle;

// Two index partitions on device
float* index_parts[2] = {d_part1, d_part2};
int sizes[2] = {5000, 3000};
int n_parts = 2;
int D = 64;

// Query data
float* d_queries;  // device pointer [100 x 64]
int n = 100;
int k = 5;

// Output buffers (pre-allocated on device)
int64_t* d_indices;  // [100 x 5]
float* d_distances;  // [100 x 5]

cumlError_t status = knn_search(handle, index_parts, sizes, n_parts, D,
                                d_queries, n, d_indices, d_distances, k,
                                true,   // rowMajorIndex
                                true,   // rowMajorQuery
                                1,      // metric_type (L2)
                                2.0f,   // metric_arg
                                false); // expanded

if (status != CUML_SUCCESS) {
    // Handle error
}

Related Pages

• Environment:Rapidsai_Cuml_CUDA_GPU