Merge branch 'branch-25.02' into remove_thrust_binary_function

rapidsai · Jan 24, 2025 · 71da983 · 71da983
2 parents 1a7dad4 + fe7852d
commit 71da983
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 8 deletions.
diff --git a/.devcontainer/cuda11.8-pip/devcontainer.json b/.devcontainer/cuda11.8-pip/devcontainer.json
@@ -5,7 +5,7 @@
     "args": {
       "CUDA": "11.8",
       "PYTHON_PACKAGE_MANAGER": "pip",
-      "BASE": "rapidsai/devcontainers:25.02-cpp-cuda11.8-ucx1.17.0-openmpi-ubuntu22.04"
+      "BASE": "rapidsai/devcontainers:25.02-cpp-cuda11.8-ucx1.18.0-openmpi-ubuntu22.04"
     }
   },
   "runArgs": [

diff --git a/.devcontainer/cuda12.5-pip/devcontainer.json b/.devcontainer/cuda12.5-pip/devcontainer.json
@@ -5,7 +5,7 @@
     "args": {
       "CUDA": "12.5",
       "PYTHON_PACKAGE_MANAGER": "pip",
-      "BASE": "rapidsai/devcontainers:25.02-cpp-cuda12.5-ucx1.17.0-openmpi-ubuntu22.04"
+      "BASE": "rapidsai/devcontainers:25.02-cpp-cuda12.5-ucx1.18.0-openmpi-ubuntu22.04"
     }
   },
   "runArgs": [

diff --git a/cpp/include/cugraph/edge_partition_edge_property_device_view.cuh b/cpp/include/cugraph/edge_partition_edge_property_device_view.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -138,9 +138,16 @@ class edge_partition_edge_property_device_view_t {
   {
     if constexpr (has_packed_bool_element) {
       static_assert(is_packed_bool, "unimplemented for thrust::tuple types.");
+      cuda::atomic_ref<uint32_t, cuda::thread_scope_device> word(
+        *(value_first_ + cugraph::packed_bool_offset(offset)));
       auto mask = cugraph::packed_bool_mask(offset);
-      auto old  = val ? atomicOr(value_first_ + cugraph::packed_bool_offset(offset), mask)
-                      : atomicAnd(value_first_ + cugraph::packed_bool_offset(offset), ~mask);
+      uint32_t old{};
+      if (compare == val) {
+        old = word.load(cuda::std::memory_order_relaxed);
+      } else {
+        old = val ? word.fetch_or(mask, cuda::std::memory_order_relaxed)
+                  : word.fetch_and(~mask, cuda::std::memory_order_relaxed);
+      }
       return static_cast<bool>(old & mask);
     } else {
       return cugraph::elementwise_atomic_cas(value_first_ + offset, compare, val);

diff --git a/cpp/include/cugraph/edge_partition_endpoint_property_device_view.cuh b/cpp/include/cugraph/edge_partition_endpoint_property_device_view.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -145,9 +145,16 @@ class edge_partition_endpoint_property_device_view_t {
     auto val_offset = value_offset(offset);
     if constexpr (has_packed_bool_element) {
       static_assert(is_packed_bool, "unimplemented for thrust::tuple types.");
+      cuda::atomic_ref<uint32_t, cuda::thread_scope_device> word(
+        *(value_first_ + cugraph::packed_bool_offset(val_offset)));
       auto mask = cugraph::packed_bool_mask(val_offset);
-      auto old  = val ? atomicOr(value_first_ + cugraph::packed_bool_offset(val_offset), mask)
-                      : atomicAnd(value_first_ + cugraph::packed_bool_offset(val_offset), ~mask);
+      uint32_t old{};
+      if (compare == val) {
+        old = word.load(cuda::std::memory_order_relaxed);
+      } else {
+        old = val ? word.fetch_or(mask, cuda::std::memory_order_relaxed)
+                  : word.fetch_and(~mask, cuda::std::memory_order_relaxed);
+      }
       return static_cast<bool>(old & mask);
     } else {
       return cugraph::elementwise_atomic_cas(value_first_ + val_offset, compare, val);