From 6b39d01c79303f41e7f457a9cb76bafb0251f5b4 Mon Sep 17 00:00:00 2001
From: Boris Fomitchev <borisfom@users.noreply.github.com>
Date: Fri, 27 May 2016 10:32:57 -0700
Subject: [PATCH] Using new CMake FindCUDA elements (#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
---
 lib/THC/CMakeLists.txt | 14 +++++++++++---
 lib/THC/THCHalf.h      |  2 +-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/lib/THC/CMakeLists.txt b/lib/THC/CMakeLists.txt
index 48f5d58c..2425a1c2 100644
--- 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
index d31d418a..b20b05c5 100644
--- 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>