Add cuda_atomic_ops_test to run_tests.sh

Pull Request resolved: pytorch#74482 Approved by: https://github.com/cpuhrsch
AnthonyBarbier · Mar 24, 2022 · 3024bcf · 3024bcf
1 parent 9270bcc
commit 3024bcf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/aten/src/ATen/test/cuda_atomic_ops_test.cu b/aten/src/ATen/test/cuda_atomic_ops_test.cu
@@ -1,6 +1,7 @@
 #include <gtest/gtest.h>
 #include <ATen/cuda/Atomic.cuh>
 #include <c10/test/util/Macros.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAException.h>
 
 #include <cmath>
@@ -75,7 +76,7 @@ void test_atomic_mul() {
   for (int i = 0; i < arraysize; ++i) {
     a[i] = 2;
     sum[i] = 2;
-    answer[i] = pow(sum[i], static_cast<T>(factor));
+    answer[i] = pow(sum[i], static_cast<T>(factor + 1));
   }
 
   cudaMalloc((void**)&ad, arraysize * sizeof(T));
@@ -98,6 +99,7 @@ void test_atomic_mul() {
 }
 
 TEST(TestAtomicOps, TestAtomicAdd) {
+  if (!at::cuda::is_available()) return;
   test_atomic_add<uint8_t>();
   test_atomic_add<int8_t>();
   test_atomic_add<int16_t>();
@@ -113,6 +115,7 @@ TEST(TestAtomicOps, TestAtomicAdd) {
 }
 
 TEST(TestAtomicOps, DISABLED_ON_WINDOWS(TestAtomicMul)) {
+  if (!at::cuda::is_available()) return;
   test_atomic_mul<at::BFloat16>();
   test_atomic_mul<at::Half>();
   test_atomic_mul<float>();

diff --git a/aten/tools/run_tests.sh b/aten/tools/run_tests.sh
@@ -64,6 +64,9 @@ fi
 if [[ -x ./cuda_cub_test ]]; then
   ./cuda_cub_test
 fi
+if [[ -x ./cuda_atomic_ops_test ]]; then
+  ./cuda_atomic_ops_test
+fi
 if [ "$VALGRIND" == "ON" ]; then
   valgrind --suppressions="$VALGRIND_SUP" --error-exitcode=1 ./basic --gtest_filter='-*CUDA'
   if [[ -x ./tensor_interop_test ]]; then