IVF-PQ Python wrappers

README.md

@@ -77,11 +77,73 @@ auto metric = raft::distance::DistanceType::L2SqrtExpanded;
 raft::distance::pairwise_distance(handle, input.view(), input.view(), output.view(), metric);
 ```
 
+It's also possible to create `raft::device_mdspan` views to invoke the same API with raw pointers and shape information:
+
+```c++
+#include <raft/core/handle.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/random/make_blobs.cuh>
+#include <raft/distance/distance.cuh>
+
+raft::handle_t handle;
+
+int n_samples = 5000;
+int n_features = 50;
+
+float *input;
+int *labels;
+float *output;
+
+...
+// Allocate input, labels, and output pointers
+...
+
+auto input_view = raft::make_device_matrix_view(input, n_samples, n_features);
+auto labels_view = raft::make_device_vector_view(labels, n_samples);
+auto output_view = raft::make_device_matrix_view(output, n_samples, n_samples);
+
+raft::random::make_blobs(handle, input_view, labels_view);
+
+auto metric = raft::distance::DistanceType::L2SqrtExpanded;
+raft::distance::pairwise_distance(handle, input_view, input_view, output_view, metric);
+```
+
+
 ### Python Example
 
 The `pylibraft` package contains a Python API for RAFT algorithms and primitives. `pylibraft` integrates nicely into other libraries by being very lightweight with minimal dependencies and accepting any object that supports the `__cuda_array_interface__`, such as [CuPy's ndarray](https://docs.cupy.dev/en/stable/user_guide/interoperability.html#rmm). The number of RAFT algorithms exposed in this package is continuing to grow from release to release.
 
-The example below demonstrates computing the pairwise Euclidean distances between CuPy arrays. `pylibraft` is a low-level API that prioritizes efficiency and simplicity over being pythonic, which is shown here by pre-allocating the output memory before invoking the `pairwise_distance` function. Note that CuPy is not a required dependency for `pylibraft`.
+The example below demonstrates computing the pairwise Euclidean distances between CuPy arrays. Note that CuPy is not a required dependency for `pylibraft`.
+
+```python
+import cupy as cp
+
+from pylibraft.distance import pairwise_distance
+
+n_samples = 5000
+n_features = 50
+
+in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
+in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
+
+output = pairwise_distance(in1, in2, metric="euclidean")
+```
+
+The `output` array supports [__cuda_array_interface__](https://numba.pydata.org/numba-doc/dev/cuda/cuda_array_interface.html#cuda-array-interface-version-2) so it is interoperable with other libraries like CuPy, Numba, and PyTorch that also support it. 
+
+Below is an example of converting the output `pylibraft.device_ndarray` to a CuPy array:
+```python
+cupy_array = cp.asarray(output)
+```
+
+And converting to a PyTorch tensor:
+```python
+import torch
+
+torch_tensor = torch.as_tensor(output, device='cuda')
+```
+
+`pylibraft` also supports writing to a pre-allocated output array so any `__cuda_array_interface__` supported array can be written to in-place:
 
 ```python
 import cupy as cp
@@ -95,9 +157,10 @@ in1 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
 in2 = cp.random.random_sample((n_samples, n_features), dtype=cp.float32)
 output = cp.empty((n_samples, n_samples), dtype=cp.float32)
 
