Merge branch 'branch-24.12' into lanczos-solver-new

ci/build_wheel_pylibraft.sh

@@ -3,6 +3,8 @@
 
 set -euo pipefail
 
+package_dir="python/pylibraft"
+
 case "${RAPIDS_CUDA_VERSION}" in
   12.*)
     EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
@@ -15,4 +17,5 @@ esac
 # Set up skbuild options. Enable sccache in skbuild config options
 export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}"
 
-ci/build_wheel.sh pylibraft python/pylibraft
+ci/build_wheel.sh pylibraft ${package_dir}
+ci/validate_wheel.sh ${package_dir} final_dist

ci/build_wheel_raft_dask.sh

@@ -3,7 +3,10 @@
 
 set -euo pipefail
 
+package_dir="python/raft-dask"
+
 # Set up skbuild options. Enable sccache in skbuild config options
 export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
 
-ci/build_wheel.sh raft-dask python/raft-dask
+ci/build_wheel.sh raft-dask ${package_dir}
+ci/validate_wheel.sh ${package_dir} final_dist

ci/validate_wheel.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION.
+
+set -euo pipefail
+
+package_dir=$1
+wheel_dir_relative_path=$2
+
+cd "${package_dir}"
+
+rapids-logger "validate packages with 'pydistcheck'"
+
+pydistcheck \
+    --inspect \
+    "$(echo ${wheel_dir_relative_path}/*.whl)"
+
+rapids-logger "validate packages with 'twine'"
+
+twine check \
+    --strict \
+    "$(echo ${wheel_dir_relative_path}/*.whl)"

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

@@ -14,12 +14,12 @@ dependencies:
 - cmake>=3.26.4,!=3.30.0
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
-- cuda-python>=11.7.1,<12.0a0
+- cuda-python>=11.7.1,<12.0a0,<=11.8.3
 - cuda-version=11.8
 - cudatoolkit
 - cupy>=12.0.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - dask-cuda==24.12.*,>=0.0.0a0
 - distributed-ucxx==0.41.*,>=0.0.0a0
 - doxygen>=1.8.20

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

@@ -14,12 +14,12 @@ dependencies:
 - cmake>=3.26.4,!=3.30.0
 - cuda-nvtx=11.8
 - cuda-profiler-api=11.8.86
-- cuda-python>=11.7.1,<12.0a0
+- cuda-python>=11.7.1,<12.0a0,<=11.8.3
 - cuda-version=11.8
 - cudatoolkit
 - cupy>=12.0.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - dask-cuda==24.12.*,>=0.0.0a0
 - distributed-ucxx==0.41.*,>=0.0.0a0
 - doxygen>=1.8.20

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

@@ -16,11 +16,11 @@ dependencies:
 - cuda-nvcc
 - cuda-nvtx-dev
 - cuda-profiler-api
-- cuda-python>=12.0,<13.0a0
+- cuda-python>=12.0,<13.0a0,<=12.6.0
 - cuda-version=12.5
 - cupy>=12.0.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - dask-cuda==24.12.*,>=0.0.0a0
 - distributed-ucxx==0.41.*,>=0.0.0a0
 - doxygen>=1.8.20

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

@@ -16,11 +16,11 @@ dependencies:
 - cuda-nvcc
 - cuda-nvtx-dev
 - cuda-profiler-api
-- cuda-python>=12.0,<13.0a0
+- cuda-python>=12.0,<13.0a0,<=12.6.0
 - cuda-version=12.5
 - cupy>=12.0.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - dask-cuda==24.12.*,>=0.0.0a0
 - distributed-ucxx==0.41.*,>=0.0.0a0
 - doxygen>=1.8.20

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

@@ -17,7 +17,7 @@ dependencies:
 - cuda-version=11.8
 - cudatoolkit
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - gcc_linux-aarch64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0

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

@@ -17,7 +17,7 @@ dependencies:
 - cuda-version=11.8
 - cudatoolkit
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - gcc_linux-64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0

conda/environments/bench_ann_cuda-120_arch-aarch64.yaml

@@ -18,7 +18,7 @@ dependencies:
 - cuda-profiler-api
 - cuda-version=12.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - gcc_linux-aarch64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0

conda/environments/bench_ann_cuda-120_arch-x86_64.yaml

@@ -18,7 +18,7 @@ dependencies:
 - cuda-profiler-api
 - cuda-version=12.0
 - cxx-compiler
-- cython>=3.0.0
+- cython>=3.0.0,<3.1.0a0
 - gcc_linux-64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0

conda/recipes/pylibraft/meta.yaml

@@ -26,6 +26,7 @@ build:
     - {{ compiler('cuda') }}
     - cuda-cudart-dev
     {% endif %}
+    - cuda-python
 
 requirements:
   build:
@@ -42,14 +43,14 @@ requirements:
     - {{ stdlib("c") }}
   host:
     {% if cuda_major == "11" %}
-    - cuda-python >=11.7.1,<12.0a0
+    - cuda-python >=11.7.1,<12.0a0,<=11.8.3
     - cudatoolkit
     {% else %}
-    - cuda-python >=12.0,<13.0a0
+    - cuda-python >=12.0,<13.0a0,<=12.6.0
     - cuda-cudart-dev
     {% endif %}
     - cuda-version ={{ cuda_version }}
-    - cython >=3.0.0
+    - cython >=3.0.0,<3.1.0a0
     - libraft {{ version }}
     - libraft-headers {{ version }}
     - python x.x
@@ -60,8 +61,10 @@ requirements:
     - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
     {% if cuda_major == "11" %}
     - cudatoolkit
+    - cuda-python >=11.7.1,<12.0a0,<=11.8.3
     {% else %}
     - cuda-cudart
+    - cuda-python >=12.0,<13.0a0,<=12.6.0
     {% endif %}
     - libraft {{ version }}
     - libraft-headers {{ version }}

conda/recipes/raft-dask/meta.yaml

@@ -26,6 +26,7 @@ build:
     - {{ compiler('cuda') }}
     - cuda-cudart-dev
     {% endif %}
+    - cuda-python
 
 requirements:
   build:
@@ -42,14 +43,14 @@ requirements:
     - {{ stdlib("c") }}
   host:
     {% if cuda_major == "11" %}
-    - cuda-python >=11.7.1,<12.0a0
+    - cuda-python >=11.7.1,<12.0a0,<=11.8.3
     - cudatoolkit
     {% else %}
-    - cuda-python >=12.0,<13.0a0
+    - cuda-python >=12.0,<13.0a0,<=12.6.0
     - cuda-cudart-dev
     {% endif %}
     - cuda-version ={{ cuda_version }}
-    - cython >=3.0.0
+    - cython >=3.0.0,<3.1.0a0
     - nccl {{ nccl_version }}
     - pylibraft {{ version }}
     - python x.x
@@ -61,8 +62,10 @@ requirements:
   run:
     {% if cuda_major == "11" %}
     - cudatoolkit
+    - cuda-python >=11.7.1,<12.0a0,<=11.8.3
     {% else %}
     - cuda-cudart
+    - cuda-python >=12.0,<13.0a0,<=12.6.0
     {% endif %}
     - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
     - dask-cuda ={{ minor_version }}

cpp/cmake/patches/faiss_override.json

@@ -1,9 +1,9 @@
 {
   "packages" : {
     "faiss" : {
-      "version": "1.7.4",
+      "version": "1.9.0",
       "git_url": "https://github.com/facebookresearch/faiss.git",
-      "git_tag": "main"
+      "git_tag": "v1.9.0"
     }
   }
 }

cpp/include/raft/core/pinned_container_policy.hpp

@@ -19,9 +19,9 @@
 
 #include <cstddef>
 #ifndef RAFT_DISABLE_CUDA
-#include <thrust/host_vector.h>
-#include <thrust/mr/allocator.h>
-#include <thrust/system/cuda/memory_resource.h>
+#include <cuda_runtime.h>
+
+#include <type_traits>
 #else
 #include <raft/core/detail/fail_container_policy.hpp>
 #endif
