Add pairwise_distance api's for C, Python and Rust (#142)

Authors:
  - Ben Frederickson (https://github.com/benfred)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: rapidsai/cuvs#142

cpp/CMakeLists.txt

@@ -545,8 +545,9 @@ target_link_options(cuvs PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")
 # * cuvs_c -------------------------------------------------------------------------------
 if(BUILD_C_LIBRARY)
   add_library(
-    cuvs_c SHARED src/core/c_api.cpp src/neighbors/brute_force_c.cpp src/neighbors/ivf_flat_c.cpp
-                  src/neighbors/ivf_pq_c.cpp src/neighbors/cagra_c.cpp
+    cuvs_c SHARED
+    src/core/c_api.cpp src/neighbors/brute_force_c.cpp src/neighbors/ivf_flat_c.cpp
+    src/neighbors/ivf_pq_c.cpp src/neighbors/cagra_c.cpp src/distance/pairwise_distance_c.cpp
   )
 
   add_library(cuvs::c_api ALIAS cuvs_c)

cpp/include/cuvs/distance/pairwise_distance.h

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 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 <dlpack/dlpack.h>
+
+#include <cuvs/core/c_api.h>
+#include <cuvs/distance/distance.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Compute pairwise distances for two matrices
+ *
+ *
+ * Usage example:
+ * @code{.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/distance/pairwise_distance.h>
+ *
+ * // Create cuvsResources_t
+ * cuvsResources_t res;
+ * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
+ *
+ * // Assume a populated `DLManagedTensor` type here
+ * DLManagedTensor x;
+ * DLManagedTensor y;
+ * DLManagedTensor dist;
+ *
+ * cuvsPairwiseDistance(res, &x, &y, &dist, L2SqrtUnexpanded, 2.0);
+ * @endcode
+ *
+ * @param[in] res cuvs resources object for managing expensive resources
+ * @param[in] x first set of points (size n*k)
+ * @param[in] y second set of points (size m*k)
+ * @param[out] dist output distance matrix (size n*m)
+ * @param[in] metric distance to evaluate
+ * @param[in] metric_arg metric argument (used for Minkowski distance)
+ */
+cuvsError_t cuvsPairwiseDistance(cuvsResources_t res,
+                                 DLManagedTensor* x,
+                                 DLManagedTensor* y,
+                                 DLManagedTensor* dist,
+                                 cuvsDistanceType metric,
+                                 float metric_arg);
+#ifdef __cplusplus
+}
+#endif

cpp/include/cuvs/neighbors/cagra.h

@@ -356,7 +356,8 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  * cuvsError_t params_create_status = cuvsCagraSearchParamsCreate(&params);
  *
  * // Search the `index` built using `cuvsCagraBuild`
- * cuvsError_t search_status = cuvsCagraSearch(res, params, index, queries, neighbors, distances);
+ * cuvsError_t search_status = cuvsCagraSearch(res, params, index, &queries, &neighbors,
+ * &distances);
  *
  * // de-allocate `params` and `res`
  * cuvsError_t params_destroy_status = cuvsCagraSearchParamsDestroy(params);

cpp/src/distance/pairwise_distance_c.cpp

@@ -0,0 +1,81 @@
+
+/*
+ * Copyright (c) 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.
+ */
+
+#include <cstdint>
+#include <dlpack/dlpack.h>
+
+#include <raft/core/error.hpp>
+#include <raft/core/mdspan_types.hpp>
+#include <raft/core/resources.hpp>
+
+#include <cuvs/core/c_api.h>
+#include <cuvs/core/exceptions.hpp>
+#include <cuvs/core/interop.hpp>
+#include <cuvs/distance/distance.hpp>
+
+namespace {
+
+template <typename T>
+void _pairwise_distance(cuvsResources_t res,
+                        DLManagedTensor* x_tensor,
+                        DLManagedTensor* y_tensor,
+                        DLManagedTensor* distances_tensor,
+                        cuvsDistanceType metric,
+                        float metric_arg)
+{
+  auto res_ptr = reinterpret_cast<raft::resources*>(res);
+
+  using mdspan_type           = raft::device_matrix_view<T const, int64_t, raft::row_major>;
+  using distances_mdspan_type = raft::device_matrix_view<T, int64_t, raft::row_major>;
+
+  auto x_mds         = cuvs::core::from_dlpack<mdspan_type>(x_tensor);
+  auto y_mds         = cuvs::core::from_dlpack<mdspan_type>(y_tensor);
+  auto distances_mds = cuvs::core::from_dlpack<distances_mdspan_type>(distances_tensor);
+
+  cuvs::distance::pairwise_distance(*res_ptr, x_mds, y_mds, distances_mds, metric, metric_arg);
+}
+}  // namespace
+
+extern "C" cuvsError_t cuvsPairwiseDistance(cuvsResources_t res,
+                                            DLManagedTensor* x_tensor,
+                                            DLManagedTensor* y_tensor,
+                                            DLManagedTensor* distances_tensor,
+                                            cuvsDistanceType metric,
+                                            float metric_arg)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto x_dt    = x_tensor->dl_tensor.dtype;
+    auto y_dt    = x_tensor->dl_tensor.dtype;
+    auto dist_dt = x_tensor->dl_tensor.dtype;
+
+    if ((x_dt.code != kDLFloat) || (y_dt.code != kDLFloat) || (dist_dt.code != kDLFloat)) {
+      RAFT_FAIL("Inputs to cuvsPairwiseDistance must all be floating point tensors");
+    }
+
+    if ((x_dt.bits != y_dt.bits) || (x_dt.bits != dist_dt.bits)) {
+      RAFT_FAIL("Inputs to cuvsPairwiseDistance must all have the same dtype");
+    }
+
+    if (x_dt.bits == 32) {
+      _pairwise_distance<float>(res, x_tensor, y_tensor, distances_tensor, metric, metric_arg);
+    } else if (x_dt.bits == 64) {
+      _pairwise_distance<double>(res, x_tensor, y_tensor, distances_tensor, metric, metric_arg);
+    } else {
+      RAFT_FAIL("Unsupported DLtensor dtype: %d and bits: %d", x_dt.code, x_dt.bits);
+    }
+  });
+}

