# Try to find the NVML library and headers Usage of this module is as follows
#
# find_package( NVML ) if(NVML_FOUND) include_directories(${NVML_INCLUDE_DIRS})
# add_executable(foo foo.cc) target_link_libraries(foo ${NVML_LIBRARIES}) endif()
#
# ===========================================================================
# Variables used by this module which can be used to change the default behaviour, and
# hence need to be set before calling find_package:
#
# NVML_ROOT_DIR The preferred installation prefix for searching for NVML Set this if the
# module has problems finding the proper NVML installation.
#
# If you don't supply NVML_ROOT_DIR, the module will search on the standard system paths.
#
# ============================================================================
# Variables set by this module:
#
# NVML_FOUND           System has NVML.
#
# NVML_INCLUDE_DIRS    NVML include directories: not cached.
#
# NVML_LIBRARIES       Link to these to use the NVML library: not cached.
#
# ===========================================================================
# If NVML is installed in a non-standard way, e.g. a non GNU-style install of
# <prefix>/{lib,include}, then this module may fail to locate the headers and libraries as
# needed. In this case, the following cached variables can be editted to point to the
# correct locations.
#
# NVML_INCLUDE_DIR    The path to the NVML include directory: cached
#
# NVML_LIBRARY        The path to the NVML library: cached
#
# You should not need to set these in the vast majority of cases
#

# ----------------------------------------------------------------------------------------#

find_path(
    NVML_ROOT_DIR
    NAMES include/nvml.h
    HINTS ENV NVML_ROOT_DIR ${CUDA_TOOLKIT_ROOT_DIR} ${CUDA_SDK_ROOT_DIR} /usr/local/cuda
    PATHS ENV NVML_ROOT_DIR ${CUDA_TOOLKIT_ROOT_DIR} ${CUDA_SDK_ROOT_DIR} /usr/local/cuda
    PATH_SUFFIXES targets/x86_64-linux targets/x86_64-darwin targets/x86_64-win32
                  targets/x86_64-win64
    DOC "NVML root installation directory")

# ----------------------------------------------------------------------------------------#

find_path(
    NVML_INCLUDE_DIR
    NAMES nvml.h
    HINTS ${NVML_ROOT_DIR} ENV NVML_ROOT_DIR ENV CPATH
    PATH_SUFFIXES include
    DOC "Path to the NVML headers")

# ----------------------------------------------------------------------------------------#

set(_NVML_PATHS ${NVML_ROOT_DIR} $ENV{NVML_ROOT_DIR} $ENV{CUDA_HOME}
                ${CUDA_TOOLKIT_ROOT_DIR} ${CUDA_SDK_ROOT_DIR})

find_library(
    NVML_LIBRARY
    NAMES nvidia-ml
    PATHS ${_NVML_PATHS}
    HINTS ${_NVML_PATHS}
    PATH_SUFFIXES lib lib64
    DOC "Path to the NVML library")

# ----------------------------------------------------------------------------------------#

# try to find cuda driver stubs library if no real driver
if(NOT NVML_LIBRARY)
    find_library(
        NVML_LIBRARY
        NAMES nvidia-ml
        PATHS ${_NVML_PATHS}
        HINTS ${_NVML_PATHS}
        PATH_SUFFIXES lib/stubs lib64/stubs stubs)
    set(NVML_LIBRARY_HAS_SYMBOLS
        OFF
        CACHE BOOL "Using stubs library" FORCE)
else()
    set(NVML_LIBRARY_HAS_SYMBOLS
        ON
        CACHE BOOL "Using stubs library" FORCE)
endif()

# ----------------------------------------------------------------------------------------#

if(NVML_ROOT_DIR)
    set(_NVML_ROOT_DIR NVML_ROOT_DIR)
endif()

include(FindPackageHandleStandardArgs)
# handle the QUIETLY and REQUIRED arguments and set NVML_FOUND to TRUE if all listed
# variables are TRUE
find_package_handle_standard_args(NVML DEFAULT_MSG ${_NVML_ROOT_DIR} NVML_INCLUDE_DIR
                                  NVML_LIBRARY)

# ----------------------------------------------------------------------------------------#

if(NVML_FOUND)
    if(NOT TARGET nvml)
        add_library(nvml INTERFACE)
        add_library(NVML::nvml ALIAS nvml)
    endif()
    target_link_libraries(nvml INTERFACE ${NVML_LIBRARY})
    target_include_directories(nvml INTERFACE ${NVML_INCLUDE_DIR})
    get_filename_component(NVML_INCLUDE_DIRS ${NVML_INCLUDE_DIR} REALPATH)
    get_filename_component(NVML_LIBRARIES ${NVML_LIBRARY} REALPATH)
endif()

mark_as_advanced(NVML_INCLUDE_DIR NVML_LIBRARY)

unset(_NVML_ROOT_DIR)