Merge branch 'branch-24.04' into doc_2404_add_cuvs_noice

build.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 
 # raft build scripts
 
@@ -512,6 +512,8 @@ fi
 
 if hasArg docs; then
     set -x
+    export RAPIDS_VERSION="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2.\3/' "${REPODIR}/VERSION")"
+    export RAPIDS_VERSION_MAJOR_MINOR="$(sed -E -e 's/^([0-9]{2})\.([0-9]{2})\.([0-9]{2}).*$/\1.\2/' "${REPODIR}/VERSION")"
     cd ${DOXYGEN_BUILD_DIR}
     doxygen Doxyfile
     cd ${SPHINX_BUILD_DIR}

ci/build_docs.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 set -euo pipefail
 
@@ -28,7 +28,9 @@ rapids-mamba-retry install \
   pylibraft \
   raft-dask
 
-export RAPIDS_VERSION_NUMBER="24.04"
+export RAPIDS_VERSION="$(rapids-version)"
+export RAPIDS_VERSION_MAJOR_MINOR="$(rapids-version-major-minor)"
+export RAPIDS_VERSION_NUMBER="$RAPIDS_VERSION_MAJOR_MINOR"
 export RAPIDS_DOCS_DIR="$(mktemp -d)"
 
 rapids-logger "Build CPP docs"

ci/release/update-version.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2020-2023, NVIDIA CORPORATION.
+# Copyright (c) 2020-2024, NVIDIA CORPORATION.
 ########################
 # RAFT Version Updater #
 ########################
@@ -36,23 +36,11 @@ function sed_runner() {
     sed -i.bak ''"$1"'' $2 && rm -f ${2}.bak
 }
 
-sed_runner "s/set(RAPIDS_VERSION .*)/set(RAPIDS_VERSION \"${NEXT_SHORT_TAG}\")/g" cpp/CMakeLists.txt
 sed_runner "s/set(RAPIDS_VERSION .*)/set(RAPIDS_VERSION \"${NEXT_SHORT_TAG}\")/g" cpp/template/cmake/thirdparty/fetch_rapids.cmake
-sed_runner "s/set(RAFT_VERSION .*)/set(RAFT_VERSION \"${NEXT_FULL_TAG}\")/g" cpp/CMakeLists.txt
-sed_runner 's/'"pylibraft_version .*)"'/'"pylibraft_version ${NEXT_FULL_TAG})"'/g' python/pylibraft/CMakeLists.txt
-sed_runner 's/'"raft_dask_version .*)"'/'"raft_dask_version ${NEXT_FULL_TAG})"'/g' python/raft-dask/CMakeLists.txt
-sed_runner 's/'"branch-.*\/RAPIDS.cmake"'/'"branch-${NEXT_SHORT_TAG}\/RAPIDS.cmake"'/g' fetch_rapids.cmake
 
 # Centralized version file update
 echo "${NEXT_FULL_TAG}" > VERSION
 