cpp/test/CMakeLists.txt

@@ -178,6 +178,10 @@ endif()
 
 if(BUILD_C_TESTS)
   ConfigureTest(NAME INTEROP_TEST PATH test/core/interop.cu C_LIB)
+  ConfigureTest(
+    NAME DISTANCE_C_TEST PATH test/distance/run_pairwise_distance_c.c
+    test/distance/pairwise_distance_c.cu C_LIB
+  )
 
   ConfigureTest(
     NAME BRUTEFORCE_C_TEST PATH test/neighbors/run_brute_force_c.c test/neighbors/brute_force_c.cu

cpp/test/distance/pairwise_distance_c.cu

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) 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.
+ */
+
+#include <cuda.h>
+
+#include <gtest/gtest.h>
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/handle.hpp>
+#include <raft/random/rng.cuh>
+
+#include <cuvs/distance/distance.h>
+
+extern "C" void run_pairwise_distance(int64_t n_rows,
+                                      int64_t n_queries,
+                                      int64_t n_dim,
+                                      float* index_data,
+                                      float* query_data,
+                                      float* distances_data,
+                                      cuvsDistanceType metric);
+
+template <typename T>
+void generate_random_data(T* devPtr, size_t size)
+{
+  raft::handle_t handle;
+  raft::random::RngState r(1234ULL);
+  raft::random::uniform(handle, r, devPtr, size, T(0.1), T(2.0));
+};
+
+TEST(PairwiseDistanceC, Distance)
+{
+  int64_t n_rows    = 8096;
+  int64_t n_queries = 128;
+  int64_t n_dim     = 32;
+
+  cuvsDistanceType metric = L2Expanded;
+
+  float *index_data, *query_data, *distances_data;
+  cudaMalloc(&index_data, sizeof(float) * n_rows * n_dim);
+  cudaMalloc(&query_data, sizeof(float) * n_queries * n_dim);
+  cudaMalloc(&distances_data, sizeof(float) * n_queries * n_rows);
+
+  generate_random_data(index_data, n_rows * n_dim);
+  generate_random_data(query_data, n_queries * n_dim);
+
+  run_pairwise_distance(n_rows, n_queries, n_dim, index_data, query_data, distances_data, metric);
+
+  // delete device memory
+  cudaFree(index_data);
+  cudaFree(query_data);
+  cudaFree(distances_data);
+}

