diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index af92e7aa1f92d2db6f0d083162d7b77147e19cb1..df005dfa4af1625a1c02e3a959062db139d7de62 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -130,12 +130,28 @@ set( headers set( libs m ) +set(CUDNN_FLAGS) + +if (USE_CUDNN) + find_library( CUDNN_LIBRARIES + NAMES cudnn libcudnn.so libcudnn.so.5 + PATHS ${CUDA_TOOLKIT_ROOT_DIR}/targets/aarch64-linux + ${CUDA_TOOLKIT_ROOT_DIR} + PATH_SUFFIXES lib lib64 + ) + + set(CUDNN_FLAGS -DCUDNN ) + + list( APPEND libs ${CUDNN_LIBRARIES} ) + add_definitions( -DCUDNN ) +endif() + if (USE_GPU) find_package(CUDA QUIET REQUIRED) include_directories( ${CUDA_TOOLKIT_ROOT_DIR}/include ) add_definitions( -DGPU ) - set( CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} + set( CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS} ${CUDNN_FLAGS} -gencode arch=compute_20,code=[sm_20,sm_21] -gencode arch=compute_30,code=sm_30 -gencode arch=compute_35,code=sm_35 @@ -151,18 +167,6 @@ if (USE_GPU) ) endif() -if (USE_CUDNN) - find_library( CUDNN_LIBRARIES - NAMES cudnn libcudnn.so libcudnn.so.5 - PATHS ${CUDA_TOOLKIT_ROOT_DIR}/targets/aarch64-linux - ${CUDA_TOOLKIT_ROOT_DIR} - PATH_SUFFIXES lib lib64 - ) - - list( APPEND libs ${CUDNN_LIBRARIES} ) - add_definitions( -DCUDNN ) -endif() - if (USE_OPENCV) find_package( OpenCV 2.4.6 REQUIRED )