Unify torch/csrc/cuda/shared/cudnn.cpp include path (#40525)

malfet · facebook-github-bot · commit c120fdc05b4a · 2020-06-24T16:40:11.000-07:00
Summary: Pull Request resolved: #40525 Move `USE_CUDNN` define under `USE_CUDA` guard, add `cuda/shared/cudnn.cpp` to filelist if either USE_ROCM or USE_CUDNN is set. This is a prep change for PyTorch CUDA src filelist unification change. Test Plan: CI Differential Revision: D22214899 fbshipit-source-id: b71b32fc603783b41cdef0e7fab2cc9cbe750a4e
diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt
@@ -124,6 +124,9 @@ if(USE_CUDA)
       ${GENERATED_THNN_CXX_CUDA}
       )
     list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDA)
+    if(USE_CUDNN)
+        list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDNN)
+    endif()
 
     if(MSVC)
       list(APPEND TORCH_PYTHON_LINK_LIBRARIES ${NVTOOLEXT_HOME}/lib/x64/nvToolsExt64_1.lib)
@@ -137,18 +140,6 @@ if(USE_CUDA)
 
 endif()
 
-if(USE_CUDNN)
-    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDNN)
-
-    list(APPEND TORCH_PYTHON_SRCS
-      ${TORCH_SRC_DIR}/csrc/cuda/shared/cudnn.cpp
-      )
-endif()
-
-if(USE_NUMPY)
-    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_NUMPY)
-endif()
-
 if(USE_ROCM)
     list(APPEND TORCH_PYTHON_SRCS
       ${TORCH_SRC_DIR}/csrc/cuda/Module.cpp
@@ -159,7 +150,6 @@ if(USE_ROCM)
       ${TORCH_SRC_DIR}/csrc/cuda/python_comm.cpp
       ${TORCH_SRC_DIR}/csrc/cuda/serialization.cpp
       ${TORCH_SRC_DIR}/csrc/cuda/shared/cudart.cpp
-      ${TORCH_SRC_DIR}/csrc/cuda/shared/cudnn.cpp
       ${TORCH_SRC_DIR}/csrc/cuda/shared/nvtx.cpp
       ${GENERATED_THNN_CXX_CUDA}
       )
@@ -172,6 +162,16 @@ if(USE_ROCM)
     list(APPEND TORCH_PYTHON_INCLUDE_DIRECTORIES ${roctracer_INCLUDE_DIRS})
 endif()
 
+if(USE_CUDNN OR USE_ROCM)
+    list(APPEND TORCH_PYTHON_SRCS
+      ${TORCH_SRC_DIR}/csrc/cuda/shared/cudnn.cpp
+      )
+endif()
+
+if(USE_NUMPY)
+    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_NUMPY)
+endif()
+
 if(USE_DISTRIBUTED)
     list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_DISTRIBUTED)
     if(NOT MSVC)
