Simplify CMake CUDNN code (#91676)

1. Move CUDNN code to seperate module. 2. Merge CUDNN public and private targets into a single private target. There is no need to expose CUDNN dependency. Pull Request resolved: https://github.com/pytorch/pytorch/pull/91676 Approved by: https://github.com/malfet
2025-10-20 12:54:11 +08:00 · 2023-02-08 01:06:07 +00:00
parent 9291f9b9e2
commit 5fa7120722
8 changed files with 60 additions and 116 deletions
--- a/aten/src/ATen/CMakeLists.txt
+++ b/aten/src/ATen/CMakeLists.txt
@ -32,6 +32,7 @@ set_bool(AT_BLAS_F2C BLAS_F2C)
 set_bool(AT_BLAS_USE_CBLAS_DOT BLAS_USE_CBLAS_DOT)
 set_bool(AT_MAGMA_ENABLED USE_MAGMA)
 set_bool(CAFFE2_STATIC_LINK_CUDA_INT CAFFE2_STATIC_LINK_CUDA)
+set_bool(AT_CUDNN_ENABLED CAFFE2_USE_CUDNN)

 configure_file(Config.h.in "${CMAKE_CURRENT_SOURCE_DIR}/Config.h")
 # TODO: Do not generate CUDAConfig.h for ROCm BUILDS
@ -622,4 +623,4 @@ set(ATen_VULKAN_INCLUDE ${ATen_VULKAN_INCLUDE} PARENT_SCOPE)
 set(ATen_CPU_DEPENDENCY_LIBS ${ATen_CPU_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_CUDA_DEPENDENCY_LIBS ${ATen_CUDA_DEPENDENCY_LIBS} PARENT_SCOPE)
 set(ATen_HIP_DEPENDENCY_LIBS ${ATen_HIP_DEPENDENCY_LIBS} PARENT_SCOPE)
-set(MEM_EFF_ATTENTION_CUDA_SOURCES ${MEM_EFF_ATTENTION_CUDA_SOURCES} PARENT_SCOPE)
+set(MEM_EFF_ATTENTION_CUDA_SOURCES ${MEM_EFF_ATTENTION_CUDA_SOURCES} PARENT_SCOPE)
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@ -1498,10 +1498,6 @@ if(USE_CUDA)
      torch_cuda PRIVATE ${Caffe2_GPU_INCLUDE})
  target_link_libraries(
      torch_cuda PRIVATE ${Caffe2_CUDA_DEPENDENCY_LIBS})
-  if(USE_CUDNN)
-    target_link_libraries(
-        torch_cuda PRIVATE  caffe2::cudnn-private)
-  endif()

  # These public dependencies must go after the previous dependencies, as the
  # order of the libraries in the linker call matters here when statically
--- a/cmake/Caffe2Config.cmake.in
+++ b/cmake/Caffe2Config.cmake.in
@ -84,7 +84,6 @@ if(@USE_CUDA@)
  # If Caffe2 was compiled with the libraries below, they must
  # be found again when including the Caffe2 target.
  set(CAFFE2_USE_CUDA @USE_CUDA@)
-  set(CAFFE2_USE_CUDNN @USE_CUDNN@)
  set(CAFFE2_USE_TENSORRT @USE_TENSORRT@)
  include("${CMAKE_CURRENT_LIST_DIR}/public/cuda.cmake")
  if(@CAFFE2_USE_CUDA@ AND NOT CAFFE2_USE_CUDA)
@ -93,12 +92,6 @@ if(@USE_CUDA@)
      "libraries. Please set the proper CUDA prefixes and / or install "
      "CUDA.")
  endif()
