Moving raft::spatial::knn -> raft::neighbors

BUILD.md

@@ -101,7 +101,7 @@ For example, to run the distance tests:
 It can take sometime to compile all of the tests. You can build individual tests by providing a semicolon-separated list to the `--limit-tests` option in `build.sh`:
 
 ```bash
-./build.sh libraft tests --limit-tests=SPATIAL_TEST;DISTANCE_TEST;MATRIX_TEST
+./build.sh libraft tests --limit-tests=NEIGHBORS_TEST;DISTANCE_TEST;MATRIX_TEST
 ```
 
 ### <a id="gbench"></a>Benchmarks
@@ -111,10 +111,10 @@ The benchmarks are broken apart by algorithm category, so you will find several
 ./build.sh libraft bench
 ```
 
-It can take sometime to compile all of the tests. You can build individual tests by providing a semicolon-separated list to the `--limit-tests` option in `build.sh`:
+It can take sometime to compile all of the benchmarks. You can build individual benchmarks by providing a semicolon-separated list to the `--limit-bench` option in `build.sh`:
 
 ```bash
-./build.sh libraft bench --limit-bench=SPATIAL_BENCH;DISTANCE_BENCH;LINALG_BENCH
+./build.sh libraft bench --limit-bench=NEIGHBORS_BENCH;DISTANCE_BENCH;LINALG_BENCH
 ```
 
 ### <a id="cpp_using_cmake"></a>C++ Using Cmake

README.md

@@ -12,19 +12,19 @@ While not exhaustive, the following general categories help summarize the accele
 | Category | Examples |
 | --- | --- |
 | **Data Formats** | sparse & dense, conversions, data generation |
-| **Dense Linear Algebra** | matrix arithmetic, norms, factorization, least squares, svd & eigenvalue problems |
+| **Dense Operations** | linear algebra, matrix and vector operations, slicing, norms, factorization, least squares, svd & eigenvalue problems |
+| **Sparse Operations** | linear algebra, eigenvalue problems, slicing, symmetrization, components & labeling |
 | **Spatial** | pairwise distances, nearest neighbors, neighborhood graph construction |
-| **Sparse Operations** | linear algebra, eigenvalue problems, slicing, symmetrization, labeling |
 | **Basic Clustering** | spectral clustering, hierarchical clustering, k-means |
 | **Solvers** | combinatorial optimization, iterative solvers |
 | **Statistics** | sampling, moments and summary statistics, metrics |
-| **Distributed Tools** | multi-node multi-gpu infrastructure |
+| **Tools & Utilities** | common utilities for developing CUDA applications, multi-node multi-gpu infrastructure |
 
 RAFT provides a header-only C++ library and pre-compiled shared libraries that can 1) speed up compile times and 2) enable the APIs to be used without CUDA-enabled compilers.
 
-RAFT also provides 2 Python libraries:
-- `pylibraft` - low-level Python wrappers around RAFT algorithms and primitives.
-- `raft-dask` - reusable infrastructure for building analytics, including tools for building both single-GPU and multi-node multi-GPU algorithms.
+In addition to the C++ library, RAFT also provides 2 Python libraries:
+- `pylibraft` - lightweight low-level Python wrappers around RAFT algorithms and primitives.
+- `raft-dask` - multi-node multi-GPU communicator infrastructure for building distributed algorithms on the GPU with Dask.
 
 ## Getting started
 
@@ -78,9 +78,9 @@ raft::distance::pairwise_distance(handle, input.view(), input.view(), output.vie
 
 ### Python Example
 
-The `pylibraft` package contains a Python API for RAFT algorithms and primitives. The package is currently limited to pairwise distances, and we will continue adding more.
+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 package is currently limited to pairwise distances and RMAT graph generation, but we will continue adding more in future releases.
 
-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.
+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`.
 
 ```python
 import cupy as cp
@@ -107,7 +107,7 @@ The easiest way to install RAFT is through conda and several packages are provid
 - `libraft-headers` RAFT headers
 - `libraft-nn` (optional) contains shared libraries for the nearest neighbors primitives.
 - `libraft-distance` (optional) contains shared libraries for distance primitives.
-- `pylibraft` (optional) Python wrappers around RAFT algorithms and primitives
+- `pylibraft` (optional) Python wrappers around RAFT algorithms and primitives.
 - `raft-dask` (optional) enables deployment of multi-node multi-GPU algorithms that use RAFT `raft::comms` in Dask clusters.
 
 Use the following command to install all of the RAFT packages with conda (replace `rapidsai` with `rapidsai-nightly` to install more up-to-date but less stable nightly packages). `mamba` is preferred over the `conda` command.
@@ -198,7 +198,25 @@ The folder structure mirrors other RAPIDS repos, with the following folders:
   - `bench`: Benchmarks source code
   - `cmake`: Cmake modules and templates
   - `doxygen`: Doxygen configuration
-  - `include`: The C++ API headers are fully-contained here
+  - `include`: The C++ API headers are fully-contained here (deprecated directories are excluded from the listing below)
+    - `cluster`: Basic clustering primitives and algorithms.
+    - `comms`: A multi-node multi-GPU communications abstraction layer for NCCL+UCX and MPI+NCCL, which can be deployed in Dask clusters using the `raft-dask` Python package.
+    - `core`: Core API headers which require minimal dependencies aside from RMM and Cudatoolkit. These are safe to expose on public APIs and do not require `nvcc` to build. This is the same for any headers in RAFT which have the suffix `*_types.hpp`. 
+    - `distance`: Distance primitives
+    - `linalg`: Dense linear algebra
+    - `matrix`: Dense matrix operations
+    - `neighbors`: Nearest neighbors and knn graph construction
+    - `random`: Random number generation, sampling, and data generation primitives
+    - `solver`: Iterative and combinatorial solvers for optimization and approximation
+    - `sparse`: Sparse matrix operations
+      - `convert`: Sparse conversion functions
+      - `distance`: Sparse distance computations
+      - `linalg`: Sparse linear algebra
+      - `neighbors`: Sparse nearest neighbors and knn graph construction
+      - `op`: Various sparse operations such as slicing and filtering (Note: this will soon be renamed to `sparse/matrix`)
+      - `solver`: Sparse solvers for optimization and approximation
+    - `stats`: Moments, summary statistics, model performance measures
+    - `util`: Various reusable tools and utilities for accelerated algorithm development
   - `scripts`: Helpful scripts for development
   - `src`: Compiled APIs and template specializations for the shared libraries
   - `test`: Googletests source code

build.sh

@@ -40,8 +40,8 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
                                   the only option to be supported)
    --minimal-deps              - disables dependencies like thrust so they can be overridden.
                                  can be useful for a pure header-only install
-   --limit-tests               - semicolon-separated list of test executables to compile (e.g. SPATIAL_TEST;CLUSTER_TEST)
-   --limit-bench               - semicolon-separated list of benchmark executables to compute (e.g. SPATIAL_BENCH;CLUSTER_BENCH)
+   --limit-tests               - semicolon-separated list of test executables to compile (e.g. NEIGHBORS_TEST;CLUSTER_TEST)
+   --limit-bench               - semicolon-separated list of benchmark executables to compute (e.g. NEIGHBORS_BENCH;CLUSTER_BENCH)
    --allgpuarch                - build for all supported GPU architectures
    --buildfaiss                - build faiss statically into raft
    --install                   - install cmake targets
@@ -72,8 +72,8 @@ COMPILE_NN_LIBRARY=OFF
 COMPILE_DIST_LIBRARY=OFF
 ENABLE_NN_DEPENDENCIES=OFF
 
-TEST_TARGETS="CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NN_TEST;SPATIAL_TEST;STATS_TEST;UTILS_TEST"
-BENCH_TARGETS="CLUSTER_BENCH;SPATIAL_BENCH;DISTANCE_BENCH;LINALG_BENCH;SPARSE_BENCH;RANDOM_BENCH"
+TEST_TARGETS="CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NEIGHBORS_TEST;NEIGHBORS_TEST;STATS_TEST;UTILS_TEST"
+BENCH_TARGETS="CLUSTER_BENCH;NEIGHBORS_BENCH;DISTANCE_BENCH;LINALG_BENCH;SPARSE_BENCH;RANDOM_BENCH"
 ENABLE_thrust_DEPENDENCY=ON
 
 CACHE_ARGS=""

cpp/bench/CMakeLists.txt

@@ -30,8 +30,6 @@ set(BENCH_NAME ${ConfigureBench_NAME})
 
 add_executable(${BENCH_NAME} ${ConfigureBench_PATH})
 
-message("BENCH PATH: ${ConfigureBench_PATH}")
-
 target_link_libraries(${BENCH_NAME}
         PRIVATE
         raft::raft
@@ -114,20 +112,20 @@ if(BUILD_BENCH)
             bench/main.cpp
             )
 
-    ConfigureBench(NAME SPATIAL_BENCH
+    ConfigureBench(NAME NEIGHBORS_BENCH
             PATH
-            bench/spatial/fused_l2_nn.cu
-            bench/spatial/knn/brute_force_float_int64_t.cu
-            bench/spatial/knn/brute_force_float_uint32_t.cu
-            bench/spatial/knn/ivf_flat_float_int64_t.cu
-            bench/spatial/knn/ivf_flat_float_uint32_t.cu
-            bench/spatial/knn/ivf_flat_int8_t_int64_t.cu
-            bench/spatial/knn/ivf_flat_uint8_t_uint32_t.cu
-            bench/spatial/knn/ivf_pq_float_int64_t.cu
-            bench/spatial/knn/ivf_pq_float_uint32_t.cu
-            bench/spatial/knn/ivf_pq_int8_t_int64_t.cu
-            bench/spatial/knn/ivf_pq_uint8_t_uint32_t.cu
-            bench/spatial/selection.cu
+            bench/neighbors/fused_l2_nn.cu
+            bench/neighbors/knn/brute_force_float_int64_t.cu
+            bench/neighbors/knn/brute_force_float_uint32_t.cu
+            bench/neighbors/knn/ivf_flat_float_int64_t.cu
+            bench/neighbors/knn/ivf_flat_float_uint32_t.cu
+            bench/neighbors/knn/ivf_flat_int8_t_int64_t.cu
+            bench/neighbors/knn/ivf_flat_uint8_t_uint32_t.cu
+            bench/neighbors/knn/ivf_pq_float_int64_t.cu
+            bench/neighbors/knn/ivf_pq_float_uint32_t.cu
+            bench/neighbors/knn/ivf_pq_int8_t_int64_t.cu
+            bench/neighbors/knn/ivf_pq_uint8_t_uint32_t.cu
+            bench/neighbors/selection.cu
             bench/main.cpp
             OPTIONAL DIST NN
             )

cpp/bench/spatial/knn.cuh → cpp/bench/neighbors/knn.cuh

@@ -20,8 +20,8 @@
 
 #include <raft/random/rng.cuh>
 
-#include <raft/spatial/knn/ivf_flat.cuh>
-#include <raft/spatial/knn/ivf_pq.cuh>
+#include <raft/neighbors/ivf_flat.cuh>
+#include <raft/neighbors/ivf_pq.cuh>
 #include <raft/spatial/knn/knn.cuh>
 
 #if defined RAFT_DISTANCE_COMPILED
@@ -143,16 +143,16 @@ template <typename ValT, typename IdxT>
 struct ivf_flat_knn {
   using dist_t = float;
 
-  std::optional<const raft::spatial::knn::ivf_flat::index<ValT, IdxT>> index;
-  raft::spatial::knn::ivf_flat::index_params index_params;
-  raft::spatial::knn::ivf_flat::search_params search_params;
+  std::optional<const raft::neighbors::ivf_flat::index<ValT, IdxT>> index;
+  raft::neighbors::ivf_flat::index_params index_params;
+  raft::neighbors::ivf_flat::search_params search_params;
   params ps;
 
   ivf_flat_knn(const raft::handle_t& handle, const params& ps, const ValT* data) : ps(ps)
   {
     index_params.n_lists = 4096;
     index_params.metric  = raft::distance::DistanceType::L2Expanded;
-    index.emplace(raft::spatial::knn::ivf_flat::build(
+    index.emplace(raft::neighbors::ivf_flat::build(
       handle, index_params, data, IdxT(ps.n_samples), uint32_t(ps.n_dims)));
   }
 
@@ -162,7 +162,7 @@ struct ivf_flat_knn {
               IdxT* out_idxs)
   {
     search_params.n_probes = 20;
-    raft::spatial::knn::ivf_flat::search(
+    raft::neighbors::ivf_flat::search(
       handle, search_params, *index, search_items, ps.n_queries, ps.k, out_idxs, out_dists);
   }
 };
@@ -171,16 +171,16 @@ template <typename ValT, typename IdxT>
 struct ivf_pq_knn {
   using dist_t = float;
 
-  std::optional<const raft::spatial::knn::ivf_pq::index<IdxT>> index;
-  raft::spatial::knn::ivf_pq::index_params index_params;
-  raft::spatial::knn::ivf_pq::search_params search_params;
+  std::optional<const raft::neighbors::ivf_pq::index<IdxT>> index;
+  raft::neighbors::ivf_pq::index_params index_params;
+  raft::neighbors::ivf_pq::search_params search_params;
   params ps;
 
   ivf_pq_knn(const raft::handle_t& handle, const params& ps, const ValT* data) : ps(ps)
   {
     index_params.n_lists = 4096;
     index_params.metric  = raft::distance::DistanceType::L2Expanded;
-    index.emplace(raft::spatial::knn::ivf_pq::build(
+    index.emplace(raft::neighbors::ivf_pq::build(
       handle, index_params, data, IdxT(ps.n_samples), uint32_t(ps.n_dims)));
   }
 
@@ -190,7 +190,7 @@ struct ivf_pq_knn {
               IdxT* out_idxs)
   {
     search_params.n_probes = 20;
-    raft::spatial::knn::ivf_pq::search(
+    raft::neighbors::ivf_pq::search(
       handle, search_params, *index, search_items, ps.n_queries, ps.k, out_idxs, out_dists);
   }
 };

cpp/include/raft/cluster/detail/connectivities.cuh

@@ -27,7 +27,7 @@
 #include <raft/distance/distance_types.hpp>
 #include <raft/sparse/convert/csr.cuh>
 #include <raft/sparse/coo.hpp>
-#include <raft/sparse/spatial/knn_graph.cuh>
+#include <raft/sparse/neighbors/knn_graph.cuh>
 
 #include <thrust/iterator/zip_iterator.h>
 #include <thrust/transform.h>
@@ -73,7 +73,7 @@ struct distance_graph_impl<raft::cluster::LinkageDistance::KNN_GRAPH, value_idx,
     // Need to symmetrize knn into undirected graph
     raft::sparse::COO<value_t, value_idx> knn_graph_coo(stream);
 
-    raft::sparse::spatial::knn_graph(handle, X, m, n, metric, knn_graph_coo, c);
+    raft::sparse::neighbors::knn_graph(handle, X, m, n, metric, knn_graph_coo, c);
 
     indices.resize(knn_graph_coo.nnz, stream);
     data.resize(knn_graph_coo.nnz, stream);

cpp/include/raft/cluster/detail/mst.cuh

@@ -19,9 +19,9 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <raft/sparse/neighbors/connect_components.cuh>
 #include <raft/sparse/op/sort.cuh>
 #include <raft/sparse/solver/mst.cuh>
-#include <raft/sparse/spatial/connect_components.cuh>
 #include <rmm/device_uvector.hpp>
 
 #include <thrust/device_ptr.h>
@@ -80,7 +80,7 @@ void connect_knn_graph(
 
   raft::sparse::COO<value_t, value_idx> connected_edges(stream);
 
-  raft::sparse::spatial::connect_components<value_idx, value_t>(
+  raft::sparse::neighbors::connect_components<value_idx, value_t>(
     handle, connected_edges, X, color, m, n, reduction_op);
 
   rmm::device_uvector<value_idx> indptr2(m + 1, stream);
@@ -153,14 +153,14 @@ void build_sorted_mst(
     handle, indptr, indices, pw_dists, (value_idx)m, nnz, color, stream, false, true);
 
   int iters        = 1;
-  int n_components = raft::sparse::spatial::get_n_components(color, m, stream);
+  int n_components = raft::sparse::neighbors::get_n_components(color, m, stream);
 
   while (n_components > 1 && iters < max_iter) {
     connect_knn_graph<value_idx, value_t>(handle, X, mst_coo, m, n, color, reduction_op);
 
     iters++;
 
-    n_components = raft::sparse::spatial::get_n_components(color, m, stream);
+    n_components = raft::sparse::neighbors::get_n_components(color, m, stream);
   }
 
   /**

cpp/include/raft/cluster/detail/single_linkage.cuh

@@ -80,7 +80,7 @@ void single_linkage(const raft::handle_t& handle,
    * 2. Construct MST, sorted by weights
    */
   rmm::device_uvector<value_idx> color(m, stream);
-  raft::sparse::spatial::FixConnectivitiesRedOp<value_idx, value_t> op(color.data(), m);
+  raft::sparse::neighbors::FixConnectivitiesRedOp<value_idx, value_t> op(color.data(), m);
   detail::build_sorted_mst<value_idx, value_t>(handle,
                                                X,
                                                indptr.data(),

cpp/include/raft/neighbors/ann_types.hpp

@@ -0,0 +1,47 @@
+/*
+ * 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/distance/distance_types.hpp>
+
+namespace raft::neighbors::ann {
+
+/** The base for approximate KNN index structures. */
+struct index {
+};
+
+/** The base for KNN index parameters. */
+struct index_params {
+  /** Distance type. */
+  raft::distance::DistanceType metric = distance::DistanceType::L2Expanded;
+  /** The argument used by some distance metrics. */
+  float metric_arg = 2.0f;
+  /**
+   * Whether to add the dataset content to the index, i.e.:
+   *
+   *  - `true` means the index is filled with the dataset vectors and ready to search after calling
+   * `build`.
+   *  - `false` means `build` only trains the underlying model (e.g. quantizer or clustering), but
+   * the index is left empty; you'd need to call `extend` on the index afterwards to populate it.
+   */
+  bool add_data_on_build = true;
+};
+
+struct search_params {
+};
+
+};  // namespace raft::neighbors::ann