From 0c6cf94e2b77f0a06819337ed0ebe4b222848e87 Mon Sep 17 00:00:00 2001
From: Maxime Schmitt <max.schmitt@unistra.fr>
Date: Thu, 5 Apr 2018 18:42:44 +0200
Subject: [PATCH] Update CMake search path for NVML lib and include

---
 CMakeLists.txt               |  4 ++++
 cmake/modules/FindNVML.cmake | 37 ++++++++++++++++++------------------
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 5b84576..e27587b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -21,6 +21,10 @@ endif()
 #                           DEPENDENCIES                            #
 #///////////////////////////////////////////////////////////////////#
 
+if(NVML_RETRIEVE_HEADER_ONLINE)
+  file(DOWNLOAD "https://raw.githubusercontent.com/NVIDIA/nvidia-settings/master/src/nvml.h" ${PROJECT_BINARY_DIR}/include/nvml.h)
+endif()
+
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules)
 
 find_package(NVML REQUIRED)
diff --git a/cmake/modules/FindNVML.cmake b/cmake/modules/FindNVML.cmake
index 7d5b52d..c1763f7 100644
--- a/cmake/modules/FindNVML.cmake
+++ b/cmake/modules/FindNVML.cmake
@@ -1,4 +1,4 @@
-# Version 1.1
+# Version 1.2
 # Public Domain
 # Written by Maxime SCHMITT <maxime.schmitt@etu.unistra.fr>
 
@@ -19,29 +19,30 @@ if (NVML_INCLUDE_DIRS AND NVML_LIBRARIES)
 endif()
 
 # Headers
-file(GLOB nvml_header_path_hint /usr/local/cuda*/include /opt/cuda*/include)
+file(GLOB nvml_header_path_hint /usr/include/nvidia*/include /usr/local/cuda*/include /opt/cuda*/include)
 find_path(NVML_INCLUDE_DIRS NAMES nvml.h
-  HINTS ${nvml_header_path_hint})
+  PATHS ${nvml_header_path_hint} ${PROJECT_BINARY_DIR}/include)
 
 # library
-find_library(NVML_LIBRARIES NAMES nvml nvidia-ml)
+file(GLOB nvml_lib_path_hint /usr/lib*/nvidia*/)
+find_library(NVML_LIBRARIES NAMES nvidia-ml libnvidia-ml.so.1
+  PATHS ${nvml_lib_path_hint})
 
 # Version
 set(filename "${NVML_INCLUDE_DIRS}/nvml.h")
-if (NOT EXISTS ${filename} AND NOT quiet)
-  message(AUTHOR_WARNING "Unable to find ${filename}")
-endif()
-file(READ "${filename}" nvml_header)
-set(nvml_api_version_match "NVML_API_VERSION")
-
-string(REGEX REPLACE ".*#[ \t]*define[ \t]*${nvml_api_version_match}[ \t]*([0-9]+).*"
-  "\\1" nvml_api_version "${nvml_header}")
-
-if (nvml_api_version STREQUAL nvml_header)
-  message(AUTHOR_WARNING "Unable to find nvml api version")
-else()
-  set(NVML_API_VERSION   "${nvml_api_version}")
-endif()
+if (EXISTS ${filename})
+  file(READ "${filename}" nvml_header)
+  set(nvml_api_version_match "NVML_API_VERSION")
+
+  string(REGEX REPLACE ".*#[ \t]*define[ \t]*${nvml_api_version_match}[ \t]*([0-9]+).*"
+    "\\1" nvml_api_version "${nvml_header}")
+
+  if (nvml_api_version STREQUAL nvml_header AND NOT quiet)
+    message(AUTHOR_WARNING "Unable to find nvml api version")
+  else()
+    set(NVML_API_VERSION   "${nvml_api_version}")
+  endif()
+endif(EXISTS ${filename})
 
 include(FindPackageHandleStandardArgs)
 FIND_PACKAGE_HANDLE_STANDARD_ARGS(NVML