-# Wheel testing script
-sed_runner "s/branch-.*/branch-${NEXT_SHORT_TAG}/g" ci/test_wheel_raft_dask.sh
-
-# Docs update
-sed_runner 's/version = .*/version = '"'${NEXT_SHORT_TAG}'"'/g' docs/source/conf.py
-sed_runner 's/release = .*/release = '"'${NEXT_FULL_TAG}'"'/g' docs/source/conf.py
-
 DEPENDENCIES=(
   dask-cuda
   pylibraft
@@ -84,9 +72,6 @@ sed_runner "/^ucx_py_version:$/ {n;s/.*/  - \"${NEXT_UCX_PY_VERSION}\"/}" conda/
 for FILE in .github/workflows/*.yaml; do
   sed_runner "/shared-workflows/ s/@.*/@branch-${NEXT_SHORT_TAG}/g" "${FILE}"
 done
-sed_runner "s/RAPIDS_VERSION_NUMBER=\".*/RAPIDS_VERSION_NUMBER=\"${NEXT_SHORT_TAG}\"/g" ci/build_docs.sh
-
-sed_runner "/^PROJECT_NUMBER/ s|\".*\"|\"${NEXT_SHORT_TAG}\"|g" cpp/doxygen/Doxyfile
 
 sed_runner "/^set(RAFT_VERSION/ s|\".*\"|\"${NEXT_SHORT_TAG}\"|g" docs/source/build.md
 sed_runner "s|branch-[0-9][0-9].[0-9][0-9]|branch-${NEXT_SHORT_TAG}|g" docs/source/build.md

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -39,7 +39,7 @@ dependencies:
 - nccl>=2.9.9
 - ninja
 - numba>=0.57
-- numpy>=1.23
+- numpy>=1.23,<2.0a0
 - numpydoc
 - nvcc_linux-aarch64=11.8
 - pre-commit

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -39,7 +39,7 @@ dependencies:
 - nccl>=2.9.9
 - ninja
 - numba>=0.57
-- numpy>=1.23
+- numpy>=1.23,<2.0a0
 - numpydoc
 - nvcc_linux-64=11.8
 - pre-commit

conda/environments/all_cuda-122_arch-aarch64.yaml

@@ -36,7 +36,7 @@ dependencies:
 - nccl>=2.9.9
 - ninja
 - numba>=0.57
-- numpy>=1.23
+- numpy>=1.23,<2.0a0
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme

conda/environments/all_cuda-122_arch-x86_64.yaml

@@ -36,7 +36,7 @@ dependencies:
 - nccl>=2.9.9
 - ninja
 - numba>=0.57
-- numpy>=1.23
+- numpy>=1.23,<2.0a0
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme

conda/recipes/pylibraft/meta.yaml

@@ -65,7 +65,7 @@ requirements:
     {% endif %}
     - libraft {{ version }}
     - libraft-headers {{ version }}
-    - numpy >=1.23
+    - numpy >=1.23,<2.0a0
     - python x.x
     - rmm ={{ minor_version }}
 

cpp/CMakeLists.txt

@@ -10,11 +10,8 @@
 # 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(RAPIDS_VERSION "24.04")
-set(RAFT_VERSION "24.04.00")
-
 cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
-include(../fetch_rapids.cmake)
+include(../rapids_config.cmake)
 include(rapids-cmake)
 include(rapids-cpm)
 include(rapids-export)
@@ -34,7 +31,7 @@ endif()
 
 project(
   RAFT
-  VERSION ${RAFT_VERSION}
+  VERSION "${RAPIDS_VERSION}"
   LANGUAGES ${lang_list}
 )
 

cpp/doxygen/Doxyfile

@@ -38,7 +38,7 @@ PROJECT_NAME           = "RAFT C++ API"
 # could be handy for archiving the generated documentation or if some version
 # control system is used.
 
-PROJECT_NUMBER         = "24.04"
+PROJECT_NUMBER         = "$(RAPIDS_VERSION_MAJOR_MINOR)"
 
 # Using the PROJECT_BRIEF tag one can provide an optional one line description
 # for a project that appears at the top of each page and should give viewer a

cpp/include/raft/neighbors/ball_cover-inl.cuh

@@ -333,7 +333,7 @@ void eps_nn(raft::resources const& handle,
     query.extent(0),
     adj.data_handle(),
     vd.data_handle(),
-    spatial::knn::detail::EuclideanFunc<value_t, int_t>());
+    spatial::knn::detail::EuclideanSqFunc<value_t, int_t>());
 }
 
 /**
@@ -392,7 +392,7 @@ void eps_nn(raft::resources const& handle,
     adj_ia.data_handle(),
     adj_ja.data_handle(),
     vd.data_handle(),
-    spatial::knn::detail::EuclideanFunc<value_t, int_t>());
+    spatial::knn::detail::EuclideanSqFunc<value_t, int_t>());
 }
 
 /**

cpp/include/raft/neighbors/ball_cover_types.hpp

@@ -69,6 +69,7 @@ class BallCoverIndex {
       R_1nn_dists(raft::make_device_vector<value_t, matrix_idx>(handle, m_)),
       R_closest_landmark_dists(raft::make_device_vector<value_t, matrix_idx>(handle, m_)),
       R(raft::make_device_matrix<value_t, matrix_idx>(handle, sqrt(m_), n_)),
+      X_reordered(raft::make_device_matrix<value_t, matrix_idx>(handle, m_, n_)),
       R_radius(raft::make_device_vector<value_t, matrix_idx>(handle, sqrt(m_))),
       index_trained(false)
   {
@@ -93,6 +94,8 @@ class BallCoverIndex {
       R_1nn_dists(raft::make_device_vector<value_t, matrix_idx>(handle, X_.extent(0))),
       R_closest_landmark_dists(raft::make_device_vector<value_t, matrix_idx>(handle, X_.extent(0))),
       R(raft::make_device_matrix<value_t, matrix_idx>(handle, sqrt(X_.extent(0)), X_.extent(1))),
+      X_reordered(
+        raft::make_device_matrix<value_t, matrix_idx>(handle, X_.extent(0), X_.extent(1))),
       R_radius(raft::make_device_vector<value_t, matrix_idx>(handle, sqrt(X_.extent(0)))),
       index_trained(false)
   {
@@ -122,6 +125,10 @@ class BallCoverIndex {
   {
     return R_closest_landmark_dists.view();
   }
+  auto get_X_reordered() const -> raft::device_matrix_view<const value_t, matrix_idx, row_major>
+  {
+    return X_reordered.view();
+  }
 
   raft::device_vector_view<value_idx, matrix_idx> get_R_indptr() { return R_indptr.view(); }
   raft::device_vector_view<value_idx, matrix_idx> get_R_1nn_cols() { return R_1nn_cols.view(); }
@@ -132,6 +139,10 @@ class BallCoverIndex {
   {
     return R_closest_landmark_dists.view();
   }
+  raft::device_matrix_view<value_t, matrix_idx, row_major> get_X_reordered()
+  {
+    return X_reordered.view();
+  }
   raft::device_matrix_view<const value_t, matrix_idx, row_major> get_X() const { return X; }
 
   raft::distance::DistanceType get_metric() const { return metric; }
@@ -162,6 +173,7 @@ class BallCoverIndex {
   raft::device_vector<value_t, matrix_idx> R_radius;
 
   raft::device_matrix<value_t, matrix_idx, row_major> R;
+  raft::device_matrix<value_t, matrix_idx, row_major> X_reordered;
 
  protected:
   bool index_trained;

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

@@ -165,6 +165,10 @@ void construct_landmark_1nn(raft::resources const& handle,
                                            index.get_R_indptr().data_handle(),
                                            index.n_landmarks + 1,
                                            resource::get_cuda_stream(handle));
+
+  // reorder X to allow aligned access
+  raft::matrix::copy_rows<value_t, value_idx>(
+    handle, index.get_X(), index.get_X_reordered(), index.get_R_1nn_cols());
 }
 
 /**
@@ -337,12 +341,6 @@ void perform_rbc_query(raft::resources const& handle,
 /**
  * Perform eps-select
  *
- * a. Map 1 row to each warp/block
- * b. Add closest k R points to heap
- * c. Iterate through batches of R, having each thread in the warp load a set
- * of distances y from R (only if d(q, r) < 3 * distance to closest r) and
- * marking the distance to be computed between x, y only
- * if knn[k].distance >= d(x_i, R_k) + d(R_k, y)
  */
 template <typename value_idx,
           typename value_t,
