From a7c4ebd906dbfa09e6509ab9aac502cce3c7695a Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 17 Mar 2021 23:44:08 +1100
Subject: [PATCH] Remove literals passed to `device_uvector::set_element_async`
 (#1453)

After rapidsai/rmm#725 is merged, this PR updates cuspatial to eliminate passing literal values to device_uvector::set_element_async.

Companion PR to rapidsai/cuspatial#367

Authors:
  - Mark Harris (@harrism)

Approvers:
  - Seunghwa Kang (@seunghwak)
  - Alex Fender (@afender)
  - Andrei Schaffer (@aschaffer)

URL: https://github.com/rapidsai/cugraph/pull/1453
---
 cpp/src/experimental/graph.cu | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/cpp/src/experimental/graph.cu b/cpp/src/experimental/graph.cu
index 498bb4eaefe..5abe141dafd 100644
--- a/cpp/src/experimental/graph.cu
+++ b/cpp/src/experimental/graph.cu
@@ -304,9 +304,15 @@ graph_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu, std::enable_if_
 
     rmm::device_uvector<vertex_t> segment_offsets(detail::num_segments_per_vertex_partition + 1,
                                                   default_stream);
-    segment_offsets.set_element_async(0, 0, default_stream);
+
+    // temporaries are necessary because the &&-overload of device_uvector is deleted
+    // Note that we must sync `default_stream` before these temporaries go out of scope to
+    // avoid use after free. (The syncs are at the end of this function)
+    auto zero_vertex  = vertex_t{0};
+    auto vertex_count = static_cast<vertex_t>(degrees.size());
+    segment_offsets.set_element_async(0, zero_vertex, default_stream);
     segment_offsets.set_element_async(
-      detail::num_segments_per_vertex_partition, degrees.size(), default_stream);
+      detail::num_segments_per_vertex_partition, vertex_count, default_stream);
 
     thrust::upper_bound(rmm::exec_policy(default_stream)->on(default_stream),
                         degrees.begin(),
@@ -454,9 +460,16 @@ graph_t<vertex_t, edge_t, weight_t, store_transposed, multi_gpu, std::enable_if_
 
     rmm::device_uvector<vertex_t> segment_offsets(detail::num_segments_per_vertex_partition + 1,
                                                   default_stream);
-    segment_offsets.set_element_async(0, 0, default_stream);
+
+    // temporaries are necessary because the &&-overload of device_uvector is deleted
+    // Note that we must sync `default_stream` before these temporaries go out of scope to
+    // avoid use after free. (The syncs are at the end of this function)
+    auto zero_vertex  = vertex_t{0};
+    auto vertex_count = static_cast<vertex_t>(this->get_number_of_vertices());
+    segment_offsets.set_element_async(0, zero_vertex, default_stream);
+
     segment_offsets.set_element_async(
-      detail::num_segments_per_vertex_partition, this->get_number_of_vertices(), default_stream);
+      detail::num_segments_per_vertex_partition, vertex_count, default_stream);
 
     thrust::upper_bound(rmm::exec_policy(default_stream)->on(default_stream),
                         degree_first,