From 4bdf038d050dfd3bc7ae7cbf4ae07e8cc4982357 Mon Sep 17 00:00:00 2001
From: Peihong Liu <mosout@qq.com>
Date: Thu, 19 Dec 2024 15:21:47 +0800
Subject: [PATCH] for compiling with cuda 12.x (#10568)

---
 cmake/cuda.cmake                  | 10 ++++++++++
 cmake/third_party/FindCUDNN.cmake |  9 ++++++++-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/cmake/cuda.cmake b/cmake/cuda.cmake
index 578786351a0..63cdb68fb2a 100644
--- a/cmake/cuda.cmake
+++ b/cmake/cuda.cmake
@@ -89,6 +89,16 @@ if(BUILD_CUDA)
       # GeForce RTX 30xx
       list(APPEND CMAKE_CUDA_ARCHITECTURES 86-real)
     endif()
+
+    if(CUDA_VERSION VERSION_GREATER_EQUAL "11.8")
+      # GeForce RTX 40xx
+      list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real)
+    endif()
+
+    if(CUDA_VERSION VERSION_GREATER_EQUAL "12.0")
+      # H100, H20
+      list(APPEND CMAKE_CUDA_ARCHITECTURES 90-real)
+    endif()
   endif()
 
   foreach(CUDA_ARCH ${CMAKE_CUDA_ARCHITECTURES})
diff --git a/cmake/third_party/FindCUDNN.cmake b/cmake/third_party/FindCUDNN.cmake
index 3e076419a3a..716898e6dd7 100644
--- a/cmake/third_party/FindCUDNN.cmake
+++ b/cmake/third_party/FindCUDNN.cmake
@@ -66,7 +66,14 @@ if(CUDNN_FOUND)
 
   set(CUDNN_INCLUDE_DIRS ${CUDNN_INCLUDE_DIR})
 
-  if(NOT CUDNN_STATIC AND CUDNN_VERSION_MAJOR GREATER_EQUAL 8)
+  if(NOT CUDNN_STATIC AND CUDNN_VERSION_MAJOR GREATER_EQUAL 9)
+    # skipping: libcudnn_adv_infer.so libcudnn_adv_train.so
+    set(CUDNN_DYNAMIC_NAMES libcudnn_cnn.so libcudnn_ops.so)
+    get_filename_component(CUDNN_LIBRARY_DIRECTORY ${CUDNN_LIBRARY} DIRECTORY)
+    foreach(CUDNN_DYNAMIC_NAME ${CUDNN_DYNAMIC_NAMES})
+      list(APPEND CUDNN_LIBRARIES ${CUDNN_LIBRARY_DIRECTORY}/${CUDNN_DYNAMIC_NAME})
+    endforeach()
+  elseif(NOT CUDNN_STATIC AND CUDNN_VERSION_MAJOR GREATER_EQUAL 8)
     # skipping: libcudnn_adv_infer.so libcudnn_adv_train.so
     set(CUDNN_DYNAMIC_NAMES libcudnn_cnn_infer.so libcudnn_cnn_train.so libcudnn_ops_infer.so
                             libcudnn_ops_train.so)