Skip test if host does not support CUDA (facebookincubator#6866)

Summary: Pull Request resolved: facebookincubator#6866 Reviewed By: r-barnes Differential Revision: D49815045 fbshipit-source-id: 3c6dadd6b37b1f2cb4fee87d31d6f189d5c2070e
laithsakka · Oct 4, 2023 · 4018ddf · 4018ddf
1 parent 074b09b
commit 4018ddf
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 2 deletions.
diff --git a/velox/experimental/wave/common/tests/CudaTest.cpp b/velox/experimental/wave/common/tests/CudaTest.cpp
@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-#include <iostream>
+#include "velox/experimental/wave/common/tests/CudaTest.h"
 
+#include <cuda_runtime.h> // @manual
 #include <folly/executors/CPUThreadPoolExecutor.h>
 #include <folly/init/Init.h>
 #include <gtest/gtest.h>
@@ -31,7 +32,8 @@
 #include "velox/common/time/Timer.h"
 #include "velox/experimental/wave/common/GpuArena.h"
 #include "velox/experimental/wave/common/tests/BlockTest.h"
-#include "velox/experimental/wave/common/tests/CudaTest.h"
+
+#include <iostream>
 
 DEFINE_int32(num_streams, 0, "Number of paralll streams");
 DEFINE_int32(op_size, 0, "Size of invoke kernel (ints read and written)");
@@ -900,5 +902,9 @@ TEST_F(CudaTest, reduceMatrix) {
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   folly::init(&argc, &argv);
+  if (int device; cudaGetDevice(&device) != cudaSuccess) {
+    LOG(WARNING) << "No CUDA detected, skipping all tests";
+    return 0;
+  }
   return RUN_ALL_TESTS();
 }
diff --git a/velox/experimental/wave/common/tests/IdMapTest.cu b/velox/experimental/wave/common/tests/IdMapTest.cu
@@ -173,5 +173,9 @@ TEST(IdMapTest, overflowNoEmptyMarker) {
 int main(int argc, char** argv) {
   testing::InitGoogleTest(&argc, argv);
   folly::Init follyInit(&argc, &argv);
+  if (int device; cudaGetDevice(&device) != cudaSuccess) {
+    LOG(WARNING) << "No CUDA detected, skipping all tests";
+    return 0;
+  }
   return RUN_ALL_TESTS();
 }
diff --git a/velox/experimental/wave/exec/tests/AggregationTest.cpp b/velox/experimental/wave/exec/tests/AggregationTest.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cuda_runtime.h> // @manual
 #include <folly/init/Init.h>
 #include <gtest/gtest.h>
 #include <random>
@@ -46,6 +47,12 @@ class AggregationTest : public OperatorTestBase {
     OperatorTestBase::SetUpTestCase();
     wave::registerWave();
   }
+
+  void SetUp() override {
+    if (int device; cudaGetDevice(&device) != cudaSuccess) {
+      GTEST_SKIP() << "No CUDA detected, skipping all tests";
+    }
+  }
 };
 
 TEST_F(AggregationTest, singleKeySingleAggregate) {

diff --git a/velox/experimental/wave/exec/tests/FilterProjectTest.cpp b/velox/experimental/wave/exec/tests/FilterProjectTest.cpp
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include <cuda_runtime.h> // @manual
 #include "velox/dwio/common/tests/utils/BatchMaker.h"
 #include "velox/exec/PlanNodeStats.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
@@ -29,6 +30,9 @@ using facebook::velox::test::BatchMaker;
 class FilterProjectTest : public OperatorTestBase {
  protected:
   void SetUp() override {
+    if (int device; cudaGetDevice(&device) != cudaSuccess) {
+      GTEST_SKIP() << "No CUDA detected, skipping all tests";
+    }
     wave::registerWave();
   }