Support Ampere devices (#3806)

* Support Ampere devices

* Extract version from cudnn_version.h
Co-authored-by: Noneflow-bot <69100618+oneflow-bot@users.noreply.github.com>

CMakeLists.txt

@@ -110,11 +110,28 @@ if (BUILD_CUDA)
     #  half is not fully supported when __CUDA_ARCH__ < 530
     #  list(APPEND __cuda_nvcc_gencodes "arch=compute_30,code=sm_30")
     #  list(APPEND __cuda_nvcc_gencodes "arch=compute_52,code=sm_52")
+    # cubin
+    # Tesla P100
     list(APPEND __cuda_nvcc_gencodes "arch=compute_60,code=sm_60")
+    # Tesla P40/P4, Quadro Pxxx/Pxxxx, GeForce GTX 10xx, TITAN X/Xp
     list(APPEND __cuda_nvcc_gencodes "arch=compute_61,code=sm_61")
+    # V100, TITAN V
     list(APPEND __cuda_nvcc_gencodes "arch=compute_70,code=sm_70")
     if(CUDA_VERSION VERSION_GREATER_EQUAL "10.0")
+      # T4, Quadro RTX xxxx, Txxxx, Geforce RTX 20xx, TITAN RTX
       list(APPEND __cuda_nvcc_gencodes "arch=compute_75,code=sm_75")
+    endif()
+    if(CUDA_VERSION VERSION_GREATER_EQUAL "11.0")
+      # A100
+      list(APPEND __cuda_nvcc_gencodes "arch=compute_80,code=sm_80")
+    endif()
+    if(CUDA_VERSION VERSION_GREATER_EQUAL "11.1")
+      # GeForce RTX 30xx
+      list(APPEND __cuda_nvcc_gencodes "arch=compute_86,code=sm_86")
+    endif()
+    if(CUDA_VERSION VERSION_GREATER_EQUAL "11.0")
+      list(APPEND __cuda_nvcc_gencodes "arch=compute_80,code=compute_80")
+    elseif(CUDA_VERSION VERSION_GREATER_EQUAL "10.0")
       list(APPEND __cuda_nvcc_gencodes "arch=compute_75,code=compute_75")
     else()
       list(APPEND __cuda_nvcc_gencodes "arch=compute_70,code=compute_70")

cmake/third_party/FindCUDNN.cmake

@@ -33,7 +33,11 @@ find_package_handle_standard_args(
 
 if(CUDNN_FOUND)
 	# get cuDNN version
-  file(READ ${CUDNN_INCLUDE_DIR}/cudnn.h CUDNN_HEADER_CONTENTS)
+  if(EXISTS "${CUDNN_INCLUDE_DIR}/cudnn_version.h")
+	  file(READ ${CUDNN_INCLUDE_DIR}/cudnn_version.h CUDNN_HEADER_CONTENTS)
+  else()
+	  file(READ ${CUDNN_INCLUDE_DIR}/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"