ci: add option to skip nvbench build

CMakeLists.txt

@@ -5,6 +5,8 @@ option(USE_CCACHE "Attempt using CCache to wrap the compilation" ON)
 option(USE_CXX11_ABI "Use the new C++-11 ABI, which is not backwards compatible." ON)
 option(USE_MANYLINUX "Build for manylinux" OFF)
 
+option(BUILD_NVBENCH "Build the nvbench binary" OFF)
+
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)

cmake/nvbench_binary.cmake

@@ -0,0 +1,40 @@
+include(CMakeParseArguments)
+
+function(nvbench_binary)
+  if(NOT BUILD_NVBENCH)
+    return()
+  endif()
+
+  cmake_parse_arguments(
+    NV_BINARY # prefix
+    "" # options
+    "NAME" # one value args
+    "HDRS;SRCS;COPTS;DEFINES;LINKOPTS;DEPS" # multi value args
+    ${ARGN}
+  )
+
+  add_executable(${NV_BINARY_NAME} "")
+  target_sources(${NV_BINARY_NAME} 
+    PRIVATE ${NV_BINARY_SRCS} ${NV_BINARY_HDRS}
+  )
+  target_link_libraries(${NV_BINARY_NAME}
+    PUBLIC 
+      ${NV_BINARY_DEPS}
+      nvbench::nvbench
+      nvbench::main
+    PRIVATE
+      ${NV_BINARY_LINKOPTS}
+  )
+  target_include_directories(${NV_BINARY_NAME}
+    PUBLIC
+      "$<BUILD_INTERFACE:${COMMON_INCLUDE_DIRS}>"      
+  )
+  target_compile_options(${NV_BINARY_NAME} 
+    PRIVATE 
+      ${NV_BINARY_COPTS}
+      -lineinfo
+  )
+  target_compile_definitions(${NV_BINARY_NAME} PUBLIC ${NV_BINARY_DEFINES})
+
+  add_executable(:${NV_BINARY_NAME} ALIAS ${NV_BINARY_NAME})
+endfunction()

src/kernels/CMakeLists.txt

@@ -28,7 +28,6 @@ cc_library(
 
 add_subdirectory(attention)
 add_subdirectory(quantization)
-add_subdirectory(bench)
 add_subdirectory(playground)
 add_subdirectory(triton)
 

src/kernels/attention/CMakeLists.txt

@@ -1,4 +1,5 @@
 include(cc_binary)
+include(nvbench_binary)
 include(cc_library)
 include(cc_test)
 
@@ -66,31 +67,23 @@ cc_test(
     torch
 )
 
-cc_binary(
+nvbench_binary(
   NAME 
     attention_sm80_bench
   SRCS 
     attention_sm80_bench.cu
   DEPS
-    nvbench::nvbench
-    nvbench::main
-    :attention.template
-  COPTS
-    -lineinfo 
+    :attention.template 
 )
 
-cc_binary(
+nvbench_binary(
   NAME 
     attention_sm80_pagedkv_bench
   SRCS 
     attention_sm80_pagedkv_bench.cu
   DEPS
     absl::random_random
-    nvbench::nvbench
-    nvbench::main
     :attention.template
-  COPTS
-    -lineinfo 
 )
 
 add_subdirectory(tools)

third_party/CMakeLists.txt

@@ -17,6 +17,8 @@ cc_library(
 )
 
 add_subdirectory(sentencepiece)
-add_subdirectory(nvbench)
+if (BUILD_NVBENCH)
+  add_subdirectory(nvbench)
+endif()
 # add_subdirectory(jinja2cpp)