@@ -104,36 +104,44 @@ link_directories(${LLVM_LIBDIR})
104104# TODO: check if this respects CMAKE_PREFIX_PATH
105105find_package (Protobuf REQUIRED)
106106
107- # cuda
108- find_package (CUDA REQUIRED)
109- include_directories (${CUDA_TOOLKIT_ROOT_DIR} /include )
110-
111- # modified CUB
112- find_path (CUB_INCLUDE_DIR NAMES cub)
113- set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DCUDA_HOME=\"\\\" ${CUDA_INCLUDE_DIRS} \\\"\" -DCUB_HOME=\"\\\" ${CUB_INCLUDE_DIR} \\\"\" " )
114-
115- # Inherited from Torch, see
116- # https://github.com/torch/cutorch/blob/master/lib/THC/cmake/select_compute_arch.cmake
117- INCLUDE (cmake/select_compute_arch.cmake)
118- CUDA_SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA)
119-
120- # TODO: Investigate these
121- set (CUDA_VERBOSE_BUILD ON )
122- set (CUDA_PROPAGATE_HOST_FLAGS ON )
123-
124- ################################################################################
125- # FindCUDA doesn't find all the libraries we need, add the extra ones
126- # Cribbed from /lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib
127107################################################################################
128- find_library (CUDA_CUDA_LIBRARIES cuda
129- PATHS ${CUDA_TOOLKIT_ROOT_DIR}
130- PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
131108
132- set (CUDA_CUDA_LIBRARY ${CUDA_CUDA_LIBRARIES} )
133-
134- # Needed to build C2 and some of our tests
135- INCLUDE (cmake/FindCuDNN.cmake)
136- include_directories (${CUDNN_INCLUDE_DIR} )
109+ # cuda
110+ set (WITH_CUDA ON CACHE BOOL "Whether to build with CUDA support" )
111+
112+ if (WITH_CUDA)
113+ find_package (CUDA REQUIRED)
114+ include_directories (${CUDA_TOOLKIT_ROOT_DIR} /include )
115+
116+ # modified CUB
117+ find_path (CUB_INCLUDE_DIR NAMES cub)
118+ set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DWITH_CUDA -DCUDA_HOME=\"\\\" ${CUDA_INCLUDE_DIRS} \\\"\" -DCUB_HOME=\"\\\" ${CUB_INCLUDE_DIR} \\\"\" " )
119+
120+ # Inherited from Torch, see
121+ # https://github.com/torch/cutorch/blob/master/lib/THC/cmake/select_compute_arch.cmake
122+ INCLUDE (cmake/select_compute_arch.cmake)
123+ CUDA_SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA)
124+
125+ # TODO: Investigate these
126+ set (CUDA_VERBOSE_BUILD ON )
127+ set (CUDA_PROPAGATE_HOST_FLAGS ON )
128+
129+ ##############################################################################
130+ # FindCUDA doesn't find all the libraries we need, add the extra ones
131+ # Cribbed from /lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib
132+ ##############################################################################
133+ find_library (CUDA_CUDA_LIBRARIES cuda
134+ PATHS ${CUDA_TOOLKIT_ROOT_DIR}
135+ PATH_SUFFIXES lib lib64 targets/x86_64-linux/lib targets/x86_64-linux/lib/stubs)
136+
137+ set (CUDA_CUDA_LIBRARY ${CUDA_CUDA_LIBRARIES} )
138+
139+ # Needed to build C2 and some of our tests
140+ INCLUDE (cmake/FindCuDNN.cmake)
141+ include_directories (${CUDNN_INCLUDE_DIR} )
142+ else ()
143+ message (STATUS "Building TC without CUDA support" )
144+ endif ()
137145
138146################################################################################
139147
@@ -143,7 +151,7 @@ if(NOT HALIDE_PREFIX)
143151endif ()
144152message (STATUS "HALIDE_PREFIX: ${HALIDE_PREFIX} " )
145153message (STATUS "Finding Halide" )
146- find_path (HALIDE_INCLUDE_DIR NAMES Halide)
154+ find_path (HALIDE_INCLUDE_DIR NAMES Halide.h )
147155message (STATUS "HALIDE_INCLUDE_DIRS: ${HALIDE_INCLUDE_DIR} " )
148156include_directories (${HALIDE_INCLUDE_DIR} )
149157find_library (HALIDE_LIBRARIES NAMES Halide PATHS ${CMAKE_INSTALL_PREFIX} PATH_SUFFIXES lib lib64)
@@ -166,8 +174,10 @@ if (WITH_CAFFE2)
166174 include_directories (${CAFFE2_INCLUDE_DIR} )
167175 find_library (CAFFE2_CPU_LIBRARIES NAMES caffe2 PATHS ${CMAKE_INSTALL_PREFIX} PATH_SUFFIXES lib lib64)
168176 message (STATUS "Found Caffe2_CPU: ${CAFFE2_CPU_LIBRARIES} " )
169- find_library (CAFFE2_GPU_LIBRARIES NAMES caffe2_gpu PATHS ${CMAKE_INSTALL_PREFIX} PATH_SUFFIXES lib lib64)
170- message (STATUS "Found Caffe2_GPU: ${CAFFE2_GPU_LIBRARIES} " )
177+ if (WITH_CUDA)
178+ find_library (CAFFE2_GPU_LIBRARIES NAMES caffe2_gpu PATHS ${CMAKE_INSTALL_PREFIX} PATH_SUFFIXES lib lib64)
179+ message (STATUS "Found Caffe2_GPU: ${CAFFE2_GPU_LIBRARIES} " )
180+ endif ()
171181else ()
172182 message (STATUS "Caffe2 installation is turned off" )
173183endif ()
@@ -236,12 +246,20 @@ message(STATUS "Found glog: ${GLOG_LIBRARIES}")
236246# endforeach()
237247
238248add_subdirectory (src)
249+
250+ # At the moment pybind is only supported in CUDA mode and compilation fails
251+ # for non-CUDA mode (CUDA_HOME and CUB_HOME undefined error).
252+ # Once the core CPU mapper is stabilized we can worry about pybind, deactivate
253+ # conditionally for now
254+ if (WITH_CUDA)
255+ add_subdirectory (tensor_comprehensions/pybinds)
256+ endif ()
257+
239258enable_testing ()
240- add_subdirectory (tensor_comprehensions/pybinds)
241259add_subdirectory (test )
242260
243- if (WITH_CAFFE2)
261+ if (WITH_CAFFE2 AND WITH_CUDA )
244262 add_subdirectory (examples)
245263else ()
246- message (STATUS "Not building examples, caffe2 not available" )
264+ message (STATUS "Not building examples, caffe2 or CUDA not available" )
247265endif ()
0 commit comments