-  if(@CAFFE2_USE_CUDNN@ AND NOT CAFFE2_USE_CUDNN)
-    message(FATAL_ERROR
-      "Your installed Caffe2 version uses cuDNN but I cannot find the cuDNN "
-      "libraries. Please set the proper cuDNN prefixes and / or install "
-      "cuDNN.")
-  endif()
  if(@CAFFE2_USE_TENSORRT@ AND NOT CAFFE2_USE_TENSORRT)
    message(FATAL_ERROR
      "Your installed Caffe2 version uses TensorRT but I cannot find the TensorRT "
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@ -53,7 +53,7 @@ if(USE_CUDA)
      caffe2_update_option(USE_NVRTC OFF)
    endif()
    if(CAFFE2_USE_CUDNN)
-      list(APPEND Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS caffe2::cudnn-public)
+      list(APPEND Caffe2_CUDA_DEPENDENCY_LIBS torch::cudnn)
    else()
      caffe2_update_option(USE_CUDNN OFF)
    endif()
@ -1236,7 +1236,7 @@ endif(USE_LLVM)
 # ---[ cuDNN
 if(USE_CUDNN)
  set(CUDNN_FRONTEND_INCLUDE_DIR ${CMAKE_CURRENT_LIST_DIR}/../third_party/cudnn_frontend/include)
-  include_directories(${CUDNN_FRONTEND_INCLUDE_DIR})
+  target_include_directories(torch::cudnn INTERFACE ${CUDNN_FRONTEND_INCLUDE_DIR})
 endif()

 # ---[ HIP
@ -1730,17 +1730,6 @@ if(NOT INTERN_BUILD_MOBILE)
    set(AT_CUDA_ENABLED 1)
  endif()

-  if(NOT USE_CUDNN)
-    message(STATUS "USE_CUDNN is set to 0. Compiling without cuDNN support")
-    set(AT_CUDNN_ENABLED 0)
-  elseif(NOT CUDNN_FOUND)
-    message(WARNING "CuDNN not found. Compiling without CuDNN support")
-    set(AT_CUDNN_ENABLED 0)
-  else()
-    include_directories(SYSTEM ${CUDNN_INCLUDE_PATH})
-    set(AT_CUDNN_ENABLED 1)
-  endif()
-
  if(NOT USE_ROCM)
    message("disabling ROCM because NOT USE_ROCM is set")
    message(STATUS "MIOpen not found. Compiling without MIOpen support")
--- a/cmake/Modules_CUDA_fix/FindCUDNN.cmake
+++ b/cmake/Modules_CUDA_fix/FindCUDNN.cmake
@ -47,4 +47,32 @@ find_library(CUDNN_LIBRARY_PATH ${CUDNN_LIBNAME}

 find_package_handle_standard_args(CUDNN DEFAULT_MSG CUDNN_LIBRARY_PATH CUDNN_INCLUDE_PATH)

-mark_as_advanced(CUDNN_ROOT CUDNN_INCLUDE_DIR CUDNN_LIBRARY)
+if(CUDNN_FOUND)
+  # Get cuDNN version
+  if(EXISTS ${CUDNN_INCLUDE_PATH}/cudnn_version.h)
+    file(READ ${CUDNN_INCLUDE_PATH}/cudnn_version.h CUDNN_HEADER_CONTENTS)
+  else()
+    file(READ ${CUDNN_INCLUDE_PATH}/cudnn.h CUDNN_HEADER_CONTENTS)
+  endif()
+  string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"
+               CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}")
+  string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1"
+               CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}")
+  string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)"
+               CUDNN_VERSION_MINOR "${CUDNN_HEADER_CONTENTS}")
+  string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1"
+               CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}")
+  string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)"
+               CUDNN_VERSION_PATCH "${CUDNN_HEADER_CONTENTS}")
+  string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1"
+               CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}")
+  # Assemble cuDNN version
+  if(NOT CUDNN_VERSION_MAJOR)
+    set(CUDNN_VERSION "?")
+  else()
+    set(CUDNN_VERSION
+        "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}")
+  endif()
+endif()
+
+mark_as_advanced(CUDNN_ROOT CUDNN_INCLUDE_DIR CUDNN_LIBRARY CUDNN_VERSION)
--- a/cmake/Summary.cmake
+++ b/cmake/Summary.cmake
@ -95,12 +95,8 @@ function(caffe2_print_configuration_summary)
    get_target_property(__tmp caffe2::curand IMPORTED_LOCATION)
    message(STATUS "    curand library      : ${__tmp}")
    if(${USE_CUDNN})
