Clamp max_iter in OOC KMeans test

cpp/src/cluster/detail/kmeans_common.cuh

@@ -43,17 +43,23 @@
 #include <cub/iterator/arg_index_input_iterator.cuh>
 #include <cuda.h>
 #include <cuda/iterator>
+#include <thrust/copy.h>
 #include <thrust/for_each.h>
 #include <thrust/iterator/transform_iterator.h>
 
 #include <raft/linalg/add.cuh>
 
 #include <algorithm>
+#include <atomic>
 #include <cmath>
 #include <cstdio>
+#include <cstring>
 #include <ctime>
 #include <optional>
 #include <random>
+#include <string>
+#include <type_traits>
+#include <vector>
 
 namespace cuvs::cluster::kmeans::detail {
 
@@ -528,9 +534,74 @@ void compute_centroid_adjustments(
 {
   cudaStream_t stream = raft::resource::get_cuda_stream(handle);
   auto n_samples      = X.extent(0);
+  auto n_features     = X.extent(1);
 
   workspace.resize(n_samples, stream);
 
+  // ----- Nondeterminism diagnostic: fingerprint INPUTS (X, labels, weights) -----
+  // Compare two test runs by diffing the resulting log. If INPUT hashes match
+  // across runs but OUTPUT hashes do not, the divergence is the non-associative
+  // atomic accumulation inside raft::linalg::reduce_rows_by_key.
+  static std::atomic<int> s_call_id{0};
+  const int call_id = s_call_id.fetch_add(1);
+  {
+    // Materialize the (possibly lazy) labels iterator into a device buffer so
+    // we can read it on host without disturbing the original sequence.
+    rmm::device_uvector<IndexT> mat_labels(static_cast<std::size_t>(n_samples), stream);
+    thrust::copy_n(
+      raft::resource::get_thrust_policy(handle), cluster_labels, n_samples, mat_labels.data());
+
+    const std::size_t n_X_elems =
+      static_cast<std::size_t>(n_samples) * static_cast<std::size_t>(n_features);
+    std::vector<IndexT> h_labels(static_cast<std::size_t>(n_samples));
+    std::vector<DataT> h_X(n_X_elems);
+    std::vector<DataT> h_w(static_cast<std::size_t>(n_samples));
+    raft::copy(h_labels.data(), mat_labels.data(), n_samples, stream);
+    raft::copy(h_X.data(), X.data_handle(), n_X_elems, stream);
+    raft::copy(h_w.data(), sample_weights.data_handle(), n_samples, stream);
+    raft::resource::sync_stream(handle, stream);
+
+    auto fnv = [](auto const& buf) {
+      std::uint64_t h = 1469598103934665603ULL;
+      for (auto v : buf) {
+        std::uint64_t bits = 0;
+        if constexpr (std::is_floating_point_v<std::decay_t<decltype(v)>>) {
+          std::memcpy(&bits, &v, sizeof(v));
+        } else {
+          bits = static_cast<std::uint64_t>(v);
+        }
+        h ^= bits;
+        h *= 1099511628211ULL;
+      }
+      return h;
+    };
+
+    auto fmt_head = [&](auto const& buf) {
+      std::string s;
+      std::size_t hd = std::min<std::size_t>(8, buf.size());
+      for (std::size_t i = 0; i < hd; ++i) {
+        if (i) s += ',';
+        s += std::to_string(static_cast<long long>(buf[i]));
+      }
+      return s;
+    };
+
+    RAFT_LOG_INFO(
+      "compute_centroid_adjustments[call=%d]: INPUTS n_samples=%lld n_features=%lld "
+      "n_clusters=%lld reset_sums=%d X_hash=0x%016llx labels_hash=0x%016llx "
+      "weights_hash=0x%016llx labels_head=[%s]",
+      call_id,
+      static_cast<long long>(n_samples),
+      static_cast<long long>(n_features),
+      static_cast<long long>(n_clusters),
+      static_cast<int>(reset_sums),
+      static_cast<unsigned long long>(fnv(h_X)),
+      static_cast<unsigned long long>(fnv(h_labels)),
+      static_cast<unsigned long long>(fnv(h_w)),
+      fmt_head(h_labels).c_str());
+  }
+  // ----- end diagnostic for INPUTS -----
+
   raft::linalg::reduce_rows_by_key(X.data_handle(),
                                    X.extent(1),
                                    cluster_labels,
@@ -551,6 +622,48 @@ void compute_centroid_adjustments(
                                    n_clusters,
                                    stream,
                                    reset_sums);
+
+  // ----- Nondeterminism diagnostic: fingerprint OUTPUTS (centroid_sums, wpc) -----
+  {
+    const std::size_t n_sums_elems =
+      static_cast<std::size_t>(n_clusters) * static_cast<std::size_t>(n_features);
+    std::vector<DataT> h_sums(n_sums_elems);
+    std::vector<DataT> h_wpc(static_cast<std::size_t>(n_clusters));
+    raft::copy(h_sums.data(), centroid_sums.data_handle(), n_sums_elems, stream);
+    raft::copy(h_wpc.data(), weight_per_cluster.data_handle(), n_clusters, stream);
+    raft::resource::sync_stream(handle, stream);
+
+    auto fnv_fp = [](std::vector<DataT> const& buf) {
+      std::uint64_t h = 1469598103934665603ULL;
+      for (auto v : buf) {
+        std::uint64_t bits = 0;
+        std::memcpy(&bits, &v, sizeof(v));
+        h ^= bits;
+        h *= 1099511628211ULL;
+      }
+      return h;
+    };
+
+    auto fmt_head_fp = [](std::vector<DataT> const& buf) {
+      std::string s;
+      std::size_t hd = std::min<std::size_t>(8, buf.size());
+      for (std::size_t i = 0; i < hd; ++i) {
+        if (i) s += ',';
+        s += std::to_string(static_cast<double>(buf[i]));
+      }
+      return s;
+    };
+
+    RAFT_LOG_INFO(
+      "compute_centroid_adjustments[call=%d]: OUTPUTS sums_hash=0x%016llx wpc_hash=0x%016llx "
+      "sums_head=[%s] wpc_head=[%s]",
+      call_id,
+      static_cast<unsigned long long>(fnv_fp(h_sums)),
+      static_cast<unsigned long long>(fnv_fp(h_wpc)),
+      fmt_head_fp(h_sums).c_str(),
+      fmt_head_fp(h_wpc).c_str());
+  }
+  // ----- end diagnostic for OUTPUTS -----
 }
 /**
  * @brief Finalize centroids by dividing accumulated sums by counts.

cpp/src/cluster/detail/minClusterDistanceCompute.cu

@@ -11,6 +11,119 @@
 
 namespace cuvs::cluster::kmeans::detail {
 
+namespace {
+
+// Determinism self-test for the fused L2 nearest-neighbor reduction used by
+// minClusterAndDistanceCompute for L2Expanded / L2SqrtExpanded / CosineExpanded
+// metrics. Runs `fusedDistanceNNMinReduce` K times on bit-identical hardcoded
+// inputs and reports whether every output (both the argmin keys and the
+// distance values) is bitwise identical across the K runs.
+//
+// This is intentionally placed alongside the function under test so the same
+// code path is exercised. Triggered once per <DataT, IndexT> instantiation
+// from the first call to minClusterAndDistanceCompute via a static atomic.
+template <typename DataT, typename IndexT>
+void runFusedNNDeterminismSelfTest(raft::resources const& handle)
+{
+  constexpr IndexT n_samples  = 64;
+  constexpr IndexT n_features = 4;
+  constexpr IndexT n_clusters = 4;
+  constexpr int K             = 4;
+
+  cudaStream_t stream = raft::resource::get_cuda_stream(handle);
+
+  auto X         = raft::make_device_matrix<DataT, IndexT>(handle, n_samples, n_features);
+  auto centroids = raft::make_device_matrix<DataT, IndexT>(handle, n_clusters, n_features);
+
+  raft::random::RngState rng_x{12345ULL};
+  raft::random::RngState rng_c{67890ULL};
+  raft::random::uniform(handle,
+                        rng_x,
+                        X.data_handle(),
+                        static_cast<IndexT>(n_samples * n_features),
+                        DataT{-1},
+                        DataT{1});
+  raft::random::uniform(handle,
+                        rng_c,
+                        centroids.data_handle(),
+                        static_cast<IndexT>(n_clusters * n_features),
+                        DataT{-1},
+                        DataT{1});
+
+  auto L2NormX       = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
+  auto centroidsNorm = raft::make_device_vector<DataT, IndexT>(handle, n_clusters);
+  raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
+    handle,
+    raft::make_device_matrix_view<const DataT, IndexT>(X.data_handle(), n_samples, n_features),
+    L2NormX.view());
+  raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(
+    handle,
+    raft::make_device_matrix_view<const DataT, IndexT>(
+      centroids.data_handle(), n_clusters, n_features),
+    centroidsNorm.view());
+
+  rmm::device_uvector<char> workspace(sizeof(int) * static_cast<std::size_t>(n_samples), stream);
+
+  std::vector<std::vector<raft::KeyValuePair<IndexT, DataT>>> host_outputs(K);
+  for (int k = 0; k < K; ++k) {
+    auto out =
+      raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(handle, n_samples);
+    raft::KeyValuePair<IndexT, DataT> initial_value(0, std::numeric_limits<DataT>::max());
+    raft::matrix::fill(handle, out.view(), initial_value);
+
+    cuvs::distance::fusedDistanceNNMinReduce<DataT, raft::KeyValuePair<IndexT, DataT>, IndexT>(
+      out.data_handle(),
+      X.data_handle(),
+      centroids.data_handle(),
+      L2NormX.data_handle(),
+      centroidsNorm.data_handle(),
+      n_samples,
+      n_clusters,
+      n_features,
+      static_cast<void*>(workspace.data()),
+      /*sqrt=*/false,
+      /*initOutBuffer=*/false,
+      /*isRowMajor=*/true,
+      cuvs::distance::DistanceType::L2Expanded,
+      0.0f,
+      stream);
+
+    host_outputs[k].resize(static_cast<std::size_t>(n_samples));
+    raft::copy(host_outputs[k].data(), out.data_handle(), n_samples, stream);
+    raft::resource::sync_stream(handle, stream);
+  }
+
+  // Compare every (k > 0) run against run 0 element-by-element. Both the
+  // argmin key (the would-be cluster label) and the distance value are
+  // checked with bitwise equality.
+  int keys_diffs   = 0;
+  int values_diffs = 0;
+  for (int k = 1; k < K; ++k) {
+    for (IndexT i = 0; i < n_samples; ++i) {
+      if (host_outputs[k][i].key != host_outputs[0][i].key) { ++keys_diffs; }
+      if (std::memcmp(&host_outputs[k][i].value, &host_outputs[0][i].value, sizeof(DataT)) != 0) {
+        ++values_diffs;
+      }
+    }
+  }
+
+  RAFT_LOG_INFO(
+    "fusedDistanceNNMinReduce determinism self-test (%d runs, %lld samples, %lld clusters, "
+    "%lld features): keys=%s (%d diffs of %d compared), values=%s (%d bit-diffs of %d compared)",
+    K,
+    static_cast<long long>(n_samples),
+    static_cast<long long>(n_clusters),
+    static_cast<long long>(n_features),
+    keys_diffs == 0 ? "ALL MATCH" : "MISMATCH",
+    keys_diffs,
+    static_cast<int>((K - 1) * n_samples),
+    values_diffs == 0 ? "ALL MATCH" : "MISMATCH",
+    values_diffs,
+    static_cast<int>((K - 1) * n_samples));
+}
+
+}  // namespace
+
 // Calculates a <key, value> pair for every sample in input 'X' where key is an
 // index to an sample in 'centroids' (index of the nearest centroid) and 'value'
 // is the distance between the sample and the 'centroids[key]'.
