fix timing bug with isend/irecv

rapidsai · Oct 6, 2020 · e0e696a · e0e696a
1 parent 97b98ed
commit e0e696a
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 10 deletions.
diff --git a/cpp/include/patterns/copy_to_adj_matrix_row_col.cuh b/cpp/include/patterns/copy_to_adj_matrix_row_col.cuh
@@ -206,6 +206,9 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
                        (graph_view.get_vertex_partition_first(comm_src_rank) -
                         graph_view.get_vertex_partition_first(row_comm_rank * col_comm_size)));
       } else {
+        CUDA_TRY(cudaStreamSynchronize(
+          handle.get_stream()));  // to ensure data to be sent are ready (FIXME: this can be removed
+                                  // if we use ncclSend in raft::comms)
         auto constexpr tuple_size = thrust_tuple_size_or_one<
           typename std::iterator_traits<VertexValueInputIterator>::value_type>::value;
         std::vector<raft::comms::request_t> requests(2 * tuple_size);
@@ -332,6 +335,10 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
                        vertex_value_input_first,
                        src_value_first);
 
+        CUDA_TRY(cudaStreamSynchronize(
+          handle.get_stream()));  // to ensure data to be sent are ready (FIXME: this can be removed
+                                  // if we use ncclSend in raft::comms)
+
         std::vector<raft::comms::request_t> value_requests(2 * (1 + tuple_size));
         device_isend<decltype(vertex_first), decltype(dst_vertices.begin())>(comm,
                                                                              vertex_first,

diff --git a/cpp/include/patterns/copy_v_transform_reduce_in_out_nbr.cuh b/cpp/include/patterns/copy_v_transform_reduce_in_out_nbr.cuh
@@ -572,6 +572,9 @@ void copy_v_transform_reduce_nbr(raft::handle_t const& handle,
                      minor_buffer_first + offset + size,
                      vertex_value_output_first);
       } else {
+        CUDA_TRY(cudaStreamSynchronize(
+          handle.get_stream()));  // to ensure data to be sent are ready (FIXME: this can be removed
+                                  // if we use ncclSend in raft::comms)
         auto constexpr tuple_size = thrust_tuple_size_or_one<
           typename std::iterator_traits<VertexValueOutputIterator>::value_type>::value;
         std::vector<raft::comms::request_t> requests(2 * tuple_size);

diff --git a/cpp/include/patterns/update_frontier_v_push_if_out_nbr.cuh b/cpp/include/patterns/update_frontier_v_push_if_out_nbr.cuh
@@ -620,11 +620,16 @@ void update_frontier_v_push_if_out_nbr(
                   buffer_requests.data() + (i + 1) * (1 + tuple_size),
                   std::numeric_limits<raft::comms::request_t>::max());
       } else {
-        comm.isend(detail::iter_to_raw_ptr(buffer_key_first + tx_offsets[i]),
-                   static_cast<size_t>(tx_counts[i]),
-                   comm_dst_rank,
-                   int{0} /* tag */,
-                   buffer_requests.data() + i * (1 + tuple_size));
+        CUDA_TRY(cudaStreamSynchronize(
+          handle.get_stream()));  // to ensure data to be sent are ready (FIXME: this can be removed
+                                  // if we use ncclSend in raft::comms)
+
+        device_isend(comm,
+                     detail::iter_to_raw_ptr(buffer_key_first + tx_offsets[i]),
+                     static_cast<size_t>(tx_counts[i]),
+                     comm_dst_rank,
+                     int{0} /* tag */,
+                     buffer_requests.data() + i * (1 + tuple_size));
         device_isend<decltype(buffer_payload_first), decltype(buffer_payload_first)>(
           comm,
           buffer_payload_first + tx_offsets[i],
@@ -654,11 +659,13 @@ void update_frontier_v_push_if_out_nbr(
                   buffer_requests.data() + (tx_counts.size() + i + 1) * (1 + tuple_size),
                   std::numeric_limits<raft::comms::request_t>::max());
       } else {
-        comm.irecv(detail::iter_to_raw_ptr(buffer_key_first + num_buffer_elements + rx_offsets[i]),
-                   static_cast<size_t>(rx_counts[i]),
-                   comm_src_rank,
-                   int{0} /* tag */,
-                   buffer_requests.data() + ((tx_counts.size() + i) * (1 + tuple_size)));
+        device_irecv(
+          comm,
+          detail::iter_to_raw_ptr(buffer_key_first + num_buffer_elements + rx_offsets[i]),
+          static_cast<size_t>(rx_counts[i]),
+          comm_src_rank,
+          int{0} /* tag */,
+          buffer_requests.data() + ((tx_counts.size() + i) * (1 + tuple_size)));
         device_irecv<decltype(buffer_payload_first), decltype(buffer_payload_first)>(
           comm,
           buffer_payload_first + num_buffer_elements + rx_offsets[i],