alpha to weight alteration for precision

rapidsai · rapids-bot · May 27, 2021 · Apr 20, 2021 · Apr 20, 2021 · Apr 23, 2021
commit e0ef5b33ca647bb9337e48ee4568f3759b0fe017
diff --git a/cpp/include/raft/sparse/mst/detail/mst_kernels.cuh b/cpp/include/raft/sparse/mst/detail/mst_kernels.cuh
@@ -285,17 +285,27 @@ template <typename vertex_t, typename edge_t, typename weight_t>
 __global__ void alteration_kernel(const vertex_t v, const edge_t e,
                                   const edge_t* offsets,
                                   const vertex_t* indices,
-                                  const weight_t* weights, weight_t max,
+                                  const weight_t* weights, double max,
                                   weight_t* random_values,
-                                  weight_t* altered_weights) {
+                                  weight_t* altered_weights, int alpha,
+                                  bool use_alpha) {
   auto row = get_1D_idx<vertex_t>();
   if (row < v) {
     auto row_begin = offsets[row];
     auto row_end = offsets[row + 1];
     for (auto i = row_begin; i < row_end; i++) {
       auto column = indices[i];
-      altered_weights[i] =
-        weights[i] + max * (random_values[row] + random_values[column]);
+      // doing the later step explicity in double for precision
+      if (use_alpha) {
+        altered_weights[i] =
+          alpha * weights[i] + alpha * max *
+                                 (static_cast<double>(random_values[row]) +
+                                  static_cast<double>(random_values[column]));
+      } else {
+        altered_weights[i] =
+          weights[i] + max * (static_cast<double>(random_values[row]) +
+                              static_cast<double>(random_values[column]));
+      }
     }
   }
 }

diff --git a/cpp/include/raft/sparse/mst/detail/mst_solver_inl.cuh b/cpp/include/raft/sparse/mst/detail/mst_solver_inl.cuh
@@ -53,7 +53,7 @@ MST_solver<vertex_t, edge_t, weight_t>::MST_solver(
   const raft::handle_t& handle_, const edge_t* offsets_,
   const vertex_t* indices_, const weight_t* weights_, const vertex_t v_,
   const edge_t e_, vertex_t* color_, cudaStream_t stream_,
-  bool symmetrize_output_, bool initialize_colors_, int iterations_)
+  bool symmetrize_output_, bool initialize_colors_, int iterations_, int alpha_)
   : handle(handle_),
     offsets(offsets_),
     indices(indices_),
