for compiling with cuda 12.x (#10568)

Oneflow-Inc · Dec 19, 2024 · 4bdf038 · 4bdf038
1 parent dd72bc9
commit 4bdf038
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git 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
@@ -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)