-
Notifications
You must be signed in to change notification settings - Fork 3.4k
/
FindCUDA.cmake
144 lines (140 loc) · 5.44 KB
/
FindCUDA.cmake
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#######################################################
# Enhanced version of find CUDA.
#
# Usage:
# find_cuda(${USE_CUDA} ${USE_CUDNN})
#
# - When USE_CUDA=ON, use auto search
# - When USE_CUDA=/path/to/cuda-path, use the cuda path
# - When USE_CUDNN=ON, use auto search
# - When USE_CUDNN=/path/to/cudnn-path, use the cudnn path
#
# Provide variables:
#
# - CUDA_FOUND
# - CUDA_INCLUDE_DIRS
# - CUDA_TOOLKIT_ROOT_DIR
# - CUDA_CUDA_LIBRARY
# - CUDA_CUDART_LIBRARY
# - CUDA_NVRTC_LIBRARY
# - CUDA_CUDNN_INCLUDE_DIRS
# - CUDA_CUDNN_LIBRARY
# - CUDA_CUBLAS_LIBRARY
#
macro(find_cuda use_cuda use_cudnn)
set(__use_cuda ${use_cuda})
if(${__use_cuda} MATCHES ${IS_TRUE_PATTERN})
find_package(CUDA QUIET)
elseif(IS_DIRECTORY ${__use_cuda})
set(CUDA_TOOLKIT_ROOT_DIR ${__use_cuda})
message(STATUS "Custom CUDA_PATH=" ${CUDA_TOOLKIT_ROOT_DIR})
set(CUDA_INCLUDE_DIRS ${CUDA_TOOLKIT_ROOT_DIR}/include)
set(CUDA_FOUND TRUE)
if(MSVC)
find_library(CUDA_CUDART_LIBRARY cudart
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
else(MSVC)
find_library(CUDA_CUDART_LIBRARY cudart
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDA_TOOLKIT_ROOT_DIR}/lib)
endif(MSVC)
endif()
# additional libraries
if(CUDA_FOUND)
if(MSVC)
find_library(CUDA_CUDA_LIBRARY cuda
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
find_library(CUDA_NVRTC_LIBRARY nvrtc
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
find_library(CUDA_CUBLAS_LIBRARY cublas
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
find_library(CUDA_CUBLASLT_LIBRARY cublaslt
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
else(MSVC)
find_library(_CUDA_CUDA_LIBRARY cuda
PATHS ${CUDA_TOOLKIT_ROOT_DIR}
PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs lib64/stubs
NO_DEFAULT_PATH)
if(_CUDA_CUDA_LIBRARY)
set(CUDA_CUDA_LIBRARY ${_CUDA_CUDA_LIBRARY})
endif()
find_library(CUDA_NVRTC_LIBRARY nvrtc
PATHS ${CUDA_TOOLKIT_ROOT_DIR}
PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs lib64/stubs lib/x86_64-linux-gnu
NO_DEFAULT_PATH)
find_library(CUDA_CURAND_LIBRARY curand
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDA_TOOLKIT_ROOT_DIR}/lib
NO_DEFAULT_PATH)
find_library(CUDA_CUBLAS_LIBRARY cublas
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDA_TOOLKIT_ROOT_DIR}/lib
NO_DEFAULT_PATH)
# search default path if cannot find cublas in non-default
find_library(CUDA_CUBLAS_LIBRARY cublas)
find_library(CUDA_CUBLASLT_LIBRARY
NAMES cublaslt cublasLt
PATHS
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDA_TOOLKIT_ROOT_DIR}/lib
NO_DEFAULT_PATH)
# search default path if cannot find cublaslt in non-default
find_library(CUDA_CUBLASLT_LIBRARY NAMES cublaslt cublasLt)
endif(MSVC)
# find cuDNN
set(__use_cudnn ${use_cudnn})
if(${__use_cudnn} MATCHES ${IS_TRUE_PATTERN})
set(CUDA_CUDNN_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS})
if(MSVC)
find_library(CUDA_CUDNN_LIBRARY cudnn
${CUDA_TOOLKIT_ROOT_DIR}/lib/x64
${CUDA_TOOLKIT_ROOT_DIR}/lib/Win32)
else(MSVC)
find_library(CUDA_CUDNN_LIBRARY cudnn
${CUDA_TOOLKIT_ROOT_DIR}/lib64
${CUDA_TOOLKIT_ROOT_DIR}/lib
NO_DEFAULT_PATH)
# search default path if cannot find cudnn in non-default
find_library(CUDA_CUDNN_LIBRARY cudnn)
endif(MSVC)
elseif(IS_DIRECTORY ${__use_cudnn})
# cuDNN doesn't necessarily live in the CUDA dir
set(CUDA_CUDNN_ROOT_DIR ${__use_cudnn})
set(CUDA_CUDNN_INCLUDE_DIRS ${CUDA_CUDNN_ROOT_DIR}/include)
find_library(CUDA_CUDNN_LIBRARY cudnn
${CUDA_CUDNN_ROOT_DIR}/lib64
${CUDA_CUDNN_ROOT_DIR}/lib
NO_DEFAULT_PATH)
endif()
message(STATUS "Found CUDA_TOOLKIT_ROOT_DIR=" ${CUDA_TOOLKIT_ROOT_DIR})
message(STATUS "Found CUDA_CUDA_LIBRARY=" ${CUDA_CUDA_LIBRARY})
message(STATUS "Found CUDA_CUDART_LIBRARY=" ${CUDA_CUDART_LIBRARY})
message(STATUS "Found CUDA_NVRTC_LIBRARY=" ${CUDA_NVRTC_LIBRARY})
message(STATUS "Found CUDA_CUDNN_INCLUDE_DIRS=" ${CUDA_CUDNN_INCLUDE_DIRS})
message(STATUS "Found CUDA_CUDNN_LIBRARY=" ${CUDA_CUDNN_LIBRARY})
message(STATUS "Found CUDA_CUBLAS_LIBRARY=" ${CUDA_CUBLAS_LIBRARY})
message(STATUS "Found CUDA_CURAND_LIBRARY=" ${CUDA_CURAND_LIBRARY})
message(STATUS "Found CUDA_CUBLASLT_LIBRARY=" ${CUDA_CUBLASLT_LIBRARY})
endif(CUDA_FOUND)
endmacro(find_cuda)