Merge branch 'branch-24.04' into ecg_c_api

.github/workflows/pr.yaml

@@ -162,7 +162,7 @@ jobs:
     with:
       build_type: pull-request
       script: ci/test_wheel_cugraph-pyg.sh
-      matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "11.8.0"))
+      matrix_filter: map(select(.ARCH == "amd64"))
   wheel-build-cugraph-equivariant:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]

.github/workflows/test.yaml

@@ -76,7 +76,7 @@ jobs:
       date: ${{ inputs.date }}
       sha: ${{ inputs.sha }}
       script: ci/test_wheel_cugraph-pyg.sh
-      matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "11.8.0"))
+      matrix_filter: map(select(.ARCH == "amd64"))
   wheel-tests-cugraph-equivariant:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]

ci/test_python.sh

@@ -166,7 +166,7 @@ if [[ "${RAPIDS_CUDA_VERSION}" == "11.8.0" ]]; then
       pylibcugraphops \
       cugraph \
       cugraph-dgl \
-      'dgl>=1.1.0.cu*' \
+      'dgl>=1.1.0.cu*,<=2.0.0.cu*' \
       'pytorch>=2.0' \
       'pytorch-cuda>=11.8'
 

ci/test_wheel_cugraph-dgl.sh

@@ -34,6 +34,6 @@ DGL_URL="https://data.dgl.ai/wheels/cu${PYTORCH_CUDA_VER}/repo.html"
 
 rapids-logger "Installing PyTorch and DGL"
 rapids-retry python -m pip install torch --index-url ${PYTORCH_URL}
-rapids-retry python -m pip install dgl --find-links ${DGL_URL}
+rapids-retry python -m pip install dgl==2.0.0 --find-links ${DGL_URL}
 
 python -m pytest python/cugraph-dgl/tests

cpp/include/cugraph/utilities/mask_utils.cuh

@@ -66,6 +66,50 @@ __device__ size_t count_set_bits(MaskIterator mask_first, size_t start_offset, s
     thrust::plus<size_t>{});
 }
 
+// @p n starts from 1
+template <typename MaskIterator>  // should be packed bool
+__device__ size_t
+find_nth_set_bits(MaskIterator mask_first, size_t start_offset, size_t num_bits, size_t n)
+{
+  static_assert(
+    std::is_same_v<typename thrust::iterator_traits<MaskIterator>::value_type, uint32_t>);
+  assert(n >= 1);
+  assert(n <= num_bits);
+
+  size_t pos{0};
+
+  mask_first   = mask_first + packed_bool_offset(start_offset);
+  start_offset = start_offset % packed_bools_per_word();
+  if (start_offset != 0) {
+    auto mask = ~packed_bool_partial_mask(start_offset);
+    if (start_offset + num_bits < packed_bools_per_word()) {
+      mask &= packed_bool_partial_mask(start_offset + num_bits);
+    }
+    auto word         = *mask_first & mask;
+    auto num_set_bits = __popc(word);
+    if (n <= num_set_bits) {
+      return static_cast<size_t>(__fns(word, start_offset, n)) - start_offset;
+    }
+    pos += __popc(mask);
+    n -= num_set_bits;
+    ++mask_first;
+  }
+
+  while (pos < num_bits) {
+    auto mask         = ((num_bits - pos) >= packed_bools_per_word())
+                          ? packed_bool_full_mask()
+                          : packed_bool_partial_mask(num_bits - pos);
+    auto word         = *mask_first & mask;
+    auto num_set_bits = __popc(word);
+    if (n <= num_set_bits) { return pos + static_cast<size_t>(__fns(word, 0, n)); }
+    pos += __popc(mask);
+    n -= num_set_bits;
+    ++mask_first;
+  }
+
+  return std::numeric_limits<size_t>::max();
+}
+
 template <typename InputIterator,
           typename MaskIterator,  // should be packed bool
           typename OutputIterator,