@@ -30,20 +30,16 @@ namespace raft {
 #ifndef RAFT_DISABLE_CUDA
 
 /**
- * @brief A thin wrapper over thrust::host_vector for implementing the pinned mdarray container
- * policy.
+ * @brief A thin wrapper over cudaMallocHost/cudaFreeHost for implementing the pinned mdarray
+ * container policy.
  *
  */
 template <typename T>
 struct pinned_container {
-  using value_type = T;
-  using allocator_type =
-    thrust::mr::stateless_resource_allocator<value_type,
-                                             thrust::cuda::universal_host_pinned_memory_resource>;
+  using value_type = std::remove_cv_t<T>;
 
  private:
-  using underlying_container_type = thrust::host_vector<value_type, allocator_type>;
-  underlying_container_type data_;
+  value_type* data_ = nullptr;
 
  public:
   using size_type = std::size_t;
@@ -57,21 +53,24 @@ struct pinned_container {
   using iterator       = pointer;
   using const_iterator = const_pointer;
 
-  ~pinned_container()                           = default;
-  pinned_container(pinned_container&&) noexcept = default;
-  pinned_container(pinned_container const& that) : data_{that.data_} {}
+  explicit pinned_container(std::size_t size)
+  {
+    RAFT_CUDA_TRY(cudaMallocHost(&data_, size * sizeof(value_type)));
+  }
+  ~pinned_container() noexcept
+  {
+    if (data_ != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaFreeHost(data_)); }
+  }
 
-  auto operator=(pinned_container<T> const& that) -> pinned_container<T>&
+  pinned_container(pinned_container&& other) { std::swap(this->data_, other.data_); }
+  pinned_container& operator=(pinned_container&& other)
   {
-    data_ = underlying_container_type{that.data_};
+    std::swap(this->data_, other.data_);
     return *this;
   }
-  auto operator=(pinned_container<T>&& that) noexcept -> pinned_container<T>& = default;
+  pinned_container(pinned_container const&) = delete;  // Copying disallowed: one array one owner
+  pinned_container& operator=(pinned_container const&) = delete;
 
-  /**
-   * @brief Ctor that accepts a size.
-   */
-  explicit pinned_container(std::size_t size, allocator_type const& alloc) : data_{size, alloc} {}
   /**
    * @brief Index operator that returns a reference to the actual data.
    */
@@ -84,15 +83,13 @@ struct pinned_container {
    * @brief Index operator that returns a reference to the actual data.
    */
   template <typename Index>
-  auto operator[](Index i) const noexcept
+  auto operator[](Index i) const noexcept -> const_reference
   {
     return data_[i];
   }
 
-  void resize(size_type size) { data_.resize(size, data_.stream()); }
-
-  [[nodiscard]] auto data() noexcept -> pointer { return data_.data().get(); }
-  [[nodiscard]] auto data() const noexcept -> const_pointer { return data_.data().get(); }
+  [[nodiscard]] auto data() noexcept -> pointer { return data_; }
+  [[nodiscard]] auto data() const noexcept -> const_pointer { return data_; }
 };
 
 /**
@@ -102,23 +99,14 @@ template <typename ElementType>
 struct pinned_vector_policy {
   using element_type          = ElementType;
   using container_type        = pinned_container<element_type>;
-  using allocator_type        = typename container_type::allocator_type;
   using pointer               = typename container_type::pointer;
   using const_pointer         = typename container_type::const_pointer;
   using reference             = typename container_type::reference;
   using const_reference       = typename container_type::const_reference;
   using accessor_policy       = std::experimental::default_accessor<element_type>;
   using const_accessor_policy = std::experimental::default_accessor<element_type const>;
 
-  auto create(raft::resources const&, size_t n) -> container_type
-  {
-    return container_type(n, allocator_);
-  }
-
-  constexpr pinned_vector_policy() noexcept(std::is_nothrow_default_constructible_v<ElementType>)
-    : allocator_{}
-  {
-  }
+  auto create(raft::resources const&, size_t n) -> container_type { return container_type(n); }
 
   [[nodiscard]] constexpr auto access(container_type& c, size_t n) const noexcept -> reference
   {
@@ -132,9 +120,6 @@ struct pinned_vector_policy {
 
   [[nodiscard]] auto make_accessor_policy() noexcept { return accessor_policy{}; }
   [[nodiscard]] auto make_accessor_policy() const noexcept { return const_accessor_policy{}; }
-
- private:
-  allocator_type allocator_;
 };
 #else
 template <typename ElementType>

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

@@ -61,14 +61,9 @@ struct pointer_residency_count<Type, Types...> {
     auto [on_device, on_host] = pointer_residency_count<Types...>::run(ptrs...);
     cudaPointerAttributes attr;
     RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr));
-    switch (attr.type) {
-      case cudaMemoryTypeUnregistered: return std::make_tuple(on_device, on_host + 1);
-      case cudaMemoryTypeHost:
-        return std::make_tuple(on_device + int(attr.devicePointer == ptr), on_host + 1);
-      case cudaMemoryTypeDevice: return std::make_tuple(on_device + 1, on_host);
-      case cudaMemoryTypeManaged: return std::make_tuple(on_device + 1, on_host + 1);
-      default: return std::make_tuple(on_device, on_host);
-    }
+    if (attr.devicePointer || attr.type == cudaMemoryTypeDevice) { ++on_device; }
+    if (attr.hostPointer || attr.type == cudaMemoryTypeUnregistered) { ++on_host; }
+    return std::make_tuple(on_device, on_host);
   }
 };
 

cpp/include/raft/util/cudart_utils.hpp

@@ -223,12 +223,13 @@ void print_vector(const char* variable_name, const T* ptr, size_t componentsCoun
 {
   cudaPointerAttributes attr;
   RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr));
-  if (attr.hostPointer != nullptr) {
+  if (attr.hostPointer) {
     print_host_vector(variable_name, reinterpret_cast<T*>(attr.hostPointer), componentsCount, out);
   } else if (attr.type == cudaMemoryTypeUnregistered) {
     print_host_vector(variable_name, ptr, componentsCount, out);
   } else {
-    print_device_vector(variable_name, ptr, componentsCount, out);
+    print_device_vector(
+      variable_name, reinterpret_cast<T*>(attr.devicePointer), componentsCount, out);
   }
 }
 /** @} */