Merge pull request #50 from ChuckHastings/fea_specify_stream_in_devic…

…e_buffer RMM is making streams explicit in the device_buffer
rapidsai · May 20, 2021 · 76a8b53 · 76a8b53
2 parents 6d2fc89 + 4c68d8c
commit 76a8b53
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/hornet/include/Core/SoA/impl/SoADataSort.i.cuh b/hornet/include/Core/SoA/impl/SoADataSort.i.cuh
@@ -310,7 +310,7 @@ cub_segmented_sort(CSoAPtr<EdgeTypes...> &soa, degree_t capacity, degree_t segme
   cub::DeviceSegmentedRadixSort::SortKeys(
       NULL, tempStorageBytes, in_edges, out_edges, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);
-  rmm::device_buffer tempStorage(tempStorageBytes);
+  rmm::device_buffer tempStorage(tempStorageBytes, cuda_stream_view{});
   cub::DeviceSegmentedRadixSort::SortKeys(
       tempStorage.data(), tempStorageBytes, in_edges, out_edges, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);
@@ -344,7 +344,7 @@ cub_segmented_sort(CSoAPtr<EdgeTypes...> &soa, degree_t capacity, degree_t segme
   cub::DeviceSegmentedRadixSort::SortPairs(
       NULL, tempStorageBytes, in_key, out_key, in_val, out_val, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);
-  rmm::device_buffer tempStorage(tempStorageBytes);
+  rmm::device_buffer tempStorage(tempStorageBytes, cuda_stream_view{});
   cub::DeviceSegmentedRadixSort::SortPairs(
       tempStorage.data(), tempStorageBytes, in_key, out_key, in_val, out_val, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);
@@ -382,7 +382,7 @@ cub_segmented_sort(CSoAPtr<EdgeTypes...> &soa, degree_t capacity, degree_t segme
   cub::DeviceSegmentedRadixSort::SortPairs(
       NULL, tempStorageBytes, in_key, out_key, in_val, out_val, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);
-  rmm::device_buffer tempStorage(tempStorageBytes);
+  rmm::device_buffer tempStorage(tempStorageBytes, cuda_stream_view{});
   cub::DeviceSegmentedRadixSort::SortPairs(
       tempStorage.data(), tempStorageBytes, in_key, out_key, in_val, out_val, capacity,
       offsets.size() - 1, offsets.data().get(), offsets.data().get() + 1);

diff --git a/xlib/include/Device/Primitives/CubWrapper.cuh b/xlib/include/Device/Primitives/CubWrapper.cuh
@@ -64,7 +64,7 @@ protected:
     explicit CubWrapper(const int num_items) noexcept : _num_items(num_items) {}
     ~CubWrapper() noexcept { release(); }
 
-    mutable rmm::device_buffer _d_temp_storage { 0 };
+    mutable rmm::device_buffer _d_temp_storage { 0, cuda_stream_view{} };
     size_t _temp_storage_bytes { 0 };
     int    _num_items          { 0 };
 };