@@ -35,6 +148,64 @@ void minClusterAndDistanceCompute(
                   metric == cuvs::distance::DistanceType::L2SqrtExpanded ||
                   metric == cuvs::distance::DistanceType::CosineExpanded;
 
+  // ----- One-shot determinism self-test for the fused L2 NN kernel -----
+  // Runs once per <DataT, IndexT> instantiation on the very first call.
+  // Feeds the fused kernel bit-identical hardcoded inputs K times and checks
+  // whether outputs are bitwise identical across the K runs.
+  {
+    static std::atomic<bool> s_self_test_done{false};
+    if (!s_self_test_done.exchange(true)) { runFusedNNDeterminismSelfTest<DataT, IndexT>(handle); }
+  }
+
+  // ----- Nondeterminism diagnostic: fingerprint INPUTS (X, centroids, L2NormX) -----
+  // Counterpart to the diagnostic in compute_centroid_adjustments. If calls
+  // with identical INPUT hashes here produce identical OUTPUT
+  // (keys_hash, values_hash), then label generation is deterministic and
+  // any downstream drift is from reduce_rows_by_key atomics, not from this
+  // kernel.
+  static std::atomic<int> s_mcad_call_id{0};
+  const int call_id = s_mcad_call_id.fetch_add(1);
+  {
+    const std::size_t n_X_elems =
+      static_cast<std::size_t>(n_samples) * static_cast<std::size_t>(n_features);
+    const std::size_t n_centroids_elems =
+      static_cast<std::size_t>(n_clusters) * static_cast<std::size_t>(n_features);
+
+    std::vector<DataT> h_X(n_X_elems);
+    std::vector<DataT> h_c(n_centroids_elems);
+    std::vector<DataT> h_norm(static_cast<std::size_t>(n_samples));
+    raft::copy(h_X.data(), X.data_handle(), n_X_elems, stream);
+    raft::copy(h_c.data(), centroids.data_handle(), n_centroids_elems, stream);
+    raft::copy(h_norm.data(), L2NormX.data_handle(), n_samples, stream);
+    raft::resource::sync_stream(handle, stream);
+
+    auto fnv_fp = [](std::vector<DataT> const& buf) {
+      std::uint64_t h = 1469598103934665603ULL;
+      for (auto v : buf) {
+        std::uint64_t bits = 0;
+        std::memcpy(&bits, &v, sizeof(v));
+        h ^= bits;
+        h *= 1099511628211ULL;
+      }
+      return h;
+    };
+
+    RAFT_LOG_INFO(
+      "minClusterAndDistanceCompute[call=%d]: INPUTS n_samples=%lld n_features=%lld "
+      "n_clusters=%lld metric=%d is_fused=%d X_hash=0x%016llx centroids_hash=0x%016llx "
+      "L2NormX_hash=0x%016llx",
+      call_id,
+      static_cast<long long>(n_samples),
+      static_cast<long long>(n_features),
+      static_cast<long long>(n_clusters),
+      static_cast<int>(metric),
+      static_cast<int>(is_fused),
+      static_cast<unsigned long long>(fnv_fp(h_X)),
+      static_cast<unsigned long long>(fnv_fp(h_c)),
+      static_cast<unsigned long long>(fnv_fp(h_norm)));
+  }
+  // ----- end diagnostic for INPUTS -----
+
   if (is_fused) {
     L2NormBuf_OR_DistBuf.resize(n_clusters, stream);
     auto centroidsNorm =
@@ -168,6 +339,49 @@ void minClusterAndDistanceCompute(
       }
     }
   }
