Skip to content

Commit

Permalink
Fix cuda graph support (#92)
Browse files Browse the repository at this point in the history
  • Loading branch information
tanmayv25 committed Jun 17, 2024
1 parent 9c5effd commit c44184e
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -82,17 +82,21 @@ FetchContent_MakeAvailable(repo-common repo-core repo-backend)
#
if(${TRITON_ENABLE_GPU})
find_package(CUDAToolkit REQUIRED)
message(STATUS "Using CUDA ${CUDA_VERSION}")
message(STATUS "Using CUDA ${CUDAToolkit_VERSION}")
if (WIN32)
set(CUDA_NVCC_FLAGS -std=c++17)
else ()
set(CUDA_NVCC_FLAGS -std=c++11)
endif ()

if(CUDA_VERSION VERSION_GREATER "10.1" OR CUDA_VERSION VERSION_EQUAL "10.1")
if(NOT DEFINED CUDAToolkit_VERSION)
message( FATAL_ERROR "Unable to determine CUDAToolkit_VERSION, CMake will exit." )
endif()

if(CUDAToolkit_VERSION VERSION_GREATER "10.1" OR CUDAToolkit_VERSION VERSION_EQUAL "10.1")
add_definitions(-DTRITON_ENABLE_CUDA_GRAPH=1)
else()
message(WARNING "CUDA ${CUDA_VERSION} does not support CUDA graphs.")
message(WARNING "CUDA ${CUDAToolkit_VERSION} does not support CUDA graphs.")
endif()
else()
message( FATAL_ERROR "TensorRT backend requires TRITON_ENABLE_GPU=1, CMake will exit." )
Expand Down

0 comments on commit c44184e

Please sign in to comment.