Turn on NVTX by default in Unix.

apache · Apr 19, 2019 · 379e2c8 · 379e2c8
1 parent da2f0c9
commit 379e2c8
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -285,13 +285,15 @@ if(USE_CUDA)
       message(WARNING "Could not find NCCL libraries")
     endif()
   endif()
-  find_package(NVTX)
-  if(NVTX_FOUND)
-    include_directories(${NVTX_INCLUDE_DIRS})
-    list(APPEND mxnet_LINKER_LIBS ${NVTX_LIBRARIES})
-    add_definitions(-DMXNET_USE_NVTX=1)
-  else()
-    message(WARNING "Could not find NVTX libraries")
+  if(UNIX)
+    find_package(NVTX)
+    if(NVTX_FOUND)
+      include_directories(${NVTX_INCLUDE_DIRS})
+      list(APPEND mxnet_LINKER_LIBS ${NVTX_LIBRARIES})
+      add_definitions(-DMXNET_USE_NVTX=1)
+    else()
+      message(WARNING "Could not find NVTX libraries")
+    endif()
   endif()
 else()
   add_definitions(-DMSHADOW_USE_CUDA=0)

diff --git a/cmake/Modules/FindNVTX.cmake b/cmake/Modules/FindNVTX.cmake
@@ -24,7 +24,7 @@ find_path(NVTX_INCLUDE_DIRS
   )
 
 find_library(NVTX_LIBRARIES
-  NAMES nvToolsExt64_1 nvToolsExt32_1 nvToolsExt
+  NAMES nvToolsExt64_1.lib nvToolsExt32_1.lib nvToolsExt
   PATHS $ENV{NVTOOLSEXT_PATH} ${NVTX_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}
   PATH_SUFFIXES lib lib64 lib/Win32 lib/x64
   )