Skip to content

Commit

Permalink
Merge pull request #150 from rapidsai/branch-24.06
Browse files Browse the repository at this point in the history
Forward-merge branch-24.06 into branch-24.08
  • Loading branch information
GPUtester authored May 23, 2024
2 parents a05fc3c + 4b260aa commit dee0d0f
Show file tree
Hide file tree
Showing 22 changed files with 677 additions and 10 deletions.
5 changes: 3 additions & 2 deletions cpp/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
62 changes: 62 additions & 0 deletions cpp/include/cuvs/distance/pairwise_distance.h
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion cpp/include/cuvs/neighbors/cagra.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
81 changes: 81 additions & 0 deletions cpp/src/distance/pairwise_distance_c.cpp
Original file line number Diff line number Diff line change
@@ -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);
}
});
}
4 changes: 4 additions & 0 deletions cpp/test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
64 changes: 64 additions & 0 deletions cpp/test/distance/pairwise_distance_c.cu
Original file line number Diff line number Diff line change
@@ -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);
}
72 changes: 72 additions & 0 deletions cpp/test/distance/run_pairwise_distance_c.c
Original file line number Diff line number Diff line change
@@ -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);
}
6 changes: 3 additions & 3 deletions docs/source/cpp_api/distance.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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*

Expand All @@ -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:
1 change: 1 addition & 0 deletions python/cuvs/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
24 changes: 24 additions & 0 deletions python/cuvs/cuvs/distance/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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_
)
Empty file.
17 changes: 17 additions & 0 deletions python/cuvs/cuvs/distance/__init__.py
Original file line number Diff line number Diff line change
@@ -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"]
Loading

0 comments on commit dee0d0f

Please sign in to comment.