[WIP] Fix batched silhouette score reduction aliasing

cpp/src/stats/detail/batched/silhouette_score.cuh

@@ -247,12 +247,15 @@ value_t silhouette_score(
 
   raft::resource::sync_stream_pool(handle);
 
-  // calculating row-wise minimum in b
+  // Keep the row-wise reduction output separate from b. The input is an
+  // n_rows x n_labels matrix, so writing an n_rows vector at b_ptr aliases
+  // matrix elements that may still be read by the reduction.
+  rmm::device_uvector<value_t> b_min(n_rows, stream);
   raft::linalg::reduce<raft::Apply::ALONG_ROWS>(
     handle,
     raft::make_device_matrix_view<const value_t, value_idx, raft::row_major>(
       b_ptr, n_rows, n_labels),
-    raft::make_device_vector_view<value_t, value_idx>(b_ptr, n_rows),
+    raft::make_device_vector_view<value_t, value_idx>(b_min.data(), n_rows),
     std::numeric_limits<value_t>::max(),
     false,
     raft::identity_op(),
@@ -265,7 +268,7 @@ value_t silhouette_score(
     cuvs::stats::detail::SilOp<value_t>(),
     raft::make_const_mdspan(raft::make_device_vector_view<const value_t, value_idx>(a_ptr, n_rows)),
     raft::make_const_mdspan(
-      raft::make_device_vector_view<const value_t, value_idx>(b_ptr, n_rows)));
+      raft::make_device_vector_view<const value_t, value_idx>(b_min.data(), n_rows)));
 
   auto sum = raft::make_device_vector<value_t, value_idx>(handle, 1);
   raft::linalg::reduce<raft::Apply::ALONG_COLUMNS>(

cpp/tests/stats/silhouette_score.cu

@@ -16,10 +16,12 @@
 #include <gtest/gtest.h>
 
 #include <algorithm>
+#include <array>
 #include <cmath>
 #include <iostream>
 #include <memory>
 #include <random>
+#include <utility>
 
 namespace cuvs {
 namespace stats {
@@ -224,40 +226,69 @@ TEST_P(silhouetteScoreTestClass, Result)
 }
 INSTANTIATE_TEST_CASE_P(silhouetteScore, silhouetteScoreTestClass, ::testing::ValuesIn(inputs));
 
-TEST(silhouetteScore, BatchedStreamPoolOrdering)
+TEST(silhouetteScore, BatchedMatchesNonBatchedAcrossConfigurations)
 {
-  constexpr int64_t n_rows = 4096;
-  constexpr int64_t n_cols = 2;
-  constexpr int n_labels   = 2;
-
+  constexpr int64_t n_rows  = 1000;
+  constexpr int64_t n_cols  = 2;
+  constexpr int n_labels    = 2;
+  constexpr float tolerance = 1e-4f;
+  constexpr std::array<int64_t, 3> chunks{n_rows, n_rows / 3, n_rows / 5};
+  constexpr std::array metrics{cuvs::distance::DistanceType::CosineExpanded,
+                               cuvs::distance::DistanceType::L2SqrtUnexpanded,
+                               cuvs::distance::DistanceType::L2Expanded,
+                               cuvs::distance::DistanceType::L1};
+
+  std::mt19937 rng(193);
+  std::uniform_real_distribution<float> centers(-1.0f, 1.0f);
+  std::normal_distribution<float> noise(0.0f, 1.5f);
+  std::array<std::array<float, n_cols>, n_labels> center{};
+  for (auto& c : center) {
+    for (auto& x : c) {
+      x = centers(rng);
+    }
+  }
+  std::vector<int64_t> order(n_rows);
+  for (int64_t i = 0; i < n_rows; ++i) {
+    order[i] = i;
+  }
+  std::shuffle(order.begin(), order.end(), rng);
   std::vector<float> X(n_rows * n_cols);
   std::vector<int> labels(n_rows);
-  for (int64_t i = 0; i < n_rows; ++i) {
-    X[2 * i]     = std::sin(0.01f * i);
-    X[2 * i + 1] = std::cos(0.013f * i);
-    labels[i]    = i % n_labels;
+  for (int64_t row = 0; row < n_rows; ++row) {
+    auto label  = static_cast<int>(order[row] / (n_rows / n_labels));
+    labels[row] = label;
+    for (int64_t col = 0; col < n_cols; ++col) {
+      X[row * n_cols + col] = center[label][col] + noise(rng);
+    }
   }
 
-  raft::resources handle;
-  raft::resource::set_cuda_stream_pool(handle, std::make_shared<rmm::cuda_stream_pool>(4));
-  auto stream = raft::resource::get_cuda_stream(handle);
+  raft::resources default_handle;
+  raft::resources pool_handle;
+  raft::resource::set_cuda_stream_pool(pool_handle, std::make_shared<rmm::cuda_stream_pool>(4));
+  auto stream = raft::resource::get_cuda_stream(default_handle);
 
   rmm::device_uvector<float> d_X(X.size(), stream);
   rmm::device_uvector<int> d_labels(labels.size(), stream);
   raft::update_device(d_X.data(), X.data(), X.size(), stream);
   raft::update_device(d_labels.data(), labels.data(), labels.size(), stream);
+  raft::resource::sync_stream(default_handle);
 
   auto X_view = raft::make_device_matrix_view<const float, int64_t>(d_X.data(), n_rows, n_cols);
   auto labels_view = raft::make_device_vector_view<const int, int64_t>(d_labels.data(), n_rows);
-  constexpr auto metric = cuvs::distance::DistanceType::L2SqrtUnexpanded;
-
-  auto expected =
-    cuvs::stats::silhouette_score(handle, X_view, labels_view, std::nullopt, n_labels, metric);
 
-  for (int repeat = 0; repeat < 8; ++repeat) {
-    auto actual = cuvs::stats::silhouette_score_batched(
-      handle, X_view, labels_view, std::nullopt, n_labels, n_rows, metric);
-    ASSERT_NEAR(actual, expected, 1e-4f);
+  for (auto metric : metrics) {
+    auto expected = cuvs::stats::silhouette_score(
+      default_handle, X_view, labels_view, std::nullopt, n_labels, metric);
+    for (auto const& handle :
+         {std::pair{"default", &default_handle}, std::pair{"pool", &pool_handle}}) {
+      for (auto chunk : chunks) {
+        SCOPED_TRACE(::testing::Message() << "handle=" << handle.first << " metric="
+                                          << static_cast<int>(metric) << " chunk=" << chunk);
+        auto actual = cuvs::stats::silhouette_score_batched(
+          *handle.second, X_view, labels_view, std::nullopt, n_labels, chunk, metric);
+        ASSERT_NEAR(actual, expected, tolerance);
+      }
+    }
   }
 }