Using new CMake FindCUDA elements (torch#415)

* Use FindCUDA module * Removed extra FindCUDA wrappers * adjusted for new CUDA_ macro name * Added fp16 recognition for mobile CUDA / new CMake * Porting to new CMake
marsbzp · May 27, 2016 · 6b39d01 · 6b39d01
1 parent 369bc80
commit 6b39d01
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/lib/THC/CMakeLists.txt b/lib/THC/CMakeLists.txt
@@ -57,8 +57,14 @@ IF(APPLE)
 ENDIF(APPLE)
 
 # Detect CUDA architecture and get best NVCC flags
-INCLUDE(${CMAKE_CURRENT_SOURCE_DIR}/../../cmake/FindCudaArch.cmake)
-SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA)
+IF(DEFINED CUDA_HAS_FP16)
+# New CMake macro, comes with FindCUDA
+  CUDA_SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA $ENV{TORCH_CUDA_ARCH_LIST})
+ELSE()
+# Old way of getting the flags, with extra include file
+  INCLUDE(${CMAKE_CURRENT_SOURCE_DIR}/../../cmake/FindCudaArch.cmake)
+  SELECT_NVCC_ARCH_FLAGS(NVCC_FLAGS_EXTRA)
+ENDIF()
 LIST(APPEND CUDA_NVCC_FLAGS ${NVCC_FLAGS_EXTRA})
 
 IF(NOT THC_INSTALL_BIN_SUBDIR
@@ -114,8 +120,10 @@ SET(src-cuda
   THCTensorTopK.cu
   )
 
-IF(NOT ${CUDA_VERSION} LESS 7.5)
+IF(${CUDA_HAS_FP16} OR NOT ${CUDA_VERSION} LESS 7.5)
   LIST(APPEND src-cuda THCHalf.cu)
+  LIST(APPEND CUDA_NVCC_FLAGS "-DCUDA_HAS_FP16")
+  SET(CMAKE_C_FLAGS "-DCUDA_HAS_FP16 ${CMAKE_C_FLAGS}")
 ENDIF()
 
 CUDA_ADD_LIBRARY(THC SHARED ${src} ${src-cuda})

diff --git a/lib/THC/THCHalf.h b/lib/THC/THCHalf.h
@@ -3,7 +3,7 @@
 
 #include "THCGeneral.h"
 
-#if CUDA_VERSION >= 7050
+#if CUDA_VERSION >= 7050 || CUDA_HAS_FP16
 
 #include <cuda_fp16.h>
 #include <stdint.h>