Merge pull request #516 from allisonvacanti/atomic_compat

Replace cuda::atomic with legacy functions for old arch compatibility.
NVIDIA · Oct 4, 2023 · 265b574 · 265b574
2 parents 3de4abb + 904c17f
commit 265b574
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/cub/test/catch2_test_device_adjacent_difference_substract_left.cu b/cub/test/catch2_test_device_adjacent_difference_substract_left.cu
@@ -278,25 +278,28 @@ CUB_TEST("DeviceAdjacentDifference::SubtractLeftCopy works with large indexes",
 }
 
 struct invocation_counter {
-  cuda::atomic_ref<cuda::std::size_t, cuda::thread_scope_device> counts_;
 
-  __host__ explicit invocation_counter(cuda::std::size_t* counts) : counts_(*counts) {}
+  __host__ explicit invocation_counter(unsigned long long* addr) : counts_(addr) {}
 
   template<class T>
   __device__ T operator()(const T& lhs, const T& rhs) const noexcept {
-    ++counts_;
+    // Use legacy atomics to support testing on older archs:
+    atomicAdd(counts_, 1ull);
     return lhs - rhs;
   }
+
+private:
+  unsigned long long *counts_;
 };
 
 CUB_TEST("DeviceAdjacentDifference::SubtractLeftCopy uses right number of invocations", "[device][adjacent_difference]")
 {
   const int num_items = GENERATE_COPY(take(2, random(1, 1000000)));
-  thrust::device_vector<cuda::std::size_t> counts(1, 0);
+  thrust::device_vector<unsigned long long> counts(1, 0);
   adjacent_difference_subtract_left_copy(thrust::counting_iterator<cuda::std::size_t>{0},
                                          thrust::discard_iterator<>(),
                                          num_items,
                                          invocation_counter{thrust::raw_pointer_cast(counts.data())});
 
-  REQUIRE(counts.front() == num_items - 1);
+  REQUIRE(counts.front() == static_cast<unsigned long long>(num_items - 1));
 }
diff --git a/cub/test/catch2_test_device_adjacent_difference_substract_right.cu b/cub/test/catch2_test_device_adjacent_difference_substract_right.cu
@@ -320,25 +320,28 @@ CUB_TEST("DeviceAdjacentDifference::SubtractRightCopy works with large indexes",
 }
 
 struct invocation_counter {
-  cuda::atomic_ref<cuda::std::size_t, cuda::thread_scope_device> counts_;
 
-  __host__ explicit invocation_counter(cuda::std::size_t* counts) : counts_(*counts) {}
+  __host__ explicit invocation_counter(unsigned long long* addr) : counts_(addr) {}
 
   template<class T>
   __device__ T operator()(const T& lhs, const T& rhs) const noexcept {
-    ++counts_;
+    // Use legacy atomics to support testing on older archs:
+    atomicAdd(counts_, 1ull);
     return lhs - rhs;
   }
+
+private:
+  unsigned long long *counts_;
 };
 
 CUB_TEST("DeviceAdjacentDifference::SubtractRightCopy uses right number of invocations", "[device][adjacent_difference]")
 {
   const int num_items = GENERATE_COPY(take(2, random(1, 1000000)));
-  thrust::device_vector<cuda::std::size_t> counts(1, 0);
+  thrust::device_vector<unsigned long long> counts(1, 0);
   adjacent_difference_subtract_right_copy(thrust::counting_iterator<cuda::std::size_t>{0},
                                           thrust::discard_iterator<>(),
                                           num_items,
                                           invocation_counter{thrust::raw_pointer_cast(counts.data())});
 
-  REQUIRE(counts.front() == num_items - 1);
+  REQUIRE(counts.front() == static_cast<unsigned long long>(num_items - 1));
 }