build: Add VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST ifdef

Add build option that disables Wave branch kernel test as it requires PTX. This test is enabled by default and can be disabled with: cmake -DVELOX_SKIP_WAVE_BRANCH_KERNEL_TEST=ON
facebookincubator · Feb 19, 2025 · 69ee8b3 · 69ee8b3
1 parent 6a9aac0
commit 69ee8b3
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 1 deletion.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -79,6 +79,7 @@ option(
 option(VELOX_MONO_LIBRARY "Build single unified library." OFF)
 option(ENABLE_ALL_WARNINGS "Enable -Wall and -Wextra compiler warnings." ON)
 option(VELOX_BUILD_SHARED "Build Velox as shared libraries." OFF)
+option(VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST "Disable Wave branch kernel test." OFF)
 # While it's possible to build both in one go we currently want to build either
 # static or shared.
 cmake_dependent_option(

diff --git a/velox/experimental/wave/common/Cuda.cu b/velox/experimental/wave/common/Cuda.cu
@@ -282,7 +282,7 @@ void Stream::deviceConstantToHostAsync(
 namespace {
 struct CallbackData {
   CallbackData(std::function<void()> callback)
-      : callback(std::move(callback)){};
+      : callback(std::move(callback)) {};
   std::function<void()> callback;
 };
 

diff --git a/velox/experimental/wave/common/tests/CMakeLists.txt b/velox/experimental/wave/common/tests/CMakeLists.txt
@@ -28,6 +28,10 @@ add_test(velox_wave_common_test velox_wave_common_test)
 set_tests_properties(velox_wave_common_test PROPERTIES LABELS cuda_driver)
 
 target_include_directories(velox_wave_common_test PRIVATE ../../../breeze)
+if(VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST)
+  target_compile_definitions(velox_wave_common_test
+                             PRIVATE VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST=1)
+endif()
 target_link_libraries(
   velox_wave_common_test
   velox_wave_common

diff --git a/velox/experimental/wave/common/tests/CudaTest.cpp b/velox/experimental/wave/common/tests/CudaTest.cpp
@@ -741,6 +741,7 @@ class RoundtripThread {
             stats.numAdds += op.param1 * op.param2 * 256;
             break;
 
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
           case OpCode::kAddBranch:
             VELOX_CHECK_LE(op.param1, kNumKB);
             if (stats.isCpu) {
@@ -754,6 +755,7 @@ class RoundtripThread {
             }
             stats.numAdds += op.param1 * op.param2 * 256;
             break;
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
 
           case OpCode::kAdd4x64:
             VELOX_CHECK_LE(op.param1, kNumKB);

diff --git a/velox/experimental/wave/common/tests/CudaTest.cu b/velox/experimental/wave/common/tests/CudaTest.cu
@@ -173,6 +173,8 @@ __global__ void addOneSwitchKernel(
   }
 }
 
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 #define BTCASE(nn, m)                              \
   asm volatile("BLK" nn ":");                      \
   temp = m + testFunc(temp, counter, flag, ptr).n; \
@@ -246,6 +248,8 @@ __global__ void addOneBranchKernel(
   }
 }
 
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 __global__ void addOneFuncStoreKernel(
     int32_t* numbers,
     int32_t size,
@@ -352,6 +356,8 @@ void TestStream::addOneFuncStore(
   CUDA_CHECK(cudaGetLastError());
 }
 
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 void TestStream::addOneBranch(
     int32_t* numbers,
     int32_t size,
@@ -371,6 +377,8 @@ void TestStream::addOneBranch(
   CUDA_CHECK(cudaGetLastError());
 }
 
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 void TestStream::addOneSwitch(
     int32_t* numbers,
     int32_t size,
@@ -779,6 +787,8 @@ __global__ void addOne4x64RegKernel(
   }
 }
 
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 #define BTCASE4(nn, m)                            \
   asm volatile("BLK" nn ":");                     \
   params = testFunc4(params, index + m, nullptr); \
@@ -844,6 +854,8 @@ void __global__ __launch_bounds__(1024) addOne4x64BranchKernel(
   }
 }
 
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
+
 __global__ void addOne4x64FuncKernel(
     int64_t* numbers,
     int32_t size,
@@ -948,10 +960,13 @@ void TestStream::addOne4x64(
       addOne4x64RegKernel<<<numBlocks, kBlockSize, smem, stream_->stream>>>(
           numbers, size, stride, repeats);
       break;
+
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
     case Add64Mode::k4Branch:
       addOne4x64BranchKernel<<<numBlocks, kBlockSize, smem, stream_->stream>>>(
           numbers, size, stride, repeats);
       break;
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
 
     case Add64Mode::k4Func:
       addOne4x64FuncKernel<<<numBlocks, kBlockSize, smem, stream_->stream>>>(
@@ -984,7 +999,9 @@ REGISTER_KERNEL("addOne", addOneKernel);
 REGISTER_KERNEL("addOneFunc", addOneFuncKernel);
 REGISTER_KERNEL("addOneWide", addOneWideKernel);
 REGISTER_KERNEL("addOneRandom", addOneRandomKernel);
+#ifndef VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
 REGISTER_KERNEL("add4x64branch", addOne4x64BranchKernel);
+#endif // !VELOX_SKIP_WAVE_BRANCH_KERNEL_TEST
 REGISTER_KERNEL("add4x64func", addOne4x64FuncKernel);
 REGISTER_KERNEL("add4x64smemfunc", addOne4x64SMemFuncKernel);