Merge branch 'main' into cuda-invalid-argument-kernel-error

Author: achirkin

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

@@ -46,7 +46,7 @@ dependencies:
 - recommonmark
 - rust
 - scikit-build-core>=0.11.0
-- scikit-learn
+- scikit-learn>=1.5
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx>=8.0.0

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

@@ -46,7 +46,7 @@ dependencies:
 - recommonmark
 - rust
 - scikit-build-core>=0.11.0
-- scikit-learn
+- scikit-learn>=1.5
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx>=8.0.0

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

@@ -47,7 +47,7 @@ dependencies:
 - recommonmark
 - rust
 - scikit-build-core>=0.11.0
-- scikit-learn
+- scikit-learn>=1.5
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx>=8.0.0

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

@@ -47,7 +47,7 @@ dependencies:
 - recommonmark
 - rust
 - scikit-build-core>=0.11.0
-- scikit-learn
+- scikit-learn>=1.5
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sphinx>=8.0.0

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

@@ -41,7 +41,7 @@ dependencies:
 - pyyaml
 - rapids-build-backend>=0.4.0,<0.5.0
 - requests
-- scikit-learn
+- scikit-learn>=1.5
 - setuptools>=77.0.0
 - sysroot_linux-aarch64==2.28
 - wheel

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

@@ -44,7 +44,7 @@ dependencies:
 - pyyaml
 - rapids-build-backend>=0.4.0,<0.5.0
 - requests
-- scikit-learn
+- scikit-learn>=1.5
 - setuptools>=77.0.0
 - sysroot_linux-64==2.28
 - wheel

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

@@ -42,7 +42,7 @@ dependencies:
 - pyyaml
 - rapids-build-backend>=0.4.0,<0.5.0
 - requests
-- scikit-learn
+- scikit-learn>=1.5
 - setuptools>=77.0.0
 - sysroot_linux-aarch64==2.28
 - wheel

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

@@ -45,7 +45,7 @@ dependencies:
 - pyyaml
 - rapids-build-backend>=0.4.0,<0.5.0
 - requests
-- scikit-learn
+- scikit-learn>=1.5
 - setuptools>=77.0.0
 - sysroot_linux-64==2.28
 - wheel

cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp

@@ -15,9 +15,17 @@
 #include <memory>
 
 struct AlgorithmLauncher {
-  AlgorithmLauncher() = default;
+  AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {}
 
-  AlgorithmLauncher(cudaKernel_t k);
+  AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib);
+
+  ~AlgorithmLauncher();
+
+  AlgorithmLauncher(const AlgorithmLauncher&)            = delete;
+  AlgorithmLauncher& operator=(const AlgorithmLauncher&) = delete;
+
+  AlgorithmLauncher(AlgorithmLauncher&& other) noexcept;
+  AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept;
 
   template <typename... Args>
   void dispatch(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
@@ -31,6 +39,7 @@ struct AlgorithmLauncher {
  private:
   void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
   cudaKernel_t kernel;
+  cudaLibrary_t library;
 };
 
 std::unordered_map<std::string, std::shared_ptr<AlgorithmLauncher>>& get_cached_launchers();

cpp/src/cluster/detail/kmeans_common.cuh

@@ -31,7 +31,10 @@
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
-#include <cub/cub.cuh>
+#include <cub/device/device_histogram.cuh>
+#include <cub/device/device_reduce.cuh>
+#include <cub/device/device_select.cuh>
+#include <cub/iterator/arg_index_input_iterator.cuh>
 #include <cuda.h>
 #include <thrust/fill.h>
 #include <thrust/for_each.h>