Update thrust/RMM deprecated calls (#1789)

PR essentially backports the fixes in #1707 for the deprecated RMM exec policy, and unpins the RMM nighlty to unblock local builds and CI and local development. Conflicts with the bigger PR should be minimal, and mainly around a single change: ``` rmm::exec_policy(handle.get_stream()) ``` to ``` handle.get_thrust_policy() ``` Authors: - Dante Gama Dessavre (https://github.com/dantegd) Approvers: - AJ Schmidt (https://github.com/ajschmidt8) - Chuck Hastings (https://github.com/ChuckHastings) - Seunghwa Kang (https://github.com/seunghwak) URL: #1789
rapidsai · Aug 25, 2021 · 80a3459 · 80a3459
1 parent fc0aef7
commit 80a3459
Show file tree

Hide file tree

Showing 83 changed files with 357 additions and 381 deletions.
diff --git a/ci/gpu/build.sh b/ci/gpu/build.sh
@@ -53,13 +53,12 @@ gpuci_logger "Activate conda env"
 . /opt/conda/etc/profile.d/conda.sh
 conda activate rapids
 
-# FIXME: return librmm and RMM to ${MINOR_VERSION}
 gpuci_logger "Install dependencies"
 gpuci_mamba_retry install -y \
       "libcudf=${MINOR_VERSION}" \
       "cudf=${MINOR_VERSION}" \
-      "librmm=21.10.00a210812" \
-      "rmm=21.10.00a210813" \
+      "librmm=${MINOR_VERSION}" \
+      "rmm=${MINOR_VERSION}" \
       "cudatoolkit=$CUDA_REL" \
       "dask-cudf=${MINOR_VERSION}" \
       "dask-cuda=${MINOR_VERSION}" \

diff --git a/conda/environments/cugraph_dev_cuda11.0.yml b/conda/environments/cugraph_dev_cuda11.0.yml
@@ -8,8 +8,8 @@ dependencies:
 - cudatoolkit=11.0
 - cudf=21.10.*
 - libcudf=21.10.*
-- rmm=21.10.00a210813
-- librmm=21.10.00a210812
+- rmm=21.10.*
+- librmm=21.10.*
 - dask>=2021.6.0
 - distributed>=2021.6.0
 - dask-cuda=21.10.*

diff --git a/conda/environments/cugraph_dev_cuda11.2.yml b/conda/environments/cugraph_dev_cuda11.2.yml
@@ -8,8 +8,8 @@ dependencies:
 - cudatoolkit=11.2
 - cudf=21.10.*
 - libcudf=21.10.*
-- rmm=21.10.00a210813
-- librmm=21.10.00a210812
+- rmm=21.10.*
+- librmm=21.10.*
 - dask>=2021.6.0
 - distributed>=2021.6.0
 - dask-cuda=21.10.*

diff --git a/conda/environments/cugraph_dev_cuda11.4.yml b/conda/environments/cugraph_dev_cuda11.4.yml
@@ -8,8 +8,8 @@ dependencies:
 - cudatoolkit=11.4
 - cudf=21.10.*
 - libcudf=21.10.*
-- rmm=21.10.00a210813
-- librmm=21.10.00a210812
+- rmm=21.10.*
+- librmm=21.10.*
 - dask>=2021.6.0
 - distributed>=2021.6.0
 - dask-cuda=21.10.*

diff --git a/conda/recipes/libcugraph/meta.yaml b/conda/recipes/libcugraph/meta.yaml
@@ -29,12 +29,11 @@ build:
     - CMAKE_CXX_COMPILER_LAUNCHER
     - CMAKE_CUDA_COMPILER_LAUNCHER
 
-# FIXME: return librmm to {{ minor_version }}.*
 requirements:
   build:
     - cmake>=3.20.1
     - cudatoolkit {{ cuda_version }}.*
-    - librmm=21.10.00a210812
+    - librmm {{ minor_version }}.*
     - boost-cpp>=1.66
     - nccl>=2.9.9
     - ucx-proc=*=gpu

diff --git a/cpp/cmake/thirdparty/get_cuhornet.cmake b/cpp/cmake/thirdparty/get_cuhornet.cmake
@@ -21,7 +21,7 @@ function(find_and_configure_cuhornet)
     FetchContent_Declare(
         cuhornet
         GIT_REPOSITORY    https://github.com/rapidsai/cuhornet.git
-        GIT_TAG           261399356e62bd76fa7628880f1a847aee713eed
+        GIT_TAG           4a1daa18405c0242370e16ce302dfa7eb5d9e857
         SOURCE_SUBDIR     hornet
     )
     FetchContent_GetProperties(cuhornet)

diff --git a/cpp/cmake/thirdparty/get_rmm.cmake b/cpp/cmake/thirdparty/get_rmm.cmake
@@ -26,16 +26,14 @@ function(find_and_configure_rmm VERSION)
         return()
     endif()
 
-    # FIXME: turn GIT_SHALLOW back to TRUE when changing GIT_TAG back
-    # to branch-${MAJOR_AND_MINOR}
     rapids_cpm_find(rmm ${VERSION}
         GLOBAL_TARGETS      rmm::rmm
         BUILD_EXPORT_SET    cugraph-exports
         INSTALL_EXPORT_SET  cugraph-exports
         CPM_ARGS
             GIT_REPOSITORY  https://github.com/rapidsai/rmm.git
-            GIT_TAG         23bbe745af1d988224b5498f7b8e3fe3720532d4
-            GIT_SHALLOW     FALSE
+            GIT_TAG         branch-${MAJOR_AND_MINOR}
+            GIT_SHALLOW     TRUE
             OPTIONS         "BUILD_TESTS OFF"
                             "BUILD_BENCHMARKS OFF"
                             "CUDA_STATIC_RUNTIME ${CUDA_STATIC_RUNTIME}"

diff --git a/cpp/include/cugraph/compute_partition.cuh b/cpp/include/cugraph/compute_partition.cuh
@@ -19,7 +19,7 @@
 
 #include <cugraph/graph.hpp>
 
-#include <rmm/thrust_rmm_allocator.h>
+#include <rmm/exec_policy.hpp>
 
 namespace cugraph {
 namespace detail {

diff --git a/cpp/include/cugraph/detail/graph_utils.cuh b/cpp/include/cugraph/detail/graph_utils.cuh
@@ -20,9 +20,9 @@
 #include <cugraph/utilities/dataframe_buffer.cuh>
 #include <cugraph/utilities/device_comm.cuh>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <raft/handle.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/sort.h>
 #include <thrust/tabulate.h>
@@ -78,19 +78,19 @@ rmm::device_uvector<edge_t> compute_major_degrees(
                                 [(detail::num_sparse_segments_per_vertex_partition + 2) * i +
                                  detail::num_sparse_segments_per_vertex_partition]
               : major_last;
-    thrust::transform(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+    thrust::transform(rmm::exec_policy(handle.get_stream()),
                       thrust::make_counting_iterator(vertex_t{0}),
                       thrust::make_counting_iterator(major_hypersparse_first - major_first),
                       local_degrees.begin(),
                       [p_offsets] __device__(auto i) { return p_offsets[i + 1] - p_offsets[i]; });
     if (use_dcs) {
       auto p_dcs_nzd_vertices   = (*adj_matrix_partition_dcs_nzd_vertices)[i];
       auto dcs_nzd_vertex_count = (*adj_matrix_partition_dcs_nzd_vertex_counts)[i];
-      thrust::fill(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+      thrust::fill(rmm::exec_policy(handle.get_stream()),
                    local_degrees.begin() + (major_hypersparse_first - major_first),
                    local_degrees.begin() + (major_last - major_first),
                    edge_t{0});
-      thrust::for_each(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+      thrust::for_each(rmm::exec_policy(handle.get_stream()),
                        thrust::make_counting_iterator(vertex_t{0}),
                        thrust::make_counting_iterator(dcs_nzd_vertex_count),
                        [p_offsets,
@@ -123,7 +123,7 @@ rmm::device_uvector<edge_t> compute_major_degrees(raft::handle_t const& handle,
                                                   vertex_t number_of_vertices)
 {
   rmm::device_uvector<edge_t> degrees(number_of_vertices, handle.get_stream());
-  thrust::tabulate(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+  thrust::tabulate(rmm::exec_policy(handle.get_stream()),
                    degrees.begin(),
                    degrees.end(),
                    [offsets] __device__(auto i) { return offsets[i + 1] - offsets[i]; });

diff --git a/cpp/include/cugraph/prims/copy_to_adj_matrix_row_col.cuh b/cpp/include/cugraph/prims/copy_to_adj_matrix_row_col.cuh
@@ -26,8 +26,8 @@
 #include <cugraph/utilities/thrust_tuple_utils.cuh>
 #include <cugraph/vertex_partition_device_view.cuh>
 
-#include <rmm/thrust_rmm_allocator.h>
 #include <raft/handle.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/copy.h>
 #include <thrust/execution_policy.h>
@@ -98,7 +98,7 @@ void copy_to_matrix_major(raft::handle_t const& handle,
     assert(graph_view.get_number_of_local_vertices() == GraphViewType::is_adj_matrix_transposed
              ? graph_view.get_number_of_local_adj_matrix_partition_cols()
              : graph_view.get_number_of_local_adj_matrix_partition_rows());
-    thrust::copy(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+    thrust::copy(rmm::exec_policy(handle.get_stream()),
                  vertex_value_input_first,
                  vertex_value_input_first + graph_view.get_number_of_local_vertices(),
                  matrix_major_value_output_first);
@@ -169,7 +169,7 @@ void copy_to_matrix_major(raft::handle_t const& handle,
           });
         // FIXME: this gather (and temporary buffer) is unnecessary if NCCL directly takes a
         // permutation iterator (and directly gathers to the internal buffer)
-        thrust::gather(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+        thrust::gather(rmm::exec_policy(handle.get_stream()),
                        map_first,
                        map_first + thrust::distance(vertex_first, vertex_last),
                        vertex_value_input_first,
@@ -190,7 +190,7 @@ void copy_to_matrix_major(raft::handle_t const& handle,
         // FIXME: this scatter is unnecessary if NCCL directly takes a permutation iterator (and
         // directly scatters from the internal buffer)
         thrust::scatter(
-          rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+          rmm::exec_policy(handle.get_stream()),
           rx_value_first,
           rx_value_first + rx_counts[i],
           map_first,
@@ -203,7 +203,7 @@ void copy_to_matrix_major(raft::handle_t const& handle,
         // FIXME: this scatter is unnecessary if NCCL directly takes a permutation iterator (and
         // directly scatters from the internal buffer)
         thrust::scatter(
-          rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+          rmm::exec_policy(handle.get_stream()),
           rx_value_first,
           rx_value_first + rx_counts[i],
           map_first,
@@ -226,7 +226,7 @@ void copy_to_matrix_major(raft::handle_t const& handle,
              ? graph_view.get_number_of_local_adj_matrix_partition_cols()
              : graph_view.get_number_of_local_adj_matrix_partition_rows());
     auto val_first = thrust::make_permutation_iterator(vertex_value_input_first, vertex_first);
-    thrust::scatter(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+    thrust::scatter(rmm::exec_policy(handle.get_stream()),
                     val_first,
                     val_first + thrust::distance(vertex_first, vertex_last),
                     vertex_first,
@@ -290,7 +290,7 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
     assert(graph_view.get_number_of_local_vertices() == GraphViewType::is_adj_matrix_transposed
              ? graph_view.get_number_of_local_adj_matrix_partition_rows()
              : graph_view.get_number_of_local_adj_matrix_partition_cols());
-    thrust::copy(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+    thrust::copy(rmm::exec_policy(handle.get_stream()),
                  vertex_value_input_first,
                  vertex_value_input_first + graph_view.get_number_of_local_vertices(),
                  matrix_minor_value_output_first);
@@ -360,7 +360,7 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
           });
         // FIXME: this gather (and temporary buffer) is unnecessary if NCCL directly takes a
         // permutation iterator (and directly gathers to the internal buffer)
-        thrust::gather(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+        thrust::gather(rmm::exec_policy(handle.get_stream()),
                        map_first,
                        map_first + thrust::distance(vertex_first, vertex_last),
                        vertex_value_input_first,
@@ -380,7 +380,7 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
           });
         // FIXME: this scatter is unnecessary if NCCL directly takes a permutation iterator (and
         // directly scatters from the internal buffer)
-        thrust::scatter(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+        thrust::scatter(rmm::exec_policy(handle.get_stream()),
                         rx_value_first,
                         rx_value_first + rx_counts[i],
                         map_first,
@@ -392,7 +392,7 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
           });
         // FIXME: this scatter is unnecessary if NCCL directly takes a permutation iterator (and
         // directly scatters from the internal buffer)
-        thrust::scatter(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+        thrust::scatter(rmm::exec_policy(handle.get_stream()),
                         rx_value_first,
                         rx_value_first + rx_counts[i],
                         map_first,
@@ -414,7 +414,7 @@ void copy_to_matrix_minor(raft::handle_t const& handle,
     assert(graph_view.get_number_of_local_vertices() ==
            graph_view.get_number_of_local_adj_matrix_partition_rows());
     auto val_first = thrust::make_permutation_iterator(vertex_value_input_first, vertex_first);
-    thrust::scatter(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+    thrust::scatter(rmm::exec_policy(handle.get_stream()),
                     val_first,
                     val_first + thrust::distance(vertex_first, vertex_last),
                     vertex_first,

diff --git a/cpp/include/cugraph/prims/copy_v_transform_reduce_in_out_nbr.cuh b/cpp/include/cugraph/prims/copy_v_transform_reduce_in_out_nbr.cuh
@@ -25,8 +25,8 @@
 #include <cugraph/utilities/host_barrier.hpp>
 
 #include <raft/cudart_utils.h>
-#include <rmm/thrust_rmm_allocator.h>
 #include <raft/handle.hpp>
+#include <rmm/exec_policy.hpp>
 
 #include <thrust/distance.h>
 #include <thrust/functional.h>
@@ -439,12 +439,12 @@ void copy_v_transform_reduce_nbr(raft::handle_t const& handle,
     }
 
     if (GraphViewType::is_multi_gpu) {
-      thrust::fill(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+      thrust::fill(rmm::exec_policy(handle.get_stream()),
                    minor_buffer_first,
                    minor_buffer_first + minor_tmp_buffer_size,
                    minor_init);
     } else {
-      thrust::fill(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+      thrust::fill(rmm::exec_policy(handle.get_stream()),
                    vertex_value_output_first,
                    vertex_value_output_first + graph_view.get_number_of_local_vertices(),
                    minor_init);
@@ -546,7 +546,7 @@ void copy_v_transform_reduce_nbr(raft::handle_t const& handle,
         if constexpr (update_major) {  // this is necessary as we don't visit every vertex in the
                                        // hypersparse segment in
                                        // for_all_major_for_all_nbr_hypersparse
-          thrust::fill(rmm::exec_policy(handle.get_stream())->on(handle.get_stream()),
+          thrust::fill(rmm::exec_policy(handle.get_stream()),
                        output_buffer_first + (*segment_offsets)[3],
                        output_buffer_first + (*segment_offsets)[4],
                        major_init);