-      get_target_property(__tmp caffe2::cudnn-public INTERFACE_LINK_LIBRARIES)
+      get_target_property(__tmp torch::cudnn INTERFACE_LINK_LIBRARIES)
      message(STATUS "    cuDNN library       : ${__tmp}")
-      if(${CUDNN_STATIC})
-        get_target_property(__tmp caffe2::cudnn-private INTERFACE_LINK_LIBRARIES)
-        message(STATUS "    cuDNN static library: ${__tmp}")
-      endif()
    endif()
    get_target_property(__tmp caffe2::nvrtc IMPORTED_LOCATION)
    message(STATUS "    nvrtc               : ${__tmp}")
--- a/cmake/public/cuda.cmake
+++ b/cmake/public/cuda.cmake
@ -107,21 +107,6 @@ if(CUDA_FOUND)
  endif()
 endif()

-# Find cuDNN.
-if(USE_STATIC_CUDNN)
-  set(CUDNN_STATIC ON CACHE BOOL "")
-else()
-  set(CUDNN_STATIC OFF CACHE BOOL "")
-endif()
-
-find_package(CUDNN)
-
-if(CAFFE2_USE_CUDNN AND NOT CUDNN_FOUND)
-  message(WARNING
-    "Caffe2: Cannot find cuDNN library. Turning the option off")
-  set(CAFFE2_USE_CUDNN OFF)
-endif()
-
 # Optionally, find TensorRT
 if(CAFFE2_USE_TENSORRT)
  find_path(TENSORRT_INCLUDE_DIR NvInfer.h
@ -153,39 +138,6 @@ if(CAFFE2_USE_TENSORRT)
  endif()
 endif()

-# ---[ Extract versions
-if(CAFFE2_USE_CUDNN)
-  # Get cuDNN version
-  if(EXISTS ${CUDNN_INCLUDE_PATH}/cudnn_version.h)
-    file(READ ${CUDNN_INCLUDE_PATH}/cudnn_version.h CUDNN_HEADER_CONTENTS)
-  else()
-    file(READ ${CUDNN_INCLUDE_PATH}/cudnn.h CUDNN_HEADER_CONTENTS)
-  endif()
-  string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)"
-               CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}")
-  string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1"
-               CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}")
-  string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)"
-               CUDNN_VERSION_MINOR "${CUDNN_HEADER_CONTENTS}")
-  string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1"
-               CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}")
-  string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)"
-               CUDNN_VERSION_PATCH "${CUDNN_HEADER_CONTENTS}")
-  string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1"
-               CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}")
-  # Assemble cuDNN version
-  if(NOT CUDNN_VERSION_MAJOR)
-    set(CUDNN_VERSION "?")
-  else()
-    set(CUDNN_VERSION
-        "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}")
-  endif()
-  message(STATUS "Found cuDNN: v${CUDNN_VERSION}  (include: ${CUDNN_INCLUDE_PATH}, library: ${CUDNN_LIBRARY_PATH})")
-  if(CUDNN_VERSION VERSION_LESS "7.0.0")
-    message(FATAL_ERROR "PyTorch requires cuDNN 7 and above.")
-  endif()
-endif()
-
 # ---[ CUDA libraries wrapper

 # find libcuda.so and lbnvrtc.so
@ -305,49 +257,37 @@ set_property(
    TARGET caffe2::cublas PROPERTY INTERFACE_INCLUDE_DIRECTORIES
    ${CUDA_INCLUDE_DIRS})

-# cudnn public and private interfaces
+# cudnn interface
 # static linking is handled by USE_STATIC_CUDNN environment variable
-# If library is linked dynamically, than private interface is no-op
-# If library is linked statically:
-#  - public interface would only reference headers
-#  - private interface will contain the actual link instructions
 if(CAFFE2_USE_CUDNN)