cpp/test/distance/run_pairwise_distance_c.c

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 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.
+ */
+#include <cuvs/distance/pairwise_distance.h>
+#include <stdint.h>
+
+void run_pairwise_distance(int64_t n_rows,
+                           int64_t n_queries,
+                           int64_t n_dim,
+                           float* index_data,
+                           float* query_data,
+                           float* distances_data,
+                           int64_t* neighbors_data,
+                           cuvsDistanceType metric)
+{
+  // create cuvsResources_t
+  cuvsResources_t res;
+  cuvsResourcesCreate(&res);
+
+  // create dataset DLTensor
+  DLManagedTensor dataset_tensor;
+  dataset_tensor.dl_tensor.data               = index_data;
+  dataset_tensor.dl_tensor.device.device_type = kDLCUDA;
+  dataset_tensor.dl_tensor.ndim               = 2;
+  dataset_tensor.dl_tensor.dtype.code         = kDLFloat;
+  dataset_tensor.dl_tensor.dtype.bits         = 32;
+  dataset_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t dataset_shape[2]                    = {n_rows, n_dim};
+  dataset_tensor.dl_tensor.shape              = dataset_shape;
+  dataset_tensor.dl_tensor.strides            = NULL;
+
+  // create queries DLTensor
+  DLManagedTensor queries_tensor;
+  queries_tensor.dl_tensor.data               = (void*)query_data;
+  queries_tensor.dl_tensor.device.device_type = kDLCUDA;
+  queries_tensor.dl_tensor.ndim               = 2;
+  queries_tensor.dl_tensor.dtype.code         = kDLFloat;
+  queries_tensor.dl_tensor.dtype.bits         = 32;
+  queries_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t queries_shape[2]                    = {n_queries, n_dim};
+  queries_tensor.dl_tensor.shape              = queries_shape;
+  queries_tensor.dl_tensor.strides            = NULL;
+
+  // create distances DLTensor
+  DLManagedTensor distances_tensor;
+  distances_tensor.dl_tensor.data               = (void*)distances_data;
+  distances_tensor.dl_tensor.device.device_type = kDLCUDA;
+  distances_tensor.dl_tensor.ndim               = 2;
+  distances_tensor.dl_tensor.dtype.code         = kDLFloat;
+  distances_tensor.dl_tensor.dtype.bits         = 32;
+  distances_tensor.dl_tensor.dtype.lanes        = 1;
+  int64_t distances_shape[2]                    = {n_rows, n_queries};
+  distances_tensor.dl_tensor.shape              = distances_shape;
+  distances_tensor.dl_tensor.strides            = NULL;
+
+  // run pairwise distances
+  cuvsPairwiseDistance(res, &dataset_tensor, &queries_tensor, &distances_tensor, metric, 2.0);
+
+  cuvsResourcesDestroy(res);
+}

docs/source/cpp_api/distance.rst

@@ -11,7 +11,7 @@ distances have been highly optimized and support a wide assortment of different
 Distance Types
 --------------
 
-``#include <cuvs/distance/distance_types.hpp>``
+``#include <cuvs/distance/distance.hpp>``
 
 namespace *cuvs::distance*
 
@@ -22,11 +22,11 @@ namespace *cuvs::distance*
 Pairwise Distances
 ------------------
 
-``include <cuvs/distance/pairwise_distance.hpp>``
+``include <cuvs/distance/distance.hpp>``
 
 namespace *cuvs::distance*
 
 .. doxygengroup:: pairwise_distance
     :project: cuvs
     :members:
-    :content-only:
+    :content-only:

python/cuvs/CMakeLists.txt

@@ -81,6 +81,7 @@ endif()
 rapids_cython_init()
 
 add_subdirectory(cuvs/common)
+add_subdirectory(cuvs/distance)
 add_subdirectory(cuvs/neighbors)
 
 if(DEFINED cython_lib_dir)

python/cuvs/cuvs/distance/CMakeLists.txt

@@ -0,0 +1,24 @@
+# =============================================================================
+# Copyright (c) 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.
+# =============================================================================
+
+# Set the list of Cython files to build
+set(cython_sources distance.pyx)
+set(linked_libraries cuvs::cuvs cuvs::c_api)
+
+# Build all of the Cython targets
+rapids_cython_create_modules(
+  CXX
+  SOURCE_FILES "${cython_sources}"
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX distance_
+)

python/cuvs/cuvs/distance/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 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.
+
+from .distance import DISTANCE_TYPES, pairwise_distance
+
+__all__ = ["DISTANCE_TYPES", "pairwise_distance"]