[Graphbolt] Support cuda compilation (#6383)

CMakeLists.txt

@@ -536,6 +536,8 @@ if(BUILD_GRAPHBOLT)
       ALL
       ${CMAKE_COMMAND} -E env
       CMAKE_COMMAND=${CMAKE_CMD}
+      CUDA_TOOLKIT_ROOT_DIR=${CUDA_TOOLKIT_ROOT_DIR}
+      USE_CUDA=${USE_CUDA}
       BINDIR=${BINDIR}
       CFLAGS=${CMAKE_C_FLAGS}
       CXXFLAGS=${CMAKE_CXX_FLAGS}
@@ -550,6 +552,8 @@ if(BUILD_GRAPHBOLT)
       ALL
       ${CMAKE_COMMAND} -E env
       CMAKE_COMMAND=${CMAKE_CMD}
+      CUDA_TOOLKIT_ROOT_DIR=${CUDA_TOOLKIT_ROOT_DIR}
+      USE_CUDA=${USE_CUDA}
       BINDIR=${CMAKE_CURRENT_BINARY_DIR}
       CFLAGS=${CMAKE_C_FLAGS}
       CXXFLAGS=${CMAKE_CXX_FLAGS}

graphbolt/CMakeLists.txt

@@ -2,6 +2,10 @@ cmake_minimum_required(VERSION 3.5)
 project(graphbolt C CXX)
 set (CMAKE_CXX_STANDARD 17)
 
+if(USE_CUDA)
+  enable_language(CUDA)
+endif()
+
 # Find PyTorch cmake files and PyTorch versions with the python interpreter
 # $PYTHON_INTERP ("python3" or "python" if empty)
 if(NOT PYTHON_INTERP)
@@ -38,6 +42,12 @@ set(BOLT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src")
 set(BOLT_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/include")
 file(GLOB BOLT_HEADERS ${BOLT_INCLUDE})
 file(GLOB BOLT_SRC ${BOLT_DIR}/*.cc)
+if(USE_CUDA)
+  file(GLOB BOLT_CUDA_SRC
+    ${BOLT_DIR}/cuda/*.cu
+  )
+  list(APPEND BOLT_SRC ${BOLT_CUDA_SRC})
+endif()
 
 add_library(${LIB_GRAPHBOLT_NAME} SHARED ${BOLT_SRC} ${BOLT_HEADERS})
 target_include_directories(${LIB_GRAPHBOLT_NAME} PRIVATE ${BOLT_DIR}

graphbolt/build.sh

@@ -12,6 +12,7 @@ else
   CPSOURCE=*.so
 fi
 
+CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DUSE_CUDA=$USE_CUDA"
 echo $CMAKE_FLAGS
 
 if [ $# -eq 0 ]; then