-pairwise_distance(in1, in2, output, metric="euclidean")
+pairwise_distance(in1, in2, out=output, metric="euclidean")
 ```
 
+
 ## Installing
 
 RAFT itself can be installed through conda, [Cmake Package Manager (CPM)](https://github.com/cpm-cmake/CPM.cmake), or by building the repository from source. Please refer to the [build instructions](docs/source/build.md) for more a comprehensive guide on building RAFT and using it in downstream projects.

cpp/CMakeLists.txt

@@ -332,6 +332,21 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/specializations/fused_l2_nn_double_int64.cu
     src/distance/specializations/fused_l2_nn_float_int.cu
     src/distance/specializations/fused_l2_nn_float_int64.cu
+    src/nn/specializations/detail/ivfpq_build.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_float_fast.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_float_no_basediff.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_float_no_smem_lut.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_fast.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_basediff.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_smem_lut.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_fast.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_basediff.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_smem_lut.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_half_fast.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_half_no_basediff.cu
+    src/nn/specializations/detail/ivfpq_compute_similarity_half_no_smem_lut.cu
+    src/nn/specializations/detail/ivfpq_search.cu
+    src/nn/specializations/detail/ivfpq_search_float_uint64_t.cu
     src/random/specializations/rmat_rectangular_generator_int_double.cu
     src/random/specializations/rmat_rectangular_generator_int64_double.cu
     src/random/specializations/rmat_rectangular_generator_int_float.cu
@@ -400,6 +415,8 @@ if(RAFT_COMPILE_NN_LIBRARY)
     src/nn/specializations/detail/ivfpq_compute_similarity_half_fast.cu
     src/nn/specializations/detail/ivfpq_compute_similarity_half_no_basediff.cu
     src/nn/specializations/detail/ivfpq_compute_similarity_half_no_smem_lut.cu
+    src/nn/specializations/detail/ivfpq_build.cu
+    src/nn/specializations/detail/ivfpq_search.cu
     src/nn/specializations/detail/ivfpq_search_float_int64_t.cu
     src/nn/specializations/detail/ivfpq_search_float_uint32_t.cu
     src/nn/specializations/detail/ivfpq_search_float_uint64_t.cu

cpp/include/raft/neighbors/specializations/ivf_pq_specialization.hpp

@@ -0,0 +1,77 @@
+/*
+ * Copyright (c) 2022, 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 <raft/neighbors/ivf_pq_types.hpp>
+
+namespace raft::neighbors ::ivf_pq {
+
+#define RAFT_INST_SEARCH(T, IdxT)   \
+  void search(const handle_t&,      \
+              const search_params&, \
+              const index<IdxT>&,   \
+              const T*,             \
+              uint32_t,             \
+              uint32_t,             \
+              IdxT*,                \
+              float*,               \
+              rmm::mr::device_memory_resource*);
+
+RAFT_INST_SEARCH(float, uint64_t);
+RAFT_INST_SEARCH(int8_t, uint64_t);
+RAFT_INST_SEARCH(uint8_t, uint64_t);
+
+#undef RAFT_INST_SEARCH
+
+// We define overloads for build and extend with void return type. This is used in the Cython
+// wrappers, where exception handling is not compatible with return type that has nontrivial
+// constructor.
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)      \
+  auto build(const handle_t& handle,         \
+             const index_params& params,     \
+             const T* dataset,               \
+             IdxT n_rows,                    \
+             uint32_t dim)                   \
+    ->index<IdxT>;                           \
+                                             \
+  auto extend(const handle_t& handle,        \
+              const index<IdxT>& orig_index, \
+              const T* new_vectors,          \
+              const IdxT* new_indices,       \
+              IdxT n_rows)                   \
+    ->index<IdxT>;                           \
+                                             \
+  void build(const handle_t& handle,         \
+             const index_params& params,     \
+             const T* dataset,               \
+             IdxT n_rows,                    \
+             uint32_t dim,                   \
+             index<IdxT>* idx);              \
+                                             \
+  void extend(const handle_t& handle,        \
+              index<IdxT>* idx,              \
+              const T* new_vectors,          \
+              const IdxT* new_indices,       \
+              IdxT n_rows);
+
+RAFT_INST_BUILD_EXTEND(float, uint64_t)
+RAFT_INST_BUILD_EXTEND(int8_t, uint64_t)
+RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t)
+
+#undef RAFT_INST_BUILD_EXTEND
+
+}  // namespace raft::neighbors::ivf_pq

cpp/include/raft/spatial/knn/detail/haversine_distance.cuh

@@ -27,7 +27,6 @@
 
 #include <raft/core/handle.hpp>
 #include <raft/distance/distance_types.hpp>
-#include <raft/spatial/knn/faiss_mr.hpp>
 
 namespace raft {
 namespace spatial {

cpp/src/nn/specializations/detail/ivfpq_build.cu

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) 2022, 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 <raft/neighbors/ivf_pq.cuh>
+#include <raft/neighbors/specializations/ivf_pq_specialization.hpp>
+
+namespace raft::neighbors::ivf_pq {
+
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                           \
+  auto build(const handle_t& handle,                                              \
+             const index_params& params,                                          \
+             const T* dataset,                                                    \
+             IdxT n_rows,                                                         \
+             uint32_t dim)                                                        \
+    ->index<IdxT>                                                                 \
+  {                                                                               \
+    return build<T, IdxT>(handle, params, dataset, n_rows, dim);                  \
+  }                                                                               \
+  auto extend(const handle_t& handle,                                             \
+              const index<IdxT>& orig_index,                                      \
+              const T* new_vectors,                                               \
+              const IdxT* new_indices,                                            \
+              IdxT n_rows)                                                        \
+    ->index<IdxT>                                                                 \
+  {                                                                               \
+    return extend<T, IdxT>(handle, orig_index, new_vectors, new_indices, n_rows); \
+  }                                                                               \
+                                                                                  \
+  void build(const handle_t& handle,                                              \
+             const index_params& params,                                          \
+             const T* dataset,                                                    \
+             IdxT n_rows,                                                         \
+             uint32_t dim,                                                        \
+             index<IdxT>* idx)                                                    \
+  {                                                                               \
+    *idx = build<T, IdxT>(handle, params, dataset, n_rows, dim);                  \
+  }                                                                               \
+  void extend(const handle_t& handle,                                             \
+              index<IdxT>* idx,                                                   \
+              const T* new_vectors,                                               \
+              const IdxT* new_indices,                                            \
+              IdxT n_rows)                                                        \
+  {                                                                               \
+    extend<T, IdxT>(handle, idx, new_vectors, new_indices, n_rows);               \
+  }
+
+RAFT_INST_BUILD_EXTEND(float, uint64_t);
+RAFT_INST_BUILD_EXTEND(int8_t, uint64_t);
+RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t);
+
+#undef RAFT_INST_BUILD_EXTEND
+
+}  // namespace raft::neighbors::ivf_pq

cpp/src/nn/specializations/detail/ivfpq_search.cu

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2022, 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 <raft/neighbors/ivf_pq.cuh>
+#include <raft/neighbors/specializations/detail/ivf_pq_search.cuh>
+#include <raft/neighbors/specializations/ivf_pq_specialization.hpp>
+
+namespace raft::neighbors::ivf_pq {
+
+#define RAFT_SEARCH_INST(T, IdxT)                                                          \
+  void search(const handle_t& handle,                                                      \
+              const search_params& params,                                                 \
+              const index<IdxT>& idx,                                                      \
+              const T* queries,                                                            \
+              uint32_t n_queries,                                                          \
+              uint32_t k,                                                                  \
+              IdxT* neighbors,                                                             \
+              float* distances,                                                            \
+              rmm::mr::device_memory_resource* mr)                                         \
+  {                                                                                        \
+    search<T, IdxT>(handle, params, idx, queries, n_queries, k, neighbors, distances, mr); \
+  }
+
+RAFT_SEARCH_INST(float, uint64_t);
+RAFT_SEARCH_INST(int8_t, uint64_t);
+RAFT_SEARCH_INST(uint8_t, uint64_t);
+
+#undef RAFT_INST_SEARCH
+
+}  // namespace raft::neighbors::ivf_pq

python/pylibraft/CMakeLists.txt

@@ -75,6 +75,7 @@ rapids_cython_init()
 
 add_subdirectory(pylibraft/common)
 add_subdirectory(pylibraft/distance)
+add_subdirectory(pylibraft/neighbors)
 add_subdirectory(pylibraft/random)
 add_subdirectory(pylibraft/cluster)
 

python/pylibraft/pylibraft/common/__init__.py

@@ -13,5 +13,7 @@
 # limitations under the License.
 #
 
+
 from .cuda import Stream
+from .device_ndarray import device_ndarray
 from .handle import Handle