-  add_library(caffe2::cudnn-public INTERFACE IMPORTED)
-  set_property(
-    TARGET caffe2::cudnn-public PROPERTY INTERFACE_INCLUDE_DIRECTORIES
-    ${CUDNN_INCLUDE_PATH})
-  add_library(caffe2::cudnn-private INTERFACE IMPORTED)
-  set_property(
-    TARGET caffe2::cudnn-private PROPERTY INTERFACE_INCLUDE_DIRECTORIES
-    ${CUDNN_INCLUDE_PATH})
+  if(USE_STATIC_CUDNN)
+    set(CUDNN_STATIC ON CACHE BOOL "")
+  else()
+    set(CUDNN_STATIC OFF CACHE BOOL "")
+  endif()
+
+  find_package(CUDNN)
+
+  if(NOT CUDNN_FOUND)
+    message(WARNING
+      "Cannot find cuDNN library. Turning the option off")
+    set(CAFFE2_USE_CUDNN OFF)
+  else()
+    if(CUDNN_VERSION VERSION_LESS "8.0.0")
+      message(FATAL_ERROR "PyTorch requires cuDNN 8 and above.")
+    endif()
+  endif()
+
+  add_library(torch::cudnn INTERFACE IMPORTED)
+  target_include_directories(torch::cudnn INTERFACE ${CUDNN_INCLUDE_PATH})
  if(CUDNN_STATIC AND NOT WIN32)
-    set_property(
-      TARGET caffe2::cudnn-private PROPERTY INTERFACE_LINK_LIBRARIES
-      ${CUDNN_LIBRARY_PATH})
-    set_property(
-      TARGET caffe2::cudnn-private APPEND PROPERTY INTERFACE_LINK_LIBRARIES
-      "${CUDA_TOOLKIT_ROOT_DIR}/lib64/libculibos.a" dl)
-    # Add explicit dependency on cublas to cudnn
-    get_target_property(__tmp caffe2::cublas INTERFACE_LINK_LIBRARIES)
-    set_property(
-      TARGET caffe2::cudnn-private APPEND PROPERTY INTERFACE_LINK_LIBRARIES
-      "${__tmp}")
-    # Lines below use target_link_libraries because we support cmake 3.5+.
-    # For cmake 3.13+, target_link_options to set INTERFACE_LINK_OPTIONS would be better.
-    # https://cmake.org/cmake/help/v3.5/command/target_link_libraries.html warns
-    # "Item names starting with -, but not -l or -framework, are treated as linker flags.
-    #  Note that such flags will be treated like any other library link item for purposes
-    #  of transitive dependencies, so they are generally safe to specify only as private
-    #  link items that will not propagate to dependents."
-    # Propagating to a dependent (torch_cuda) is exactly what we want here, so we are
-    # flouting the warning, but I can't think of a better (3.5+ compatible) way.
-    target_link_libraries(caffe2::cudnn-private INTERFACE
+    target_link_options(torch::cudnn INTERFACE
        "-Wl,--exclude-libs,libcudnn_static.a")
  else()
-  set_property(
-    TARGET caffe2::cudnn-public PROPERTY INTERFACE_LINK_LIBRARIES
-    ${CUDNN_LIBRARY_PATH})
+    target_link_libraries(torch::cudnn INTERFACE ${CUDNN_LIBRARY_PATH})
  endif()
+else()
+  message(STATUS "USE_CUDNN is set to 0. Compiling without cuDNN support")
 endif()

 # curand
--- a/torch/CMakeLists.txt
+++ b/torch/CMakeLists.txt
@ -132,6 +132,7 @@ if(USE_CUDA)

    list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDA)
    if(USE_CUDNN)
+        list(APPEND TORCH_PYTHON_LINK_LIBRARIES torch::cudnn)
        list(APPEND TORCH_PYTHON_COMPILE_DEFINITIONS USE_CUDNN)
    endif()