+
+  // ----- Nondeterminism diagnostic: fingerprint OUTPUTS (keys, values) -----
+  // Split the KeyValuePair vector into keys (cluster indices = future labels)
+  // and values (distances). The keys are the inputs to
+  // compute_centroid_adjustments downstream.
+  {
+    std::vector<raft::KeyValuePair<IndexT, DataT>> h_kvp(static_cast<std::size_t>(n_samples));
+    raft::copy(h_kvp.data(), minClusterAndDistance.data_handle(), n_samples, stream);
+    raft::resource::sync_stream(handle, stream);
+
+    std::uint64_t keys_hash   = 1469598103934665603ULL;
+    std::uint64_t values_hash = 1469598103934665603ULL;
+    for (auto const& kv : h_kvp) {
+      keys_hash ^= static_cast<std::uint64_t>(kv.key);
+      keys_hash *= 1099511628211ULL;
+      std::uint64_t vbits = 0;
+      std::memcpy(&vbits, &kv.value, sizeof(kv.value));
+      values_hash ^= vbits;
+      values_hash *= 1099511628211ULL;
+    }
+
+    std::string keys_head;
+    std::string values_head;
+    const std::size_t hd = std::min<std::size_t>(8, h_kvp.size());
+    for (std::size_t i = 0; i < hd; ++i) {
+      if (i) {
+        keys_head += ',';
+        values_head += ',';
+      }
+      keys_head += std::to_string(static_cast<long long>(h_kvp[i].key));
+      values_head += std::to_string(static_cast<double>(h_kvp[i].value));
+    }
+
+    RAFT_LOG_INFO(
+      "minClusterAndDistanceCompute[call=%d]: OUTPUTS keys_hash=0x%016llx values_hash=0x%016llx "
+      "keys_head=[%s] values_head=[%s]",
+      call_id,
+      static_cast<unsigned long long>(keys_hash),
+      static_cast<unsigned long long>(values_hash),
+      keys_head.c_str(),
+      values_head.c_str());
+  }
+  // ----- end diagnostic for OUTPUTS -----
 }
 
 #define INSTANTIATE_MIN_CLUSTER_AND_DISTANCE(DataT, IndexT)                                    \

cpp/tests/cluster/kmeans.cu

@@ -381,6 +381,9 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam<KmeansBatchedInputs
     params.tol                 = testparams.tol;
     params.rng_state.seed      = 1;
     params.oversampling_factor = 0;
+    // Limit the number of iterations to ensure same number of iterations for reference and batched
+    // code paths.
+    params.max_iter = 3;
 
     auto stream = raft::resource::get_cuda_stream(handle);