@@ -75,7 +75,8 @@ MST_solver<vertex_t, edge_t, weight_t>::MST_solver(
     stream(stream_),
     symmetrize_output(symmetrize_output_),
     initialize_colors(initialize_colors_),
-    iterations(iterations_) {
+    iterations(iterations_),
+    alpha(alpha_) {
   max_blocks = handle_.get_device_properties().maxGridSize[0];
   max_threads = handle_.get_device_properties().maxThreadsPerBlock;
   sm_count = handle_.get_device_properties().multiProcessorCount;
@@ -114,7 +115,9 @@ MST_solver<vertex_t, edge_t, weight_t>::solve() {
   timer0 = duration_us(stop - start);
 #endif
 
-  Graph_COO<vertex_t, edge_t, weight_t> mst_result(2 * v - 2, stream);
+  auto n_expected_edges = symmetrize_output ? 2 * v - 2 : v - 1;
+
+  Graph_COO<vertex_t, edge_t, weight_t> mst_result(n_expected_edges, stream);
 
   // Boruvka original formulation says "while more than 1 supervertex remains"
   // Here we adjust it to support disconnected components (spanning forest)
@@ -150,6 +153,10 @@ MST_solver<vertex_t, edge_t, weight_t>::solve() {
     timer3 += duration_us(stop - start);
 #endif
 
+    RAFT_EXPECTS(mst_edge_count[0] == n_expected_edges,
+                 "Number of edges found by MST is invalid. This may be due to "
+                 "loss in precision. Try increasing precision of weights.")
+
     if (prev_mst_edge_count[0] == mst_edge_count[0]) {
 #ifdef MST_TIME
       std::cout << "Iterations: " << i << std::endl;
@@ -209,7 +216,7 @@ struct alteration_functor {
 
 // Compute the uper bound for the alteration
 template <typename vertex_t, typename edge_t, typename weight_t>
-weight_t MST_solver<vertex_t, edge_t, weight_t>::alteration_max() {
+double MST_solver<vertex_t, edge_t, weight_t>::alteration_max() {
   auto policy = rmm::exec_policy(stream);
   rmm::device_vector<weight_t> tmp(e);
   thrust::device_ptr<const weight_t> weights_ptr(weights);
@@ -229,7 +236,7 @@ weight_t MST_solver<vertex_t, edge_t, weight_t>::alteration_max() {
   auto max =
     thrust::transform_reduce(policy, begin, end, alteration_functor<weight_t>(),
                              init, thrust::minimum<weight_t>());
-  return max / static_cast<weight_t>(2);
+  return max / static_cast<double>(2);
 }
 
 // Compute the alteration to make all undirected edge weight unique
@@ -240,7 +247,7 @@ void MST_solver<vertex_t, edge_t, weight_t>::alteration() {
   auto nblocks = std::min((v + nthreads - 1) / nthreads, max_blocks);
 
   // maximum alteration that does not change realtive weights order
-  weight_t max = alteration_max();
+  double max = alteration_max();
 
   // pool of rand values
   rmm::device_vector<weight_t> rand_values(v);
@@ -258,10 +265,12 @@ void MST_solver<vertex_t, edge_t, weight_t>::alteration() {
   RAFT_EXPECTS(curand_status == CURAND_STATUS_SUCCESS,
                "MST: CURAND cleanup failed");
 
+  bool use_alpha = max < 1e-3 && sizeof(weight_t) == 4;
+
   //Alterate the weights, make all undirected edge weight unique while keeping Wuv == Wvu
   detail::alteration_kernel<<<nblocks, nthreads, 0, stream>>>(
     v, e, offsets, indices, weights, max, rand_values.data().get(),
-    altered_weights.data().get());
+    altered_weights.data().get(), alpha, use_alpha);
 }
 
 // updates colors of vertices by propagating the lower color to the higher

diff --git a/cpp/include/raft/sparse/mst/mst.cuh b/cpp/include/raft/sparse/mst/mst.cuh
@@ -27,7 +27,7 @@ raft::Graph_COO<vertex_t, edge_t, weight_t> mst(
   const raft::handle_t& handle, edge_t const* offsets, vertex_t const* indices,
   weight_t const* weights, vertex_t const v, edge_t const e, vertex_t* color,
   cudaStream_t stream, bool symmetrize_output = true,
-  bool initialize_colors = true, int iterations = 0) {
+  bool initialize_colors = true, int iterations = 0, int alpha = 1e6) {
   MST_solver<vertex_t, edge_t, weight_t> mst_solver(
     handle, offsets, indices, weights, v, e, color, stream, symmetrize_output,
     initialize_colors, iterations);

diff --git a/cpp/include/raft/sparse/mst/mst_solver.cuh b/cpp/include/raft/sparse/mst/mst_solver.cuh
@@ -43,7 +43,7 @@ class MST_solver {
              const vertex_t* indices_, const weight_t* weights_,
              const vertex_t v_, const edge_t e_, vertex_t* color_,
              cudaStream_t stream_, bool symmetrize_output_,
-             bool initialize_colors_, int iterations_);
+             bool initialize_colors_, int iterations_ int alpha_);
 
   raft::Graph_COO<vertex_t, edge_t, weight_t> solve();
 
@@ -54,6 +54,7 @@ class MST_solver {
   cudaStream_t stream;
   bool symmetrize_output, initialize_colors;
   int iterations;
+  int alpha;
 
   //CSR
   const edge_t* offsets;
@@ -90,7 +91,7 @@ class MST_solver {
   void min_edge_per_supervertex();
   void check_termination();
   void alteration();
-  weight_t alteration_max();
+  double alteration_max();
   void append_src_dst_pair(vertex_t* mst_src, vertex_t* mst_dst,
                            weight_t* mst_weights);
 };