hipvs/distance/
mod.rs

1/*
2 * Copyright (c) 2024-2025, NVIDIA CORPORATION.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::distance_type::DistanceType;
18use crate::dlpack::ManagedTensor;
19use crate::error::{check_cuvs, Result};
20use crate::resources::Resources;
21
22/// Compute pairwise distances between X and Y
23///
24/// # Arguments
25///
26/// * `res` - Resources to use
27/// * `x` - A matrix in device memory - shape (m, k)
28/// * `y` - A matrix in device memory - shape (n, k)
29/// * `distances` - A matrix in device memory that receives the output distances - shape (m, n)
30/// * `metric` - DistanceType to use for building the index
31/// * `metric_arg` - Optional value of `p` for Minkowski distances
32pub fn pairwise_distance(
33    res: &Resources,
34    x: &ManagedTensor,
35    y: &ManagedTensor,
36    distances: &ManagedTensor,
37    metric: DistanceType,
38    metric_arg: Option<f32>,
39) -> Result<()> {
40    unsafe {
41        check_cuvs(ffi::cuvsPairwiseDistance(
42            res.0,
43            x.as_ptr(),
44            y.as_ptr(),
45            distances.as_ptr(),
46            metric,
47            metric_arg.unwrap_or(2.0),
48        ))
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55    use ndarray_rand::rand_distr::Uniform;
56    use ndarray_rand::RandomExt;
57
58    #[test]
59    fn test_pairwise_distance() {
60        let res = Resources::new().unwrap();
61
62        // Create a new random dataset to index
63        let n_datapoints = 256;
64        let n_features = 16;
65        let dataset =
66            ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));
67        let dataset_device = ManagedTensor::from(&dataset).to_device(&res).unwrap();
68
69        let mut distances_host = ndarray::Array::<f32, _>::zeros((n_datapoints, n_datapoints));
70        let distances = ManagedTensor::from(&distances_host)
71            .to_device(&res)
72            .unwrap();
73
74        pairwise_distance(
75            &res,
76            &dataset_device,
77            &dataset_device,
78            &distances,
79            DistanceType::L2Expanded,
80            None,
81        )
82        .unwrap();
83
84        // Copy back to host memory
85        distances.to_host(&res, &mut distances_host).unwrap();
86
87        // Self distance should be 0
88        assert_eq!(distances_host[[0, 0]], 0.0);
89    }
90}
hipvs/distance/mod.rs

hipvs/distance/
mod.rs
