Fix cudnn.make include path (#73005)

co63oc · web-flow · commit 4c802cc579f7 · 2025-06-03T17:13:35.000+08:00
* Fix cudnn.make

* Fix
diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake
@@ -10,20 +10,24 @@ else()
       CACHE PATH "CUDNN ROOT")
 endif()
 
+set(TARGET_ARCH "x86_64")
+if(NOT ${CMAKE_SYSTEM_PROCESSOR})
+  set(TARGET_ARCH ${CMAKE_SYSTEM_PROCESSOR})
+endif()
+
 find_path(
   CUDNN_INCLUDE_DIR cudnn.h
-  PATHS ${CUDNN_ROOT} ${CUDNN_ROOT}/include $ENV{CUDNN_ROOT}
-        $ENV{CUDNN_ROOT}/include ${CUDA_TOOLKIT_INCLUDE}
+  PATHS ${CUDNN_ROOT}
+        ${CUDNN_ROOT}/include
+        ${CUDNN_ROOT}/include/${TARGET_ARCH}-linux-gnu
+        $ENV{CUDNN_ROOT}
+        $ENV{CUDNN_ROOT}/include
+        ${CUDA_TOOLKIT_INCLUDE}
         /usr/local/lib/python${PY_VERSION}/dist-packages/nvidia/cudnn/include/
   NO_DEFAULT_PATH)
 
 get_filename_component(__libpath_hist ${CUDA_CUDART_LIBRARY} PATH)
 
-set(TARGET_ARCH "x86_64")
-if(NOT ${CMAKE_SYSTEM_PROCESSOR})
-  set(TARGET_ARCH ${CMAKE_SYSTEM_PROCESSOR})
-endif()
-
 list(
   APPEND
   CUDNN_CHECK_LIBRARY_DIRS