@@ -355,7 +353,7 @@ void perform_rbc_eps_nn_query(
   const value_t* query,
   value_int n_query_pts,
   value_t eps,
-  const value_t* landmark_dists,
+  const value_t* landmarks,
   dist_func dfunc,
   bool* adj,
   value_idx* vd)
@@ -367,7 +365,7 @@ void perform_rbc_eps_nn_query(
   resource::sync_stream(handle);
 
   rbc_eps_pass<value_idx, value_t, value_int, matrix_idx>(
-    handle, index, query, n_query_pts, eps, landmark_dists, dfunc, adj, vd);
+    handle, index, query, n_query_pts, eps, landmarks, dfunc, adj, vd);
 
   resource::sync_stream(handle);
 }
@@ -384,14 +382,14 @@ void perform_rbc_eps_nn_query(
   value_int n_query_pts,
   value_t eps,
   value_int* max_k,
-  const value_t* landmark_dists,
+  const value_t* landmarks,
   dist_func dfunc,
   value_idx* adj_ia,
   value_idx* adj_ja,
   value_idx* vd)
 {
   rbc_eps_pass<value_idx, value_t, value_int, matrix_idx>(
-    handle, index, query, n_query_pts, eps, max_k, landmark_dists, dfunc, adj_ia, adj_ja, vd);
+    handle, index, query, n_query_pts, eps, max_k, landmarks, dfunc, adj_ia, adj_ja, vd);
 
   resource::sync_stream(handle);
 }
@@ -664,15 +662,9 @@ void rbc_eps_nn_query(raft::resources const& handle,
 {
   ASSERT(index.is_index_trained(), "index must be previously trained");
 
-  auto R_dists =
-    raft::make_device_matrix<value_t, matrix_idx>(handle, index.n_landmarks, n_query_pts);
-
-  // find all landmarks that might have points in range
-  compute_landmark_dists(handle, index, query, n_query_pts, R_dists.data_handle());
-
   // query all points and write to adj
   perform_rbc_eps_nn_query(
-    handle, index, query, n_query_pts, eps, R_dists.data_handle(), dfunc, adj, vd);
+    handle, index, query, n_query_pts, eps, index.get_R().data_handle(), dfunc, adj, vd);
 }
 
 template <typename value_idx = std::int64_t,
@@ -693,20 +685,14 @@ void rbc_eps_nn_query(raft::resources const& handle,
 {
   ASSERT(index.is_index_trained(), "index must be previously trained");
 
-  auto R_dists =
-    raft::make_device_matrix<value_t, matrix_idx>(handle, index.n_landmarks, n_query_pts);
-
-  // find all landmarks that might have points in range
-  compute_landmark_dists(handle, index, query, n_query_pts, R_dists.data_handle());
-
   // query all points and write to adj
   perform_rbc_eps_nn_query(handle,
                            index,
                            query,
                            n_query_pts,
                            eps,
                            max_k,
-                           R_dists.data_handle(),
+                           index.get_R().data_handle(),
                            dfunc,
                            adj_ia,
                            adj_ja,

cpp/include/raft/spatial/knn/detail/ball_cover/registers-ext.cuh

@@ -190,7 +190,7 @@ instantiate_raft_spatial_knn_detail_rbc_low_dim_pass_two(
   std::int64_t, float, std::int64_t, std::int64_t, 3, raft::spatial::knn::detail::DistFunc);
 
 instantiate_raft_spatial_knn_detail_rbc_eps_pass(
-  std::int64_t, float, std::int64_t, std::int64_t, raft::spatial::knn::detail::EuclideanFunc);
+  std::int64_t, float, std::int64_t, std::int64_t, raft::spatial::knn::detail::EuclideanSqFunc);
 
 #undef instantiate_raft_spatial_knn_detail_rbc_low_dim_pass_two
 #undef instantiate_raft_spatial_knn_detail_rbc_low_dim_pass_one