[BUG] Fix CAGRA graph optimization bug

cpp/src/neighbors/detail/cagra/graph_core.cuh

@@ -1102,7 +1102,6 @@ void optimize(
     }
   }
 
-  auto pruned_graph = raft::make_host_matrix<uint32_t, int64_t>(graph_size, output_graph_degree);
   {
     //
     // Prune kNN graph
@@ -1191,13 +1190,14 @@ void optimize(
     const auto num_full = host_stats.data_handle()[1];
 
     // Create pruned kNN graph
+    bool invalid_neighbor_list = false;
 #pragma omp parallel for
     for (uint64_t i = 0; i < graph_size; i++) {
       // Find the `output_graph_degree` smallest detourable count nodes by checking the detourable
       // count of the neighbors while increasing the target detourable count from zero.
-      uint64_t pk         = 0;
+      uint64_t pk         = mst_graph_num_edges_ptr[i];
       uint32_t num_detour = 0;
-      while (pk < output_graph_degree) {
+      for (uint32_t l = 0; l < input_graph_degree && pk < output_graph_degree; l++) {
         uint32_t next_num_detour = std::numeric_limits<uint32_t>::max();
         for (uint64_t k = 0; k < input_graph_degree; k++) {
           const auto num_detour_k = detour_count.data_handle()[k + (input_graph_degree * i)];
@@ -1208,23 +1208,41 @@ void optimize(
 
           // Store the neighbor index if its detourable count is equal to `num_detour`.
           if (num_detour_k != num_detour) { continue; }
-          output_graph_ptr[pk + (output_graph_degree * i)] =
-            input_graph_ptr[k + (input_graph_degree * i)];
-          pk += 1;
+
+          // Check duplication and append
+          const auto candidate_node = input_graph_ptr[k + (input_graph_degree * i)];
+          bool dup                  = false;
+          for (uint32_t dk = 0; dk < pk; dk++) {
+            if (candidate_node == output_graph_ptr[i * output_graph_degree + dk]) {
+              dup = true;
+              break;
+            }
+          }
+          if (!dup && candidate_node < graph_size) {
+            output_graph_ptr[i * output_graph_degree + pk] = candidate_node;
+            pk += 1;
+          }
           if (pk >= output_graph_degree) break;
         }
         if (pk >= output_graph_degree) break;
 
         assert(next_num_detour != std::numeric_limits<uint32_t>::max());
         num_detour = next_num_detour;
       }
-      RAFT_EXPECTS(
-        pk == output_graph_degree,
-        "Couldn't find the output_graph_degree (%lu) smallest detourable count nodes for "
-        "node %lu in the rank-based node reranking process",
-        output_graph_degree,
-        i);
+      if (pk != output_graph_degree) {
+        RAFT_LOG_DEBUG(
+          "Couldn't find the output_graph_degree (%lu) smallest detourable count nodes for "
+          "node %lu in the rank-based node reranking process",
+          output_graph_degree,
+          i);
+        invalid_neighbor_list = true;
+      }
     }
+    RAFT_EXPECTS(
+      !invalid_neighbor_list,
+      "Could not generate an intermediate CAGRA graph because the initial kNN graph contains too "
+      "many invalid or duplicated neighbor nodes. This error can occur, for example, if too many "
+      "overflows occur during the norm computation between the dataset vectors.");
 
     const double time_prune_end = cur_time();
     RAFT_LOG_DEBUG(
@@ -1313,36 +1331,14 @@ void optimize(
 
 #pragma omp parallel for
     for (uint64_t i = 0; i < graph_size; i++) {
-      auto my_fwd_graph = pruned_graph.data_handle() + (output_graph_degree * i);
       auto my_rev_graph = rev_graph.data_handle() + (output_graph_degree * i);
       auto my_out_graph = output_graph_ptr + (output_graph_degree * i);
-      uint32_t kf       = 0;
       uint32_t k        = mst_graph_num_edges_ptr[i];
 
       const auto num_protected_edges = std::max<uint64_t>(k, output_graph_degree / 2);
       assert(num_protected_edges <= output_graph_degree);
       if (num_protected_edges == output_graph_degree) continue;
 
-      // Append edges from the pruned graph to output graph
-      while (k < output_graph_degree && kf < output_graph_degree) {
-        if (my_fwd_graph[kf] < graph_size) {
-          auto flag_match = false;
-          for (uint32_t kk = 0; kk < k; kk++) {
-            if (my_out_graph[kk] == my_fwd_graph[kf]) {
-              flag_match = true;
-              break;
-            }
-          }
-          if (!flag_match) {
-            my_out_graph[k] = my_fwd_graph[kf];
-            k += 1;
-          }
-        }
-        kf += 1;
-      }
-      assert(k == output_graph_degree);
-      assert(kf <= output_graph_degree);
-
       // Replace some edges of the output graph with edges of the reverse graph.
       auto kr = std::min<uint32_t>(rev_graph_count.data_handle()[i], output_graph_degree);
       while (kr) {
@@ -1419,6 +1415,36 @@ void optimize(
                      (double)sum_hist / graph_size);
     }
   }
+
+  // Check duplication and out-of-range indices
+  {
+    uint64_t num_dup = 0;
+    uint64_t num_oor = 0;
+#pragma omp parallel for reduction(+ : num_dup) reduction(+ : num_oor)
+    for (uint64_t i = 0; i < graph_size; i++) {
+      auto my_out_graph = output_graph_ptr + (output_graph_degree * i);
+      for (uint32_t j = 0; j < output_graph_degree; j++) {
+        const auto neighbor_a = my_out_graph[j];
+
+        // Check oor
+        if (neighbor_a > graph_size) {
+          num_oor++;
+          continue;
+        }
+
+        // Check duplication
+        for (uint32_t k = j + 1; k < output_graph_degree; k++) {
+          const auto neighbor_b = my_out_graph[k];
+          if (neighbor_a == neighbor_b) { num_dup++; }
+        }
+      }
+    }
+    RAFT_EXPECTS(
+      num_dup == 0, "%lu duplicated node(s) are found in the generated CAGRA graph", num_dup);
+    RAFT_EXPECTS(num_oor == 0,
+                 "%lu out-of-range index node(s) are found in the generated CAGRA graph",
+                 num_oor);
+  }
 }
 
 }  // namespace cuvs::neighbors::cagra::detail::graph

cpp/test/neighbors/ann_cagra/bug_extreme_inputs_oob.cu

@@ -40,8 +40,14 @@ class cagra_extreme_inputs_oob_test : public ::testing::Test {
     ix_ps.graph_degree              = 64;
     ix_ps.intermediate_graph_degree = 128;
 
-    [[maybe_unused]] auto ix = cagra::build(res, ix_ps, raft::make_const_mdspan(dataset->view()));
-    raft::resource::sync_stream(res);
+    try {
+      [[maybe_unused]] auto ix = cagra::build(res, ix_ps, raft::make_const_mdspan(dataset->view()));
+      raft::resource::sync_stream(res);
+    } catch (const std::exception&) {
+      SUCCEED();
+      return;
+    }
+    FAIL();
   }
 
   void SetUp() override