Improvement of GPU AiSaQ

cmake/libs/libdiskann.cmake

@@ -33,6 +33,7 @@ set(DISKANN_LINKER_LIBS PUBLIC ${AIO_LIBRARIES} ${DISKANN_BOOST_PROGRAM_OPTIONS_
 if (WITH_CUVS)
     list(APPEND DISKANN_LINKER_LIBS PRIVATE cuvs::cuvs)
     list(APPEND DISKANN_SOURCES thirdparty/DiskANN/src/diskann_gpu.cpp)
+    list(APPEND DISKANN_SOURCES thirdparty/DiskANN/src/mean_center_gpu.cu)
 endif()
 
 add_library(diskann STATIC ${DISKANN_SOURCES})

thirdparty/DiskANN/include/diskann/diskann_gpu.h

@@ -42,6 +42,8 @@
 #ifdef KNOWHERE_WITH_CUVS
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/device_resources_manager.hpp>
+#include <raft/core/pinned_mdarray.hpp>
 
 template <typename T, typename idxT = std::size_t>
 raft::device_matrix<T, idxT> read_bin_dataset(const raft::device_resources& dev_resources,
@@ -57,9 +59,20 @@ void vamana_build_and_write(raft::device_resources const& dev_resources,
                             int iters);
 bool is_gpu_available();
 
+// train function for centroids computation
+void kmeans_gpu_device_data(
+    const raft::resources& dev_resources,
+    const float* d_chunk_data,
+    size_t num_train,
+    size_t chunk_size,
+    size_t num_centers,
+    int max_iter,
+    float* d_centroids_out,
+	bool is_balanced=false);
+
 // train function for centroids computation
 void kmeans_gpu(
-    raft::resources& dev_resources,
+    const raft::resources& dev_resources,
     const float* h_chunk_data,
     size_t num_train,
     size_t chunk_size,
@@ -68,9 +81,29 @@ void kmeans_gpu(
     float* h_centroids_out,
 	bool is_balanced=false);
 
+// train function for centroids combined with medoids computation
+void kmeans_medoids_gpu(
+    const raft::resources& dev_resources,
+    const float* h_chunk_data,
+    size_t num_train,
+    size_t dim,
+    size_t num_centers,
+    int max_iter,
+    float* h_centroids_out,
+	bool is_balanced,
+	uint32_t* h_medoids);
+
+// mean-centers the (num_train x dim) row-major training matrix in place and
+// writes the per-dimension centroid (length dim) to d_centroid_out
+void mean_center_gpu(const raft::resources& dev_resources,
+                     float* d_train_data,
+                     size_t num_train,
+                     size_t dim,
+                     float* d_centroid_out);
+
 // assign function for computes the compressed PQ vector per each vector in the dataset
 // based on centroids computation
-int predict_gpu(raft::resources& dev_resources,
+int predict_gpu(const raft::resources& dev_resources,
 				const float* h_data,
                 size_t n_samples,
                 size_t dim,
@@ -79,7 +112,7 @@ int predict_gpu(raft::resources& dev_resources,
 				int* h_labels);
 
 template <typename T>
-int brute_force_gpu(raft::resources& dev_resources,
+int brute_force_gpu(const raft::resources& dev_resources,
 		const T* h_data,
         size_t n_samples,
         size_t dim,
@@ -88,7 +121,7 @@ int brute_force_gpu(raft::resources& dev_resources,
         size_t n_queries,
 		int64_t* h_labels);
 
-void gpu_get_mem_info(raft::resources &dev_resources, size_t &gpu_free_mem,size_t &gpu_total_mem);
+void gpu_get_mem_info(const raft::resources &dev_resources, size_t &gpu_free_mem,size_t &gpu_total_mem);
 
 #endif
 

thirdparty/DiskANN/src/diskann_gpu.cpp

@@ -41,6 +41,7 @@
 #include <fstream>
 #include <iostream>
 #include <raft/core/resources.hpp>
+#include <raft/core/host_mdarray.hpp>
 #include <cuvs/neighbors/brute_force.hpp>
 
 template<typename T, typename idxT>
@@ -122,9 +123,9 @@ void vamana_build_and_write(raft::device_resources const &dev_resources,
 	serialize(dev_resources, out_fname, index, false);
 }
 
-void kmeans_gpu(raft::resources &dev_resources, const float *h_chunk_data,
+void kmeans_gpu_device_data(const raft::resources &dev_resources, const float *d_chunk_data,
 		size_t num_train, size_t dim, size_t num_centers, int max_iter,
-		float *h_centroids_out, bool is_balanced/*=false*/) {
+		float *d_centroids_out, bool is_balanced/*=false*/) {
 	// KMeans parameters
 	cuvs::cluster::kmeans::params km_params;
 	int64_t n_iter;
@@ -133,34 +134,73 @@ void kmeans_gpu(raft::resources &dev_resources, const float *h_chunk_data,
 	km_params.n_init = 1;
 	km_params.n_clusters = num_centers;
 
-	// Allocate device matrices
-	auto d_data = raft::make_device_matrix<float, int64_t>(dev_resources, num_train, dim);
-	auto d_centroids = raft::make_device_matrix<float, int64_t>(dev_resources, num_centers, dim);
+	auto d_data = raft::make_device_matrix_view<const float, int64_t>(d_chunk_data, num_train, dim);
+	auto d_centroids = raft::make_device_matrix_view<float, int64_t>(d_centroids_out, num_centers, dim);
 
-	// Copy input chunk to device asynchronously
-	raft::copy(d_data.data_handle(), h_chunk_data, num_train * dim, raft::resource::get_cuda_stream(dev_resources));
 	// Run KMeans (fit uses RAFT-managed streams internally)
 	if(!is_balanced) {
 		cuvs::cluster::kmeans::fit(dev_resources, km_params,
-				raft::make_const_mdspan(d_data.view()), std::nullopt,
-				d_centroids.view(),
+				d_data, std::nullopt,
+				d_centroids,
 				raft::make_host_scalar_view<float>(&inertia),
 				raft::make_host_scalar_view<int64_t>(&n_iter));
-	}else {
+	} else {
 		//use balance kmeans
 		cuvs::cluster::kmeans::balanced_params b_p;
 		b_p.n_iters=100;
-		cuvs::cluster::kmeans::fit(dev_resources, b_p,
-				raft::make_const_mdspan(d_data.view()), d_centroids.view());
+		cuvs::cluster::kmeans::fit(dev_resources, b_p, d_data, d_centroids);
 	}
-	// Copy centroids back to host asynchronously
+}
+
+void kmeans_gpu(const raft::resources &dev_resources, const float *h_chunk_data,
+		size_t num_train, size_t dim, size_t num_centers, int max_iter,
+		float *h_centroids_out, bool is_balanced/*=false*/) {
+
+	// Allocate device matrices
+	auto d_data = raft::make_device_matrix<float, int64_t>(dev_resources, num_train, dim);
+	auto d_centroids = raft::make_device_matrix<float, int64_t>(dev_resources, num_centers, dim);
+
+
+	raft::copy(d_data.data_handle(), h_chunk_data, num_train * dim, raft::resource::get_cuda_stream(dev_resources));
+	kmeans_gpu_device_data(dev_resources, d_data.data_handle(), num_train, dim, num_centers, max_iter, d_centroids.data_handle(), is_balanced);
 	raft::copy(h_centroids_out, d_centroids.data_handle(), num_centers * dim,
 			raft::resource::get_cuda_stream(dev_resources));
     raft::resource::sync_stream(dev_resources);
 }
 
+void kmeans_medoids_gpu(const raft::resources &dev_resources, const float *h_chunk_data,
+		size_t num_train, size_t dim, size_t num_centers, int max_iter,
+		float *h_centroids_out, bool is_balanced, uint32_t* h_medoids) {
+
+	// Allocate device matrices
+	auto d_data = raft::make_device_matrix<float, int64_t>(dev_resources, num_train, dim);
+	auto d_centroids = raft::make_device_matrix<float, int64_t>(dev_resources, num_centers, dim);
+    auto d_medoids_i64    = raft::make_device_matrix<int64_t, int64_t>(dev_resources, num_centers, 1);
+	auto h_medoids_i64    = raft::make_host_vector<int64_t, int64_t>(num_centers);
+
+	// Copy input chunk to device asynchronously
+	raft::copy(d_data.data_handle(), h_chunk_data, num_train * dim, raft::resource::get_cuda_stream(dev_resources));
+	kmeans_gpu_device_data(dev_resources, d_data.data_handle(), num_train, dim, num_centers, max_iter, d_centroids.data_handle(), is_balanced);
+
+	// Compute medoids using brute force
+	using namespace cuvs::neighbors;
+	brute_force::index_params bf_index_params;
+	brute_force::search_params bf_search_params;
+	auto index = brute_force::build(dev_resources, bf_index_params, raft::make_const_mdspan(d_data.view()));
+    auto distances = raft::make_device_matrix<float, int64_t>(dev_resources, num_centers, 1);
+
+    brute_force::search(dev_resources, bf_search_params, index,
+		raft::make_const_mdspan(d_centroids.view()), d_medoids_i64.view(), distances.view());
+	raft::copy(h_centroids_out, d_centroids.data_handle(), num_centers * dim,
+			raft::resource::get_cuda_stream(dev_resources));
+	raft::copy(h_medoids_i64.data_handle(), d_medoids_i64.data_handle(), num_centers, raft::resource::get_cuda_stream(dev_resources));
+    raft::resource::sync_stream(dev_resources);
+	for (size_t i = 0; i < num_centers; ++i) {
+		h_medoids[i] = static_cast<uint32_t>(h_medoids_i64(i));
+	}
+}
 
-int predict_gpu(raft::resources& handle,
+int predict_gpu(const raft::resources& handle,
                 const float* h_data,
                 size_t n_samples,
                 size_t dim,
@@ -214,7 +254,7 @@ bool is_gpu_available() {
 	return (cudaGetDeviceCount(&count) == cudaSuccess) && (count > 0);
 }
 
-void gpu_get_mem_info(raft::resources &dev_resources, size_t &gpu_free_mem,
+void gpu_get_mem_info(const raft::resources &dev_resources, size_t &gpu_free_mem,
 		size_t &gpu_total_mem) {
     int dev = raft::resource::get_device_id(dev_resources);
 
@@ -230,7 +270,7 @@ void gpu_get_mem_info(raft::resources &dev_resources, size_t &gpu_free_mem,
 }
 
 template<typename T>
-int brute_force_gpu(raft::resources &dev_resources, const T *h_data,
+int brute_force_gpu(const raft::resources &dev_resources, const T *h_data,
 		size_t n_samples, size_t dim, size_t k, const T *queries,
 		size_t n_queries, int64_t *h_labels) {
 	using namespace cuvs::neighbors;
@@ -281,7 +321,7 @@ template void vamana_build_and_write<uint8_t>(
 		std::string out_fname, int degree, int visited_size, float max_fraction,
 		int iters);
 
-template int brute_force_gpu<float>(raft::resources& dev_resources,
+template int brute_force_gpu<float>(const raft::resources &dev_resources,
 		const float* h_data,
         size_t n_samples,
         size_t dim,

thirdparty/DiskANN/src/mean_center_gpu.cu

@@ -0,0 +1,49 @@
+// Copyright (c) KIOXIA Corporation. All rights reserved.
+// Licensed under the MIT license.
+/*
+ * Copyright (c) 2024, 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 <cstddef>
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/mdspan.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/core/types.hpp>
+#include <raft/stats/mean.cuh>
+#include <raft/stats/mean_center.cuh>
+
+#include "diskann/diskann_gpu.h"
+
+// GPU implementation of the make_zero_mean step in generate_pq_pivots:
+// computes the per-dimension centroid of the (num_train x dim) row-major
+// training matrix and subtracts it from every vector in place. This mirrors
+// the CPU loops:
+//   centroid[d] = (1/num_train) * sum_p train_data[p * dim + d];
+//   train_data[p * dim + d] -= centroid[d];
+void mean_center_gpu(const raft::resources& dev_resources, float* d_train_data, size_t num_train, size_t dim,
+                     float* d_centroid_out) {
+    auto stream = raft::resource::get_cuda_stream(dev_resources);
+
+    // Population mean per column (divides by num_train, matching the CPU path).
+    auto d_train_data_view = raft::make_device_matrix_view<float, int64_t>(d_train_data, num_train, dim);
+    auto d_centroid_out_view = raft::make_device_vector_view<float, int64_t>(d_centroid_out, dim);
+    raft::stats::mean(dev_resources, raft::make_const_mdspan(d_train_data_view), d_centroid_out_view);
+
+    // Subtract the centroid from every row, writing the result back in place.
+    raft::stats::mean_center<raft::Apply::ALONG_ROWS>(dev_resources, raft::make_const_mdspan(d_train_data_view),
+                                                      raft::make_const_mdspan(d_centroid_out_view), d_train_data_view);
+}